{eval-rst}

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

torch.nn.attention

{eval-rst}

.. automodule:: torch.nn.attention

Utils

{eval-rst}

.. 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

{eval-rst}

.. autosummary::
    :nosignatures:

    flex_attention
    bias
    experimental
    varlen

{eval-rst}

.. toctree::
    :hidden:

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