AttentionBackend:注意力后端抽象层
职责
vLLM 支持的注意力 (attention) 实现多到一只手数不过来:FlashAttention 2/3/4、FlashInfer、Triton、MLA、Mamba、ROCm AITER、XPU……每家在 KV cache 布局、kernel block size、cudagraph 支持上都不一样。AttentionBackend 这一层的存在,就是把这些差异吸收掉,让上层 GPUModelRunner 只调一个统一的 forward,不必关心当前跑在哪一种 kernel 上。
抽象分两半:AttentionBackend(backend.py:56)本身是类属性 + 静态方法的工厂,回答「我叫什么名字、用哪个 impl 类、用哪个 metadata builder、KV cache 该长成什么形状」;AttentionMetadataBuilder(backend.py:600)则负责在每个调度 step 把 CommonAttentionMetadata(跨后端共享的 per-batch 元数据,如 query_start_loc、seq_lens、block_table_tensor)翻译成当前后端 kernel 真正要吃的 AttentionMetadata。再下面 AttentionImpl(backend.py:858)才是真正持有 forward 方法、被模型层每层调一次的那个对象。
后端是可插拔的:AttentionBackendEnum(registry.py:34)把所有后端的「类路径」列成枚举,运行时 get_attn_backend(selector.py:54)根据 head_size、dtype、kv_cache_dtype、是否 MLA、平台算力等条件挑一个,懒加载对应模块。模型层只要在初始化时拿到 type[AttentionBackend],后面的事就交给这套抽象。
设计动机
- 跨硬件统一:同一份模型代码在 NVIDIA / AMD / Intel XPU / CPU 上都能跑,差别只在选哪个后端,selector 把决策集中到一处。
- KV cache 形状解耦:每个后端的
get_kv_cache_shape决定 block 池的物理 layout(FlashAttention 是[num_blocks, 2, block_size, num_kv_heads, head_size],MLA 是[num_blocks, block_size, head_size]),KVCacheManager只看 spec 不关心具体后端。 - cudagraph 支持分级:
AttentionCGSupport(backend.py:583)分 ALWAYS / UNIFORM_BATCH / UNIFORM_SINGLE_TOKEN_DECODE / NEVER 四档,让GPUModelRunner知道当前后端能不能塞进 cudagraph。 - argmax 不变性分类:
AttentionImplBase(backend.py:769)区分「能否返回 softmax lse」「lse 是 ln 还是 log2」「支持 Prefill Context Parallelism」等属性,DCP 跨分片合并 kernel 靠这些标志选分支。 - MLA 走单独通路:
MLAAttentionImpl(backend.py:941)不实现forward,而是forward_mqa/forward_mha,因为 MLA 的 KV cache 存的是压缩后的 latent,不是普通 K/V。
关键文件
AttentionBackend ABC:56— 抽象基类,定义get_name/get_impl_cls/get_builder_cls/get_kv_cache_shape等静态方法。四个核心 abstractmethod:74-97—get_name、get_impl_cls、get_builder_cls、get_kv_cache_shape,子类必须实现。CommonAttentionMetadata:395— 跨后端共享的 per-batch 元数据,query_start_loc、seq_lens、block_table_tensor等都在这里。AttentionMetadataBuilder ABC:600—build方法把CommonAttentionMetadata翻译成后端自己的 metadata。AttentionImplBase:769— 真正forward的基类,持can_return_lse_for_decode、lse_base_on_e、supports_pcp等能力标志。AttentionImpl:858— 标准注意力实现 ABC,定义forward(layer, query, key, value, kv_cache, attn_metadata, output)。MLAAttentionImpl:941— MLA 专用 ABC,接口是forward_mqa/forward_mha,不走标准forward。get_attn_backend:54— 按 head_size/dtype/kv_cache_dtype/use_mla 等条件选后端,带@cache。AttentionBackendEnum:34— 把所有后端类路径列成枚举,默认值指向vllm.v1.attention.backends.*。register_backend:233— 运行时覆盖某枚举成员指向的类,用于插件。
数据流
一个 step 里 GPUModelRunner 先把整批请求的 query/key/value 和 CommonAttentionMetadata 准备好,然后逐层调 AttentionImpl.forward。每一层之前,后端的 AttentionMetadataBuilder.build 已经把公共元数据加工成了 kernel 需要的形状。下面是 AttentionMetadataBuilder.build 的签名,所有后端都实现这同一个接口:
# 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.后端选哪条路由 get_attn_backend 决定,它根据当前 vllm_config + 模型 dtype/head_size 等参数,从 AttentionBackendEnum 里捞出后端类路径并懒加载:
# 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,
)模型层在初始化每个 Attention / MLAAttention 时拿到这个 backend 类,然后 backend.get_impl_cls()(...) 构造 impl,backend.get_builder_cls()(...) 构造 metadata builder,后面整轮推理都复用这两个对象。
边界与失败
- 不支持的头尺寸:
AttentionBackend.supports_head_size(backend.py:159)会先问后端,FlashAttention 要求 head_size 是 8 的倍数且不超 256(FA4 才到 512),不满足直接换后端。 - 不支持的数据类型 / KV cache dtype:
supports_dtype/supports_kv_cache_dtype(backend.py:163-173)分别检查,FlashInfer 支持 fp8/nvfp4,FlashAttention 只在 Hopper + FA3 才接 fp8。 - block size 必须 16 对齐:FlashAttention / Triton / MLA 的
get_supported_kernel_block_sizes都返回MultipleOf(16),block_size 不对齐直接抛ValueError(flash_attn.py:130-131)。 - MLA 后端不复用
forward:MLAAttentionImpl(backend.py:941)只声明forward_mqa/forward_mha,模型层的MLAAttention会根据上下文调用对应方法,误用forward会直接AttributeError。 - DCP lse 基数不能错:
AttentionImplBase.lse_base_on_e(backend.py:795)区分 ln 与 log2,选错会让跨分片 softmax 分母静默错乱。 - selector 缓存假设配置不变:
_cached_get_attn_backend(selector.py:114)用@cache,运行中改 dtype/head_size 不会重选后端,需要重启进程。
小结
AttentionBackend 这层抽象是 vLLM 多硬件、多 kernel 共存的关键:模型层只认 AttentionImpl.forward 这一个接口,后端怎么排布 KV cache、怎么生成 metadata、走 FlashAttention 还是 Triton,全靠 selector + enum 在初始化时定下来。具体每种后端怎么实现,见 /attention/flash(FlashAttention / FlashInfer)和 /attention/triton-mla(Triton / MLA)。上层在哪里调这些 impl,见 /worker/gpu-model-runner。