nn.attention.rst

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

Submodules

.. autosummary::
    :nosignatures:

    flex_attention
    bias
    experimental
    varlen

.. toctree::
    :hidden:

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