FlashAttention and FlashInfer: the two GPU attention backends

源码版本v0.25.1

Responsibilities

Two backends dominate on NVIDIA GPUs: FlashAttentionBackend(flash_attn.py:67) wraps Dao Labs' flash_attn_varlen_func and adapts across FA2, FA3, and FA4; FlashInferBackend(flashinfer.py:330) wraps FlashInfer's prefill / decode kernels and has native support for GQA, FP8 KV cache, and page size 128+ on Hopper/Blackwell. Both inherit from AttentionBackend and implement the same abstraction, but their internal KV cache shapes, metadata fields, and cudagraph paths differ.

FlashAttentionMetadataBuilder.build(flash_attn.py:428) slices CommonAttentionMetadata into prefill / decode segments, fills in query_start_loc, seq_lens, block_table, and computes AOT scheduling along the way. FlashInferMetadataBuilder.build(flashinfer.py:1057) additionally does split_decodes_and_prefills, reordering the batch into a "decode first, then prefill" layout, because FlashInfer's decode kernel assumes the front segment of the batch consists entirely of same-length decode queries.

FlashAttentionImpl.forward(flash_attn.py:749) and FlashInferImpl.forward(flashinfer.py:1599) are where the kernels are actually invoked: they take Q/K/V, kv_cache, and attn_metadata, first reshape them into the [num_tokens, num_heads, head_size] layout the kernel expects, and then call either flash_attn_varlen_func or FlashInfer's BatchPrefillWithRaggedKVCacheWrapper / BatchDecodeWithPagedKVCacheWrapper. forward_includes_kv_cache_update is False for FA (the FA kernel writes the KV cache itself), and similarly for FlashInfer, but the KV cache write path goes through the unified_kv_cache_update op.

Design motivation

• FA version auto-selection: get_flash_attn_version(flash_attn.py:714) chooses among FA2/3/4 based on head_size, alibi, and platform compute capability; only FA3 supports sink, FP8 KV cache, and per-head quant scales.
• block size 16 alignment: FA's get_supported_kernel_block_sizes returns MultipleOf(16)(flash_attn.py:76-77), ensuring kernel tiles are not split.
• NHD / HND layouts: get_kv_cache_stride_order(flash_attn.py:134-153) decides the dimension permutation based on KVCacheLayoutType; FA3 prefers NHD, while FlashInfer prefers different layouts on different kernels.
• FlashInfer large page support: get_supported_kernel_block_sizes returns [16, 32, 64, 128, 256, 512, 1024] on Blackwell + GQA(flashinfer.py:343-362), letting large models use larger KV pages and reduce block table lookups.
• tiered cudagraph: the FA builder defaults to _cudagraph_support = UNIFORM_SINGLE_TOKEN_DECODE, while FlashInfer can reach ALWAYS on GQA + uniform batches, letting spec-decode also go through cudagraph.
• cascade attention shares prefixes: FA uses use_cascade_attention(flash_attn.py:670-671) to let multiple requests share the same KV prefix, while FlashInfer goes through its own cascade path; both extract the in-batch prefix and compute it only once.

Key files

• FlashAttentionBackend:67 — backend class, declares supported 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), the 2 is the K/V slot.
• FlashAttentionMetadataBuilder.build:428 — prefill + decode segment split, AOT schedule decisions.
• FlashAttentionImpl:674 — impl class, forward_includes_kv_cache_update = False.
• FlashAttentionImpl.forward:749 — main entry calling flash_attn_varlen_func, with Q/K/V shape [num_tokens, num_heads, head_size].
• FlashInferBackend:330 — FlashInfer backend class, supported_kv_cache_dtypes includes fp8/nvfp4.
• FlashInfer get_kv_cache_shape:380-389 — returns different shapes based on dtype; nvfp4 goes through a packed layout.
• FlashInferMetadataBuilder:600 — holds FIPrefill / FIDecode sub-builders, dispatching by phase.
• FlashInferMetadataBuilder.build:1057 — reorders the batch via split_decodes_and_prefills then fills metadata.
• FlashInferImpl.forward:1599 — selects the prefill or decode wrapper to call the FlashInfer kernel.

Data flow

Both backends' forward receive tensors of the same shape: Q is [num_tokens, num_heads, head_size], and the KV cache shape is decided by each backend's own get_kv_cache_shape. Below is FA's KV cache shape and layout choice:

# 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 exposes larger page sizes on Blackwell + GQA, reducing the number of KV blocks and saving block table lookups:

# 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]

After FlashInferMetadataBuilder.build receives CommonAttentionMetadata, it first reorders the batch, moving the decode segment to the front so the decode kernel receives a contiguous run of same-length queries and does not need varlen indexing. FA does not reorder, because it goes directly through flash_attn_varlen_func, which natively supports arbitrary-length segments.

Boundaries and failure modes

• Only FA4 supports 512 head_size: supports_head_size(flash_attn.py:155-163) only returns True for is_fa_version_supported(4) when head_size > 256.
• FP8 KV cache limits: FA only accepts fp8/fp8_e4m3 on FA3 + Hopper (is_device_capability_family(90))(flash_attn.py:166-176); other combinations fall back to other backends.
• block_size must be a multiple of 16: both FA and FI's get_supported_kernel_block_sizes require MultipleOf(16); non-16-multiple block sizes cause the selector to skip the backend.
• FlashInfer cannot run non-causal query-query attention: when causal=False in build, it falls back to native prefill(flashinfer.py:1074-1080), because FlashInfer's decode/TRTLLM path does not express bidirectional attention.
• FlashInfer does not return post-top-k logits: TopKTopPSampler does not go through forward_cuda when logprobs_mode is processed_logits / processed_logprobs(topk_topp_sampler.py:86-95), because the FlashInfer sampler does not expose the filtered logits.
• DCP combine branch: FlashAttentionImpl.__init__ selects cp_lse_ag_out_rs or dcp_a2a_lse_reduce based on dcp_comm_backend(flash_attn.py:738-743); choosing wrong leads to incorrect cross-shard softmax merges.

Summary

FlashAttention is vLLM's default backend on NVIDIA GPUs, covering the widest range of dtype / head_size combinations; FlashInfer has more native support for GQA, FP8 KV cache, and large pages on Hopper/Blackwell, delivering more stable performance on decode batches. Both are exposed to the upper layer through the AttentionBackend abstraction, and a selector picks one based on model and hardware. See other backends at /attention/backend (abstraction layer) and /attention/triton-mla (Triton / MLA). For the integration with the sampling layer, see /sampling/sampler.

See official docs: vLLM docs · README