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.