Attention

Purpose

Models do not implement attention inline. They declare an _attn_implementation and call into a dispatcher that picks the right backend at runtime. The library supports SDPA (PyTorch built-in), FlashAttention 2 / 3 / 4, FlexAttention, an "eager" reference implementation, and paged variants of each for continuous batching.

Key abstractions

Backend	Source	When picked
eager	src/transformers/modeling_utils.py (in eager_attention_forward)	Fallback; reference correctness
sdpa	src/transformers/integrations/sdpa_attention.py	Default on PyTorch 2+; fast and memory-efficient
flash_attention_2	src/transformers/modeling_flash_attention_utils.py	Requires flash-attn package; fastest on supported GPUs
flash_attention_3	same	Requires flash-attn v3
flash_attention_4	same	Requires flash-attn v4 (Hopper+)
flex_attention	src/transformers/integrations/flex_attention.py	torch.nn.attention.flex_attention; programmable masks
eager_paged, flash_paged, sdpa_paged	src/transformers/integrations/eager_paged.py, flash_paged.py, sdpa_paged.py	Paged-KV variants for continuous batching
npu_flash_attention	src/transformers/integrations/npu_flash_attention.py	Ascend NPU

How dispatch works

graph TD
    User[from_pretrained] --> Choose{attn_implementation kwarg or auto}
    Choose -->|auto| Detect[Probe FA / SDPA availability]
    Detect --> Pick[Pick fastest available]
    Choose -->|explicit| Pick
    Pick --> Set[model.config._attn_implementation]
    Set --> Forward[Layer forward calls dispatcher]
    Forward --> Backend[(Backend implementation)]

Set explicitly:

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B",
    attn_implementation="flash_attention_2",
    dtype=torch.bfloat16,
)

If the requested backend is unavailable, from_pretrained raises a clear error.

What every backend must support

A backend implementation receives:

• query_states, key_states, value_states tensors (after rotary, after grouped-query expansion).
• attention_mask — either an additive mask or a Boolean mask.
• dropout_p, scaling, is_causal.
• The cache (past_key_values) — but the backend reads it indirectly through key_states / value_states after cache.update.

It returns the attention output and (optionally) attention weights when output_attentions=True. Backends that cannot return weights (FlashAttention) emit a warning and fall back to eager when weights are requested.

Masking utilities

src/transformers/masking_utils.py (77K LOC) builds the masks that backends consume. It handles:

• Standard causal masks.
• Sliding-window masks (Mistral).
• Hybrid (per-layer alternating) masks.
• Block / packed attention for sample packing.
• Custom masks for VLMs (image-text-image patterns).

src/transformers/modeling_attn_mask_utils.py provides legacy helpers retained for backward compatibility.

Paged attention

Continuous batching uses a paged KV cache (block-allocated). The matching attention kernels are in:

• src/transformers/integrations/eager_paged.py — pure PyTorch reference.
• src/transformers/integrations/sdpa_paged.py — wraps SDPA with a block table.
• src/transformers/integrations/flash_paged.py — wraps FlashAttention with paged KV.

These are dispatched when the scheduler in src/transformers/generation/continuous_batching/ is active.

FlexAttention

flex_attention exposes a Python API to define custom attention scores (e.g., ALiBi, document masks, soft-cap). Useful for research models. The integration at src/transformers/integrations/flex_attention.py adapts the dispatcher to FlexAttention's score_mod callbacks.

Visualizing masks

src/transformers/utils/attention_visualizer.py renders masks as ASCII grids in the terminal. Helpful when debugging hybrid or sliding-window cache + mask interactions. See Debugging.

Integration points

• Modeling — PreTrainedModel sets _attn_implementation and routes layers.
• Cache — cache.update returns the keys/values that go into attention.
• Continuous batching — uses paged variants.
• Tensor parallelism — heads are sharded; backends are TP-aware.

Entry points for modification

• New attention backend → add a file in src/transformers/integrations/ exporting an <name>_attention_forward(...) function and register it in the dispatcher.
• New mask shape → add a builder in src/transformers/masking_utils.py and tests in tests/test_modeling_common.py.
• New paged variant → update both the integrations file and src/transformers/generation/continuous_batching/ scheduler.