AttentionBackend: attention backend abstraction layer

源码版本v0.25.1

Responsibilities

vLLM supports more attention implementations than you can count on one hand: FlashAttention 2/3/4, FlashInfer, Triton, MLA, Mamba, ROCm AITER, XPU... Each differs in KV cache layout, kernel block size, and cudagraph support. The AttentionBackend layer exists to absorb those differences so the upper-layer GPUModelRunner only calls a unified forward without caring which kernel is currently running.

The abstraction has two halves. AttentionBackend(backend.py:56) is itself a factory of class attributes and static methods, answering "what is my name, which impl class, which metadata builder, and what shape should the KV cache be"; AttentionMetadataBuilder(backend.py:600) is responsible, at each scheduling step, for translating CommonAttentionMetadata(per-batch metadata shared across backends, such as query_start_loc, seq_lens, block_table_tensor) into the AttentionMetadata that the current backend's kernel actually consumes. Below that, AttentionImpl(backend.py:858) is the object that actually holds the forward method and is called once per layer by the model.

Backends are pluggable: AttentionBackendEnum(registry.py:34) lists the "class path" of every backend as an enum, and at runtime get_attn_backend(selector.py:54) picks one based on head_size, dtype, kv_cache_dtype, whether MLA is in use, platform capability, and so on, then lazy-loads the corresponding module. The model layer only needs to obtain type[AttentionBackend] at initialization; the rest is handled by this abstraction.

Design motivation

• Cross-hardware unification: the same model code runs on NVIDIA / AMD / Intel XPU / CPU; the only difference is which backend is selected, and the selector centralizes that decision.
• KV cache shape decoupling: each backend's get_kv_cache_shape determines the physical layout of the block pool (FlashAttention uses [num_blocks, 2, block_size, num_kv_heads, head_size], MLA uses [num_blocks, block_size, head_size]), and KVCacheManager only looks at the spec without caring about the specific backend.
• Tiered cudagraph support: AttentionCGSupport(backend.py:583) classifies into ALWAYS / UNIFORM_BATCH / UNIFORM_SINGLE_TOKEN_DECODE / NEVER, so GPUModelRunner knows whether the current backend can fit inside a cudagraph.
• argmax invariance classification: AttentionImplBase(backend.py:769) distinguishes properties such as "can return softmax lse", "is lse ln or log2", "supports Prefill Context Parallelism"; DCP cross-shard merge kernels rely on these flags to pick a branch.
• MLA uses a separate path: MLAAttentionImpl(backend.py:941) does not implement forward; instead it exposes forward_mqa / forward_mha, because MLA's KV cache stores the compressed latent, not ordinary K/V.

Key files

• AttentionBackend ABC:56 — abstract base class defining static methods such as get_name / get_impl_cls / get_builder_cls / get_kv_cache_shape.
• four core abstractmethods:74-97 — get_name, get_impl_cls, get_builder_cls, get_kv_cache_shape; subclasses must implement them.
• CommonAttentionMetadata:395 — per-batch metadata shared across backends; query_start_loc, seq_lens, block_table_tensor etc. live here.
• AttentionMetadataBuilder ABC:600 — the build method translates CommonAttentionMetadata into the backend's own metadata.
• AttentionImplBase:769 — base class for the real forward; holds capability flags such as can_return_lse_for_decode, lse_base_on_e, supports_pcp.
• AttentionImpl:858 — ABC for the standard attention implementation; defines forward(layer, query, key, value, kv_cache, attn_metadata, output).
• MLAAttentionImpl:941 — MLA-specific ABC; the interface is forward_mqa / forward_mha and does not go through the standard forward.
• get_attn_backend:54 — selects a backend by head_size/dtype/kv_cache_dtype/use_mla etc., with @cache.
• AttentionBackendEnum:34 — lists every backend class path as an enum; the default values point to vllm.v1.attention.backends.*.
• register_backend:233 — overrides the class an enum member points to at runtime, used for plugins.

