FlashAttention 與 FlashInfer:兩大 GPU 注意力後端

源码版本v0.25.1

職責

NVIDIA GPU 上跑得最多的是兩個後端:FlashAttentionBackend(flash_attn.py:67)封裝 Dao Labs 的 flash_attn_varlen_func,從 FA2 到 FA3、FA4 都能自適應;FlashInferBackend(flashinfer.py:330)封裝 FlashInfer 的 prefill / decode kernel,在 Hopper/Blackwell 上對 GQA、FP8 KV cache、page size 128+ 有原生支持。兩者都繼承自 AttentionBackend,實現同一套抽象,但內部 KV cache 形狀、metadata 欄位、cudagraph 路徑都不一樣。

FlashAttentionMetadataBuilder.build(flash_attn.py:428)做的事是把 CommonAttentionMetadata 切成 prefill / decode 兩個段,分別填 query_start_loc、seq_lens、block_table,順帶算 AOT scheduling。FlashInferMetadataBuilder.build(flashinfer.py:1057)額外要做 split_decodes_and_prefills,把 batch 重排成「先 decode 再 prefill」的佈局,因為 FlashInfer 的 decode kernel 假設 batch 前段全是同長度的 decode query。

FlashAttentionImpl.forward(flash_attn.py:749)和 FlashInferImpl.forward(flashinfer.py:1599)是真正調 kernel 的地方:它們拿到 Q/K/V、kv_cache、attn_metadata,先 reshape 成 kernel 要的 [num_tokens, num_heads, head_size],然後調對應的 flash_attn_varlen_func 或 FlashInfer 的 BatchPrefillWithRaggedKVCacheWrapper / BatchDecodeWithPagedKVCacheWrapper。forward_includes_kv_cache_update 在 FA 是 False(FA kernel 自己寫 KV cache),在 FlashInfer 也類似,但 KV cache 寫入路徑走 unified_kv_cache_update op。

設計動機

• FA 版本自適應:get_flash_attn_version(flash_attn.py:714)根據 head_size、alibi、平臺算力在 FA2/3/4 間選,FA3 才支持 sink、FP8 KV cache、per-head quant scales。
• block size 16 對齊:FA 的 get_supported_kernel_block_sizes 返回 MultipleOf(16)(flash_attn.py:76-77),保證 kernel tile 不被切碎。
• NHD / HND 兩種 layout:get_kv_cache_stride_order(flash_attn.py:134-153)根據 KVCacheLayoutType 決定維度排列,FA3 偏好 NHD,FlashInfer 在不同 kernel 上偏好不同 layout。
• FlashInfer 大 page 支持:get_supported_kernel_block_sizes 在 Blackwell + GQA 時返回 [16, 32, 64, 128, 256, 512, 1024](flashinfer.py:343-362),讓大模型用更大的 KV 頁減少 block table 查詢。
• cudagraph 分級:FA builder 預設 _cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE,FlashInfer 在 GQA + uniform batch 時能到 ALWAYS,讓 spec-decode 也能走 cudagraph。
• cascade attention 共享 prefix:FA 通過 use_cascade_attention(flash_attn.py:670-671)讓多請求共享同一個 KV prefix,FlashInfer 走自己的 cascade 路徑,兩邊都把 batch 內 prefix 抽出來只算一次。

關鍵檔案

• FlashAttentionBackend:67 — 後端類,聲明支持 dtype / kv_cache_dtype / block size / head_size。
• FlashAttention get_kv_cache_shape:122-132 — (num_blocks, 2, block_size, num_kv_heads, head_size),2 是 K/V 分槽。
• FlashAttentionMetadataBuilder.build:428 — prefill + decode 段切分、AOT schedule 決策。
• FlashAttentionImpl:674 — impl 類,forward_includes_kv_cache_update = False。
• FlashAttentionImpl.forward:749 — 調 flash_attn_varlen_func 的主入口,Q/K/V shape [num_tokens, num_heads, head_size]。
• FlashInferBackend:330 — FlashInfer 後端類,supported_kv_cache_dtypes 含 fp8/nvfp4。
• FlashInfer get_kv_cache_shape:380-389 — 根據 dtype 返回不同形狀,nvfp4 走 packed layout。
• FlashInferMetadataBuilder:600 — 持 FIPrefill / FIDecode 子 builder,按 phase 分流。
• FlashInferMetadataBuilder.build:1057 — split_decodes_and_prefills 重排 batch 後填 metadata。
• FlashInferImpl.forward:1599 — 選 prefill 或 decode wrapper 調 FlashInfer kernel。

