Attention backends

Active contributors: Lucas Wilkinson, Matthew Bonanni, Wentao Ye, Woosuk Kwon.

Purpose

Attention is the hot path. vLLM ships ~20 attention backends so that on every supported hardware, every model topology (standard MHA, MLA, sliding-window, sparse, mamba, hybrid), and every quantization scheme there is at least one path that compiles, runs, and is fast.

Directory layout

vllm/v1/attention/
├── backend.py              # AttentionBackend base
├── selector.py             # Runtime selection based on capabilities
└── backends/
    ├── registry.py         # AttentionBackendEnum + register_backend
    ├── flash_attn.py       # FlashAttention 2/3 (NVIDIA)
    ├── flash_attn_diffkv.py
    ├── flashinfer.py       # FlashInfer (NVIDIA, persistent kernels)
    ├── triton_attn.py      # Triton fallback
    ├── flex_attention.py   # PyTorch flex_attention
    ├── tree_attn.py        # Tree-decoded spec attention
    ├── rocm_attn.py        # ROCm baseline
    ├── rocm_aiter_fa.py    # ROCm AITER FlashAttention
    ├── rocm_aiter_unified_attn.py
    ├── cpu_attn.py         # CPU
    ├── linear_attn.py      # Linear attention
    ├── short_conv_attn.py
    ├── mamba_attn.py / mamba1_attn.py / mamba2_attn.py
    ├── gdn_attn.py         # Gated DeltaNet (Qwen3-Next)
    ├── turboquant_attn.py
    ├── fa_utils.py         # FlashAttention helpers
    ├── utils.py            # CommonAttentionMetadata, batch reshapes
    └── mla/                # Multi-head Latent Attention (DeepSeek-style)
        ├── flashattn_mla.py
        ├── flashinfer_mla.py
        ├── flashinfer_mla_sparse.py
        ├── flashmla.py / flashmla_sparse.py
        ├── cutlass_mla.py
        ├── triton_mla.py
        ├── aiter_triton_mla.py
        ├── rocm_aiter_mla.py / rocm_aiter_mla_sparse.py
        ├── xpu_mla_sparse.py
        ├── indexer.py         # Sparse-index builder
        └── sparse_swa.py / sparse_utils.py / compressor_utils.py

Key abstractions

Abstraction	File	Role
AttentionBackend	vllm/v1/attention/backend.py	Abstract base — capabilities, metadata builder, layer factory
AttentionBackendEnum	vllm/v1/attention/backends/registry.py	All known backends, addressable by name
register_backend(name, cls)	vllm/v1/attention/backends/registry.py	Override a backend at runtime (used by plugins)
AttentionImpl	per-backend	The kernel-calling forward
AttentionMetadata	per-backend	Per-step metadata (block tables, slot mapping, seq lens, …)
MLAAttentionImpl	vllm/v1/attention/backends/mla/*	MLA variants
Attention layer	vllm/model_executor/layers/attention/	Model-facing layer; resolves backend from vllm_config

How a backend is chosen

graph TD
    Cfg[ModelConfig + AttentionConfig + KernelConfig]
    Plat[Platform.detect_supported_attn_backends]
    Sel[AttentionBackend.select]
    Reg[AttentionBackendEnum / registry]
    BE[Concrete backend class]

    Cfg --> Sel
    Plat --> Sel
    Sel -->|list of candidates by priority| Reg
    Reg --> BE

The selector (vllm/v1/attention/selector.py) consults:

• The active platform (CUDA / ROCm / CPU / XPU) to filter what's compilable here.
• The model's attention spec (head dim, num heads, dtype, sliding window, MLA latent dim).
• Quantization mode (e.g., FP8 KV cache requires backends that support kv_cache_dtype).
• User overrides via --attention-backend or KernelConfig.attn_backend.

It returns the highest-priority backend that satisfies all constraints. Plugins register backends through register_backend(name, "module.path:Class").

NVIDIA backend overview

Backend	When chosen
FLASH_ATTN	Default for FA2/FA3 hardware. Uses vllm-flash-attn package.
FLASHINFER	Persistent kernels, FP8 KV, sparse attention; chosen for large models
TRITON_ATTN	Triton-only fallback, used when CUDA toolchain unavailable
FLEX_ATTENTION	Uses torch.flex_attention; useful for unusual attention patterns
TREE_ATTN	Required for tree-style speculative decoding (EAGLE)
FLASH_ATTN_DIFFKV	For models that have different K and V dimensions

MLA (Multi-head Latent Attention)

DeepSeek-V2/V3/V4 introduced MLA, where K and V are derived from a low-rank latent. vLLM implements many MLA variants because the workload's optimal kernel changes with sparsity, batch shape, and hardware:

• flashmla.py / flashmla_sparse.py — kernels co-developed with DeepSeek
• flashattn_mla.py — FlashAttention 3 + MLA wrapping
• flashinfer_mla.py / flashinfer_mla_sparse.py — FlashInfer paths
• cutlass_mla.py — CUTLASS-based path
• triton_mla.py / aiter_triton_mla.py — Triton (NVIDIA / ROCm)
• rocm_aiter_mla.py / rocm_aiter_mla_sparse.py — ROCm AITER
• xpu_mla_sparse.py — Intel XPU
• indexer.py — builds sparse attention indices

compressor_utils.py and sparse_utils.py are shared helpers for the sparse variants.

ROCm

rocm_attn.py, rocm_aiter_fa.py, rocm_aiter_unified_attn.py and the MLA rocm_aiter_* variants cover AMD GPUs. The AITER (AMD ITERative kernels) library ships its own optimized paths; vLLM dispatches into them via vllm/_aiter_ops.py.

SSM / hybrid layers

Mamba, GatedDeltaNet (GDN), Lightning, and short-convolution layers also live here because they share the per-step state-management pattern with attention:

• mamba_attn.py (shared utilities)
• mamba1_attn.py, mamba2_attn.py
• gdn_attn.py (Qwen3-Next)
• linear_attn.py
• short_conv_attn.py

These backends manage their own per-layer state cache (analogous to KV cache) and hook into KVCacheConfig via custom AttentionSpec subtypes.

Key source files

File	Purpose
vllm/v1/attention/backend.py	AttentionBackend base class
vllm/v1/attention/backends/registry.py	The enum + plugin override registry
vllm/v1/attention/selector.py	Backend selection logic
vllm/v1/attention/backends/utils.py	CommonAttentionMetadata, batch reshaping helpers
vllm/model_executor/layers/attention/	Model-facing Attention layer
vllm/model_executor/layers/mla.py	The MLA layer wrapper
csrc/attention/	Custom CUDA attention kernels
csrc/mamba/	Mamba CUDA kernels

Entry points for modification

• New backend: subclass AttentionBackend, implement make_metadata_builder and forward, register via AttentionBackendEnum (in-tree) or register_backend(...) (plugin).
• Quantized KV support: implement the kv_cache_dtype capability and provide a quant-aware kernel path; declare it in the backend's supported_kv_cache_dtypes.
• New sparse pattern (e.g., new MLA variant): fork the closest mla/*.py, implement the sparse index builder, and register.
• Per-platform tuning: edit vllm/platforms/{cuda,rocm,xpu,cpu}.py to expose hardware capabilities the selector should use.

For how the layer above attention works, see Model executor. For how the KV blocks fed to the kernel are managed, see KV cache.