FlashAttention 与 FlashInfer:两大 GPU 注意力后端
职责
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。