fix(fmha_v2): enable flash_attention for Q_PAGED_KV regardless of s_kv

[blake-snc]

Problem

determine_launch_params() in csrc/fmha_v2_run.cu disables flash_attention when max_kv_len < 16 via the s >= 16 check. For Q_PAGED_KV layouts this falls back to a non-flash-attention kernel that does not correctly read paged KV — producing wrong output (max_diff 4+ vs reference, or NaN) for any prefill whose total KV length is under 16 tokens.

Reproducer

import torch
from flashinfer.prefill import trtllm_fmha_v2_prefill

device = torch.device("cuda")
torch.manual_seed(42)
# Qwen2.5-0.5B-style GQA: 14 Q heads, 2 KV heads, D=64, page_size=16
# seq_len=9 (< page_size), q_len=8 (prefix-cached scenario)
seq_len, q_len, PS = 9, 8, 16
full_k = torch.randn(seq_len, 2, 64, device=device, dtype=torch.float16)
full_v = torch.randn(seq_len, 2, 64, device=device, dtype=torch.float16)
q_full = torch.randn(seq_len, 14, 64, device=device, dtype=torch.float16)

kv = torch.zeros(1, 2, PS, 2, 64, device=device, dtype=torch.float16)
kv[0, 0, :seq_len] = full_k
kv[0, 1, :seq_len] = full_v
# ... (call trtllm_fmha_v2_prefill with Q_PAGED_KV_NHD)
# Result before fix: max_diff = 2.27 vs reference attention
# Result after fix:  max_diff = 0.002

Root cause

The s >= 16 gate was added for pre-flash-attention fused kernels that required a minimum sequence length for tiling. Q_PAGED_KV layouts only have a flash-attention code path, so this gate incorrectly disqualifies them. Per-request bounds are already enforced via seq_lens; s is the padded max_kv_len used for tile scheduling, not a correctness bound.

Fix

Exempt Q_PAGED_KV layouts from the s >= 16 check:

bool const is_paged_kv = input_layout == Attention_input_layout::Q_PAGED_KV;
launch_params.flash_attention =
    (data_type == DATA_TYPE_FP16 || data_type == DATA_TYPE_BF16 || data_type == DATA_TYPE_E4M3) &&
    (is_paged_kv || (s >= 16 && d >= 16)) && !force_non_flash_attention;

Validation

Validated on SM121a (DGX Spark GB10, NVIDIA GB10):

• Reproducer (seq_len=9, page_size=16): max_diff 2.27 → 0.002 vs reference
• seq_len in {1, 5, 9, 15, 16, 17, 25, 33}: all match reference (max_diff < 0.01)
• GQA (14 Q heads, 2 KV heads): max_diff 0.002 vs reference
• vLLM end-to-end (Qwen2.5-0.5B, trl-internal-testing/tiny-GptOssForCausalLM) with prefix caching: deterministic output across repeated calls, no regressions on prior page-aligned cases

Contributed by Second Nature Computing (https://joinsecondnature.com)

Summary by CodeRabbit

• Bug Fixes
  • Improved flash attention dispatch for paged-KV attention layouts, removing sequence length and dimension constraints.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 63f71669-e663-4585-b39b-b8ca7f117bf1

📥 Commits

Reviewing files that changed from the base of the PR and between 24f2032 and 5192ec5.

📒 Files selected for processing (1)

• csrc/fmha_v2_run.cu

---

📝 Walkthrough

Walkthrough

Modified the flash_attention dispatch selection in determine_launch_params to unconditionally enable flash_attention for paged-KV layouts, bypassing the sequence length and dimension threshold constraints that normally apply.

Changes

Cohort / File(s)	Summary
Flash Attention Paged-KV Dispatch csrc/fmha_v2_run.cu	Introduced is_paged_kv boolean to detect Q_PAGED_KV layout and updated flash_attention condition to enable kernel selection for paged-KV regardless of (s >= 16 && d >= 16) threshold.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• [Fmha] support nvfp4 output keepsMmaAb generation kernels #2988: Adds paged-KV index support and updates generation-kernel selection, directly related to paged-KV dispatch optimization.

Suggested labels

op: attention

Suggested reviewers

• yzh119
• bkryu
• jimmyzho
• cyx-6
• nvmbreughe

Poem

🐰 A paged-KV takes the stage so bright,
No size constraints to dim its light,
Flash kernels dance without a care,
Attention flows through memory fair! ✨

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The pull request title clearly and specifically describes the main change: enabling flash_attention for Q_PAGED_KV layouts regardless of sequence length constraints.
Description check	✅ Passed	The pull request description provides comprehensive coverage: it clearly states the problem, includes a code reproducer, explains the root cause, details the fix with code snippets, and documents extensive validation results.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request modifies the determine_launch_params function in csrc/fmha_v2_run.cu to ensure that flash attention is always enabled for Q_PAGED_KV layouts, regardless of the sequence length. This change addresses an issue where falling back to non-flash paths for small sequence lengths resulted in incorrect outputs, as the paged KV dispatch path only supports flash attention kernels. I have no feedback to provide.