TopolearnAttention kernels and placement decide whether a long-context model is fast, slow or fits at all.
Attention is the most expensive operation in a transformer. Transformers decouples the model code from the attention implementation, selected with attn_implementation at load time:
"sdpa"(default when supported): PyTorch'sscaled_dot_product_attention, which dispatches to fused FlashAttention-style or memory-efficient kernels automatically."flash_attention_2"/"flash_attention_3": the FlashAttention packages; very fast on supported NVIDIA GPUs, require fp16/bf16, and enable padding-free training where sequences are packed without pad tokens."flex_attention": PyTorch's programmable attention for custom masks."eager": the reference Python implementation, slow but needed to return attention weights (output_attentions=True) or for debugging.
Through the attention interface you can register your own function (AttentionInterface.register) or load kernels from the Hub, and switch an existing model with model.set_attn_implementation(...).
Placement: device_map="auto" uses Accelerate to fill GPUs in order, then CPU RAM, then disk; max_memory={0: "20GiB", "cpu": "60GiB"} sets budgets. Offloaded layers work but are slow, because weights stream over PCIe every forward pass. model.hf_device_map shows the result. device_map="auto" splits layers across GPUs sequentially (pipeline-style, one GPU busy at a time); real tensor parallelism uses tp_plan="auto" under torchrun.
Benchmark before switching: gains depend on sequence length, batch size, dtype and GPU generation.
Example
import torch
from transformers import AutoModelForCausalLM
repo = "Qwen/Qwen2.5-1.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
repo, dtype=torch.bfloat16, attn_implementation="sdpa", device_map="auto",
max_memory={0: "6GiB", "cpu": "24GiB"}, # spill to CPU beyond 6 GiB of GPU 0
)
print(model.config._attn_implementation) # sdpa
print(set(model.hf_device_map.values())) # e.g. {0} or {0, 'cpu'}
# eager attention is required to inspect attention weights
eager = AutoModelForCausalLM.from_pretrained(repo, attn_implementation="eager")
ids = torch.tensor([[1, 2, 3]])
att = eager(ids, output_attentions=True).attentions
print(len(att), att[0].shape) # num_layers, (batch, heads, seq, seq)
# FlashAttention (needs the flash-attn package, fp16/bf16, supported GPU):
# AutoModelForCausalLM.from_pretrained(repo, dtype=torch.bfloat16,
# attn_implementation="flash_attention_2")Remember
attn_implementationselects sdpa, flash_attention_2/3, flex_attention or eager- Only eager returns attention weights; fused kernels never materialize them
device_map="auto"withmax_memoryplaces layers across GPUs, CPU and disk
Watch out for
- Assuming device_map='auto' across two GPUs doubles speed; it splits layers sequentially, so only one GPU works at a time.