.. role:: hidden :class: hidden-section

torch.nn.attention

.. automodule:: torch.nn.attention

Utils

.. autosummary:: :toctree: generated :nosignatures:

sdpa_kernel
SDPBackend
register_flash_attention_impl
activate_flash_attention_impl
list_flash_attention_impls
current_flash_attention_impl
restore_flash_attention_impl

Submodules

.. autosummary:: :nosignatures:

flex_attention
bias
experimental
varlen

.. toctree:: :hidden:

nn.attention.flex_attention
nn.attention.bias
nn.attention.experimental
nn.attention.varlen