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。