torch.nn.attention#

Created On: Jan 24, 2024 | Last Updated On: Nov 12, 2025

This module contains functions and classes that alter the behavior of torch.nn.functional.scaled_dot_product_attention

Utils#

`sdpa_kernel`	Context manager to select which backend to use for scaled dot product attention.
`SDPBackend`	An enum-like class that contains the different backends for scaled dot product attention.
`register_flash_attention_impl`	Register the callable that activates a flash attention impl.
`activate_flash_attention_impl`	Activate into the dispatcher a previously registered flash attention impl.
`list_flash_attention_impls`	Return the names of all available flash attention implementations.
`current_flash_attention_impl`	Return the currently activated flash attention impl name, if any.

`flex_attention`	This module implements the user facing API for flex_attention in PyTorch.
`bias`	Defines bias subclasses that work with scaled_dot_product_attention
`experimental`
`varlen`	Variable-length attention implementation using Flash Attention.