資料流

兩個後端的 forward 接收的 tensor 形狀一樣:Q 是 [num_tokens, num_heads, head_size],KV cache 由後端自己的 get_kv_cache_shape 決定。下面是 FA 的 KV cache 形狀和佈局選擇:

# vllm/v1/attention/backends/flash_attn.py L122-L132
@staticmethod
def get_kv_cache_shape(
    num_blocks: int,
    block_size: int,
    num_kv_heads: int,
    head_size: int,
    cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
    if block_size % 16 != 0:
        raise ValueError("Block size must be a multiple of 16.")
    return (num_blocks, 2, block_size, num_kv_heads, head_size)

FlashInfer 在 Blackwell + GQA 時會暴露更大的 page size,讓 KV block 數變少、block table 查詢更省:

# vllm/v1/attention/backends/flashinfer.py L342-L362
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
    # Page sizes >= 128 only run on the trtllm-gen dynamic kernel (GQA/MQA
    # on Blackwell); advertise them only when usable so selection never
    # picks a large kernel block we cannot serve.
    use_large_pages = False
    vllm_config = get_current_vllm_config_or_none()
    if vllm_config is not None and vllm_config.model_config is not None:
        pc = vllm_config.parallel_config
        mc = vllm_config.model_config
        num_qo_heads = mc.get_num_attention_heads(pc)
        num_kv_heads = mc.get_num_kv_heads(pc)
        use_large_pages = (
            num_kv_heads > 0
            and num_qo_heads // num_kv_heads > 1
            and current_platform.is_device_capability_family(100)
            and can_use_trtllm_attention(num_qo_heads, num_kv_heads)
        )
    if not use_large_pages:
        return [16, 32, 64]
    return [16, 32, 64, 128, 256, 512, 1024]

FlashInferMetadataBuilder.build 在拿到 CommonAttentionMetadata 後先做 batch 重排,把 decode 段挪到 batch 前面,這樣 decode kernel 拿到的就是一段連續的同長度 query,不用做 varlen 索引。FA 不重排,因為它直接走 flash_attn_varlen_func,varlen 本來就支持任意長度段。

邊界與失敗

• FA4 才支持 512 head_size:supports_head_size(flash_attn.py:155-163)在 head_size > 256 時只有 is_fa_version_supported(4) 才返回 True。
• FP8 KV cache 限制:FA 僅在 FA3 + Hopper(is_device_capability_family(90))(flash_attn.py:166-176)接受 fp8/fp8_e4m3,其它組合回退到其它後端。
• block_size 必須是 16 的倍數:FA 和 FI 的 get_supported_kernel_block_sizes 都要求 MultipleOf(16),非 16 倍數 block size 會讓 selector 跳過該後端。
• FlashInfer 不能跑非因果 query-query 注意力:build 裡 causal=False 時直接走 native prefill(flashinfer.py:1074-1080),因為 FlashInfer decode/TRTLLM 路徑不表達雙向注意力。
• FlashInfer 不返回 post-top-k logits:TopKTopPSampler 在 logprobs_mode 為 processed_logits / processed_logprobs 時不走 forward_cuda(topk_topp_sampler.py:86-95),因為 FlashInfer sampler 不暴露過濾後的 logits。
• DCP combine 分支:FlashAttentionImpl.__init__ 根據 dcp_comm_backend 選 cp_lse_ag_out_rs 或 dcp_a2a_lse_reduce(flash_attn.py:738-743),選錯會讓跨分片 softmax 合併錯。

小結

FlashAttention 是 vLLM 在 NVIDIA GPU 上的預設後端,覆蓋最全的 dtype / head_size 組合;FlashInfer 在 Hopper/Blackwell 上對 GQA、FP8 KV cache、大 page 有更原生支持,跑 decode batch 性能更穩。兩者都通過 AttentionBackend 抽象暴露給上層,selector 按模型和硬件挑一個。其它後端見 /attention/backend(抽象層)和 /attention/triton-mla(Triton / MLA)。採樣層銜接見 /sampling/sampler。

對照官方資料:vLLM 文件 · README