[sfc-gh-goliaro]

[FA4] Support paged/decode compile specs in compile_flash_attn_varlen_func_from_specs

Repo: vllm-project/flash-attention · Branch: fa4-dense-paged-compile-specs

Summary

Extend the CuTeDSL compile-only warmup helper so the dense (non-MLA) paged
decode/prefill specializations can be precompiled, instead of JIT-compiling during
serving. The helper previously exposed only the contiguous varlen (cu_seqlens) axes,
which is enough for MLA prefill but not for vLLM's standard FLASH_ATTN (GQA) path.

New optional parameters (all default to None/False, fully backward compatible):

• seqused_k_shape — paged decode uses seqused_k (not cu_seqlens_k)
• page_table_shape — paged KV
• k_stride — key_cache is a non-contiguous view of a (num_blocks, 2, page_size, …)
  cache, so its num_blocks stride is not the contiguous default; the compiled kernel
  bakes strides, so this must be settable
• q_descale / k_descale / v_descale (bool) — fp8 KV cache
• softcap

All of these are already parameters of _flash_attn_fwd; this only threads them through
the compile-only entry point.

Why this is needed

vLLM warms FA4 kernels at startup via compile_flash_attn_varlen_func_from_specs. For a
standard GQA model the dense forward (FlashAttentionForwardSm100, q_stage=2, paged,
seqused_k) is a different compile key than the varlen/MLA specs this helper could
express, so it could not be precompiled and JIT-compiled on the first serving request.

Test

• compile_flash_attn_varlen_func_from_specs(...) with a paged decode/prefill spec now
  compiles the exact FlashAttentionForwardSm100 specialization vLLM serves with
  (verified by matching the compile_key the engine reports).
• Existing MLA-prefill callers are unaffected (new args default off).

Note for integrators

Paired with vllm-project/vllm#48156, which calls these new args (dense-FA4 CuTeDSL
warmup) and bumps the flash-attn pin to this commit — see that PR for the end-to-end
performance results. Land this PR first (or together).

🤖 AI-assisted change.