Data flow

In a step, GPUModelRunner first prepares the batch's query/key/value and CommonAttentionMetadata, then calls AttentionImpl.forward layer by layer. Before each layer, the backend's AttentionMetadataBuilder.build has already shaped the common metadata into what the kernel needs. Below is the signature of AttentionMetadataBuilder.build; every backend implements this same interface:

# vllm/v1/attention/backend.py L666-L678
@abstractmethod
def build(
    self,
    common_prefix_len: int,
    common_attn_metadata: CommonAttentionMetadata,
    fast_build: bool = False,
) -> M:
    """
    Central method that builds attention metadata.
    Some builders (MLA) require reorder_batch to be called prior to build.

    Args:
        common_prefix_len: The length of the common prefix of the batch.
        common_attn_metadata: The common attention metadata.

Which backend gets routed is decided by get_attn_backend, which takes the current vllm_config + model dtype/head_size etc., pulls the backend class path out of AttentionBackendEnum and lazy-loads it:

# vllm/v1/attention/selector.py L75-L110
from vllm.config import get_current_vllm_config

vllm_config = get_current_vllm_config()

cache_config = vllm_config.cache_config
if cache_config is not None and cache_config.user_specified_block_size:
    block_size = cache_config.block_size
else:
    block_size = None

kv_transfer_config = vllm_config.kv_transfer_config
use_kv_connector = (
    kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance
)

attn_selector_config = AttentionSelectorConfig(
    head_size=head_size,
    dtype=dtype,
    kv_cache_dtype=cast(CacheDType | None, kv_cache_dtype),
    block_size=block_size,
    use_mla=use_mla,
    ...
)

return _cached_get_attn_backend(
    backend=vllm_config.attention_config.backend,
    attn_selector_config=attn_selector_config,
    num_heads=num_heads,
)

When initializing each Attention / MLAAttention, the model layer takes this backend class, then constructs the impl via backend.get_impl_cls()(...) and the metadata builder via backend.get_builder_cls()(...); both objects are reused for the rest of inference.

Boundaries and failure modes

• Unsupported head sizes: AttentionBackend.supports_head_size(backend.py:159) asks the backend first; FlashAttention requires head_size to be a multiple of 8 and no greater than 256 (FA4 reaches 512), otherwise the backend is swapped.
• Unsupported dtypes / KV cache dtypes: supports_dtype / supports_kv_cache_dtype(backend.py:163-173) check separately; FlashInfer supports fp8/nvfp4, while FlashAttention only accepts fp8 on Hopper + FA3.
• block size must be 16-aligned: get_supported_kernel_block_sizes for FlashAttention / Triton / MLA all return MultipleOf(16); a misaligned block_size raises ValueError(flash_attn.py:130-131).
• MLA backends do not reuse forward: MLAAttentionImpl(backend.py:941) only declares forward_mqa / forward_mha; the model layer's MLAAttention calls the matching method based on context, and misuse of forward raises AttributeError directly.
• DCP lse base must not be wrong: AttentionImplBase.lse_base_on_e(backend.py:795) distinguishes ln from log2; picking the wrong one silently corrupts the cross-shard softmax denominator.
• selector cache assumes config is immutable: _cached_get_attn_backend(selector.py:114) uses @cache; changing dtype/head_size at runtime will not reselect a backend, the process must be restarted.

Summary

The AttentionBackend abstraction is what lets vLLM keep multiple hardware targets and multiple kernels side by side: the model layer only knows the single AttentionImpl.forward interface, while how KV cache is laid out, how metadata is generated, and whether FlashAttention or Triton is used are all pinned down by the selector + enum at initialization. For how each backend is implemented, see /attention/flash(FlashAttention / FlashInfer) and /attention/triton-mla(Triton / MLA). Where the upper layer calls these impls, see /worker/gpu-model-runner.

See official docs: vLLM docs · README