feat(attention): asymmetric VO-split NVFP4 paged prefill (qk=512/vo=256) for Gemma-4 on SM120/121 - #3684
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR updates NVFP4 KV-cache handling, attention JIT generation, CTA tile selection, Python wrapper logic, and SM12x build/runtime messages. It separates K/V strides, threads explicit SF strides through prefill paths, adjusts output sizing, and updates SM121-related architecture and error text. ChangesNVFP4 KV-cache and SM12x support
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for independent K and V cache strides, asymmetric QK/VO head dimension configurations, and explicit scale-factor strides for NVFP4 KV cache in FlashInfer's attention kernels and JIT compilation modules. It also updates the FP4 GEMM heuristics to support SM12x (including GB10) and adds shared memory validation guards to prevent cryptic launch failures. Feedback on the code changes highlights a critical bug in flashinfer/cute_dsl/utils.py where the cute module is referenced without being imported, which will lead to a NameError at runtime.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if not hasattr(cute.nvgpu, "OperandMajorMode"): | ||
| try: | ||
| cute.nvgpu.OperandMajorMode = cute.nvgpu.tcgen05.OperandMajorMode | ||
| except AttributeError: | ||
| pass |
There was a problem hiding this comment.
The cute module is referenced here but it is not imported in this file. This will raise a NameError when the module is loaded. Please import cute from cutlass.
| if not hasattr(cute.nvgpu, "OperandMajorMode"): | |
| try: | |
| cute.nvgpu.OperandMajorMode = cute.nvgpu.tcgen05.OperandMajorMode | |
| except AttributeError: | |
| pass | |
| import cutlass.cute as cute | |
| if not hasattr(cute.nvgpu, "OperandMajorMode"): | |
| try: | |
| cute.nvgpu.OperandMajorMode = cute.nvgpu.tcgen05.OperandMajorMode | |
| except AttributeError: | |
| pass |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
include/flashinfer/attention/prefill.cuh (3)
657-672:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winInitialize RoPE frequencies over
HEAD_DIM_QK, notHEAD_DIM_VO.Asymmetric dispatch makes
NUM_MMA_D_QK != NUM_MMA_D_VOreachable.q_smem_inplace_apply_rotary/k_smem_inplace_apply_rotaryindexrope_frequp toNUM_MMA_D_QK / 2, but this initializer only fillsNUM_MMA_D_VO / 2; forqk=512, vo=256half the table is uninitialized.🐛 Proposed fix
- for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_VO / 2; ++mma_d) { + for (uint32_t mma_d = 0; mma_d < KTraits::NUM_MMA_D_QK / 2; ++mma_d) {🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@include/flashinfer/attention/prefill.cuh` around lines 657 - 672, The outer for loop in the init_rope_freq function is using KTraits::NUM_MMA_D_VO as the loop bound instead of KTraits::NUM_MMA_D_QK. This causes the rope_freq array to be only partially initialized when NUM_MMA_D_QK is larger than NUM_MMA_D_VO, since the calling functions q_smem_inplace_apply_rotary and k_smem_inplace_apply_rotary index the rope_freq array up to NUM_MMA_D_QK / 2. Change the loop condition from mma_d < KTraits::NUM_MMA_D_VO / 2 to mma_d < KTraits::NUM_MMA_D_QK / 2 to ensure the entire rope_freq table is properly initialized.
542-588:⚠️ Potential issue | 🟠 Major | ⚡ Quick winGuard rounded-up SF loader lanes before writing shared memory.
NUM_SF_ITERSis rounded up, so threads withflat_byte >= SF_TOTAL_BYTESexist. In theFLASHINFER_PAGED_V_SF_DESWIZZLEbranch, those lanes still execute*(uint32_t*)(sf_smem + flat_byte) = packed, which writes pastv_sf_smem.🛡️ Proposed fix
const uint32_t flat_uint32_idx = thread_id + k * THREADS_PER_CTA; const uint32_t flat_byte = flat_uint32_idx * 4; + if (flat_byte >= SF_TOTAL_BYTES) { + continue; + } @@ - const bool in_bounds = (flat_byte < SF_TOTAL_BYTES) && (kv_idx_base + sf_smem_row < kv_len); + const bool in_bounds = (kv_idx_base + sf_smem_row < kv_len);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@include/flashinfer/attention/prefill.cuh` around lines 542 - 588, In the FLASHINFER_PAGED_V_SF_DESWIZZLE branch of the code, the statement that writes packed data to shared memory `*reinterpret_cast<uint32_t*>(sf_smem + flat_byte) = packed;` is executed unconditionally for all threads, even those where in_bounds is false. This causes out-of-bounds writes past the v_sf_smem buffer for threads with flat_byte >= SF_TOTAL_BYTES. Guard this write operation with the in_bounds condition by wrapping the assignment in an `if (in_bounds)` check to prevent threads outside the valid bounds from writing to shared memory.
129-136:⚠️ Potential issue | 🟠 Major | ⚡ Quick winInclude FP4 SF smem in the
NUM_MMA_KVbudget.
SharedStorageQKVOaddsk_sf_smem/v_sf_smemproportional toCTA_TILE_KV, but the dispatch budgets only count KV data/repack/VO-split terms. On tight-smem GPUs this can select a largerNUM_MMA_KVthat fails the finalsizeof(SharedStorage*)check, even though a smaller KV tile would fit.🛠️ Suggested adjustment
+ constexpr uint32_t kFP4SfSmemPerMmaKV = + is_fp4_type_v<DTypeKV> ? (NUM_WARPS_KV * (HEAD_DIM_QK + HEAD_DIM_VO)) : 0u; constexpr uint32_t kKVSmemPerMmaKV = (kKVShared ? (HEAD_DIM_QK * 16 * NUM_WARPS_KV * sizeof(DTypeKV)) : ((HEAD_DIM_QK + HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeKV))) + + kFP4SfSmemPerMmaKV + (kUseRepack ? ((HEAD_DIM_QK > HEAD_DIM_VO ? HEAD_DIM_QK : HEAD_DIM_VO) * 16 * NUM_WARPS_KV * sizeof(DTypeQ)) : 0u);Apply the same term in the single, ragged, and paged dispatch budget blocks.
Also applies to: 2099-2104, 3516-3521, 3704-3710
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@include/flashinfer/attention/prefill.cuh` around lines 129 - 136, The dispatch budget calculations for single, ragged, and paged dispatch blocks do not account for the FP4 scale factor shared memory (k_sf_smem and v_sf_smem) that are added to SharedStorageQKVO when DTypeKV is an FP4 type. This causes the dispatch logic to select a NUM_MMA_KV value that may exceed actual shared memory capacity. Add the FP4 SF memory terms (proportional to CTA_TILE_KV multiplied by the scale factor dimensions divided by NVFP4_SF_VEC_SIZE) to the shared memory budget calculations in all three dispatch budget blocks wherever NUM_MMA_KV is constrained by the total SharedStorage size.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@flashinfer/decode.py`:
- Around line 1555-1561: The `out_head_dim` calculation for NVFP4 at lines
1555-1561 incorrectly determines whether to double the dimension based only on
whether `kv_cache_sf` is not None, rather than checking the actual dtype of the
KV cache. Modify the logic to derive `out_head_dim` from the KV cache dtype
(specifically checking if it is uint8) instead of relying on the presence of
`kv_cache_sf`. This ensures the output width is correctly computed based on the
actual data type and prevents miscomputation if `kv_cache_sf` is accidentally
provided for non-uint8 KV caches, maintaining consistency with the validation
logic.
In `@flashinfer/jit/attention/modules.py`:
- Around line 1915-1917: The gen_customize_batch_attention_module function at
line 1915 is missing the FP4 guard/flag pattern that exists in
gen_customize_batch_prefill_module (lines 1572-1681). Apply the same FP4 compile
gating logic from the batch-prefill module to the batch-attention module to
ensure that FP4-typed code generation is properly protected by FP4 compile
flags, making both paths consistent in how they handle FP4 customization.
In `@include/flashinfer/attention/prefill.cuh`:
- Around line 204-207: The CTA_TILE_Q pruning logic in the IsInvalid() method on
line 207 is inconsistent with FA2DetermineCtaTileQ's dispatcher logic. Update
the condition to check both HEAD_DIM_QK and HEAD_DIM_VO together (not just
HEAD_DIM_QK alone) to handle three distinct cases: true VO-split cases where
HEAD_DIM_VO >= 512, asymmetric cases where HEAD_DIM_QK >= 512 but HEAD_DIM_VO <=
256, and cases where HEAD_DIM_VO >= 512 but HEAD_DIM_QK < 512. Reorder the
pruning conditions to mirror the exact branch order used in FA2DetermineCtaTileQ
to ensure the dispatcher's selections are never incorrectly rejected by this
validity check.
---
Outside diff comments:
In `@include/flashinfer/attention/prefill.cuh`:
- Around line 657-672: The outer for loop in the init_rope_freq function is
using KTraits::NUM_MMA_D_VO as the loop bound instead of KTraits::NUM_MMA_D_QK.
This causes the rope_freq array to be only partially initialized when
NUM_MMA_D_QK is larger than NUM_MMA_D_VO, since the calling functions
q_smem_inplace_apply_rotary and k_smem_inplace_apply_rotary index the rope_freq
array up to NUM_MMA_D_QK / 2. Change the loop condition from mma_d <
KTraits::NUM_MMA_D_VO / 2 to mma_d < KTraits::NUM_MMA_D_QK / 2 to ensure the
entire rope_freq table is properly initialized.
- Around line 542-588: In the FLASHINFER_PAGED_V_SF_DESWIZZLE branch of the
code, the statement that writes packed data to shared memory
`*reinterpret_cast<uint32_t*>(sf_smem + flat_byte) = packed;` is executed
unconditionally for all threads, even those where in_bounds is false. This
causes out-of-bounds writes past the v_sf_smem buffer for threads with flat_byte
>= SF_TOTAL_BYTES. Guard this write operation with the in_bounds condition by
wrapping the assignment in an `if (in_bounds)` check to prevent threads outside
the valid bounds from writing to shared memory.
- Around line 129-136: The dispatch budget calculations for single, ragged, and
paged dispatch blocks do not account for the FP4 scale factor shared memory
(k_sf_smem and v_sf_smem) that are added to SharedStorageQKVO when DTypeKV is an
FP4 type. This causes the dispatch logic to select a NUM_MMA_KV value that may
exceed actual shared memory capacity. Add the FP4 SF memory terms (proportional
to CTA_TILE_KV multiplied by the scale factor dimensions divided by
NVFP4_SF_VEC_SIZE) to the shared memory budget calculations in all three
dispatch budget blocks wherever NUM_MMA_KV is constrained by the total
SharedStorage size.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 97e44c80-6c09-4648-b2fb-a58e620e2f3f
📒 Files selected for processing (25)
.github/workflows/nightly-release.yml.github/workflows/release.ymlcsrc/batch_decode.cucsrc/batch_prefill.cucsrc/batch_prefill_customize_config.jinjacsrc/batch_prefill_paged_kernel_inst.jinjacsrc/batch_prefill_ragged_kernel_inst.jinjadocs/installation.rstflashinfer/cute_dsl/utils.pyflashinfer/decode.pyflashinfer/gemm/gemm_base.pyflashinfer/jit/attention/modules.pyflashinfer/jit/attention/utils.pyflashinfer/jit/utils.pyflashinfer/mla/_core.pyflashinfer/prefill.pyflashinfer/xqa.pyinclude/flashinfer/attention/persistent.cuhinclude/flashinfer/attention/prefill.cuhinclude/flashinfer/attention/scheduler.cuhinclude/flashinfer/page.cuhinclude/flashinfer/utils.cuhinclude/flashinfer/vec_dtypes.cuhtests/gemm/test_mm_fp4.pytests/jit/test_attention_utils.py
| + generate_sf_stride_setter_lines( | ||
| get_sf_stride_tensor_names(additional_tensor_names), prefix="params[i]." | ||
| ) |
There was a problem hiding this comment.
Mirror FP4 compile gating in batch-attention JIT path.
Line [1915] adds NVFP4 SF-stride wiring for gen_customize_batch_attention_module, but this generator still lacks the FP4 guard/flag pattern used in gen_customize_batch_prefill_module (Line [1572]-Line [1681]). That leaves an inconsistent path where FP4 KV customization can emit FP4-typed code without enabling FP4 compile flags.
Suggested fix
def gen_customize_batch_attention_module(
@@
):
+ require_fp4_kv_cache = dtype_map_kv[dtype_kv] == "__nv_fp4x2_e2m1"
+ if require_fp4_kv_cache:
+ missing_sf_tensors = [
+ name
+ for name in ("maybe_k_cache_sf", "maybe_v_cache_sf")
+ if name not in additional_tensor_names
+ ]
+ if missing_sf_tensors:
+ raise ValueError(
+ "NVFP4 KV batch-attention JIT modules require scale-factor tensors "
+ f"{missing_sf_tensors}; pass maybe_k_cache_sf and maybe_v_cache_sf "
+ "as additional tensors."
+ )
@@
- return gen_jit_spec(
+ extra_cuda_cflags = ["-DFLASHINFER_ENABLE_PROFILER"] if use_profiler else []
+ if require_fp4_kv_cache:
+ extra_cuda_cflags += common_nvcc_flags
+
+ return gen_jit_spec(
uri,
source_paths,
- extra_cuda_cflags=["-DFLASHINFER_ENABLE_PROFILER"] if use_profiler else [],
+ extra_cuda_cflags=extra_cuda_cflags,
)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/jit/attention/modules.py` around lines 1915 - 1917, The
gen_customize_batch_attention_module function at line 1915 is missing the FP4
guard/flag pattern that exists in gen_customize_batch_prefill_module (lines
1572-1681). Apply the same FP4 compile gating logic from the batch-prefill
module to the batch-attention module to ensure that FP4-typed code generation is
properly protected by FP4 compile flags, making both paths consistent in how
they handle FP4 customization.
| static constexpr bool IsInvalid() { | ||
| // The first clause prunes (CTA_TILE_Q, head_dim) pairs FA2DetermineCtaTileQ | ||
| // never selects: {16, 32} for head_dim_vo >= 512, {16, 64, 128} otherwise. | ||
| return ((HEAD_DIM_VO >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (NUM_MMA_D_VO < 4) || | ||
| // never selects: {16, 32} for head_dim_qk >= 512, {16, 64, 128} otherwise. | ||
| return ((HEAD_DIM_QK >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (NUM_MMA_D_VO < 4) || |
There was a problem hiding this comment.
Keep CTA pruning consistent with VO-vs-QK tile selection.
Line 207 allows CTA_TILE_Q == 32 whenever HEAD_DIM_QK >= 512, but FA2DetermineCtaTileQ only allows 32 for true HEAD_DIM_VO >= 512 VO-split cases and forces 16 for asymmetric QK >= 512, VO <= 256. It also rejects CTA_TILE_Q == 32 for VO >= 512, QK < 512, even though the dispatcher can select it. Mirror the dispatcher’s branch order here.
🐛 Proposed fix
- return ((HEAD_DIM_QK >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (NUM_MMA_D_VO < 4) ||
+ constexpr bool invalid_cta_tile_q =
+ (HEAD_DIM_VO >= 512)
+ ? (CTA_TILE_Q > 32)
+ : (HEAD_DIM_QK >= 512 ? (CTA_TILE_Q != 16) : (CTA_TILE_Q == 32));
+ return (invalid_cta_tile_q || (NUM_MMA_D_VO < 4) ||📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| static constexpr bool IsInvalid() { | |
| // The first clause prunes (CTA_TILE_Q, head_dim) pairs FA2DetermineCtaTileQ | |
| // never selects: {16, 32} for head_dim_vo >= 512, {16, 64, 128} otherwise. | |
| return ((HEAD_DIM_VO >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (NUM_MMA_D_VO < 4) || | |
| // never selects: {16, 32} for head_dim_qk >= 512, {16, 64, 128} otherwise. | |
| return ((HEAD_DIM_QK >= 512 ? (CTA_TILE_Q > 32) : (CTA_TILE_Q == 32)) || (NUM_MMA_D_VO < 4) || | |
| static constexpr bool IsInvalid() { | |
| // The first clause prunes (CTA_TILE_Q, head_dim) pairs FA2DetermineCtaTileQ | |
| // never selects: {16, 32} for head_dim_qk >= 512, {16, 64, 128} otherwise. | |
| constexpr bool invalid_cta_tile_q = | |
| (HEAD_DIM_VO >= 512) | |
| ? (CTA_TILE_Q > 32) | |
| : (HEAD_DIM_QK >= 512 ? (CTA_TILE_Q != 16) : (CTA_TILE_Q == 32)); | |
| return (invalid_cta_tile_q || (NUM_MMA_D_VO < 4) || |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@include/flashinfer/attention/prefill.cuh` around lines 204 - 207, The
CTA_TILE_Q pruning logic in the IsInvalid() method on line 207 is inconsistent
with FA2DetermineCtaTileQ's dispatcher logic. Update the condition to check both
HEAD_DIM_QK and HEAD_DIM_VO together (not just HEAD_DIM_QK alone) to handle
three distinct cases: true VO-split cases where HEAD_DIM_VO >= 512, asymmetric
cases where HEAD_DIM_QK >= 512 but HEAD_DIM_VO <= 256, and cases where
HEAD_DIM_VO >= 512 but HEAD_DIM_QK < 512. Reorder the pruning conditions to
mirror the exact branch order used in FA2DetermineCtaTileQ to ensure the
dispatcher's selections are never incorrectly rejected by this validity check.
|
A heads-up for reviewers before you spend time here: I'm extending this work to add multimodal support for Gemma 3/4, which I erroneously omitted from the initial version. Gemma 3/4 are natively multimodal, and serving image (and, for Gemma 4, audio) inputs through the NVFP4 VO-split prefill requires the model's span-level bidirectional image masking to flow through the kernel's custom-mask path. Without it the contribution only covers text prompts, which isn't a complete Gemma 3/4 claim. I'm validating the multimodal path now and will push the update shortly. I'd appreciate holding final review until it lands — apologies for the churn. |
|
Update — the multimodal support has landed, on the vLLM side: vllm-project/vllm#46329 (commit "Gemma 3/4 multimodal: NVFP4 KV mm-prefix span masking on FA2"). It serves image-token spans bidirectionally via FlashInfer's existing FA2 packed Validated reading images back under NVFP4 KV, matching bf16: |
Address review on flashinfer-ai#3684: the paged-decode output allocation doubled out_head_dim whenever kv_cache_sf was non-None. Gate it on the KV cache actually being uint8-packed (the NVFP4 layout that stores VO at half width), so a stray scale-factor tensor on a non-uint8 cache fails the shape check instead of silently allocating a mis-sized output.
|
Thanks for the reviews. Pushed one fix (7f549fe) and notes on the rest:
|
|
Pushed a follow-up commit ( Found this validating the kernel under prefix caching. NVFP4 split-KV (flash-decoding) corrupts reads when a short query attends a long cached KV ( Worth flagging for the companion vLLM PR and anyone else consuming this kernel: since the fix lives in |
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
flashinfer/prefill.py (2)
1372-1374: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winGate NVFP4 output-width doubling on the packed KV dtype, not
kv_cache_sf.A stray
kv_cache_sfwith non-uint8KV still doubles the output width in prefill paths, silently allocating/checking the wrong shape. Match the decode-side fix by using the V-cache dtype as the packed-NVFP4 signal.Suggested shape guard
- out_head_dim = ( - v_cache.shape[-1] * 2 if kv_cache_sf is not None else v_cache.shape[-1] - ) + out_head_dim = ( + v_cache.shape[-1] * 2 if v_cache.dtype == torch.uint8 else v_cache.shape[-1] + )- out_head_dim = v.shape[-1] * 2 if kv_cache_sf is not None else v.shape[-1] + out_head_dim = v.shape[-1] * 2 if v.dtype == torch.uint8 else v.shape[-1]Also applies to: 2457-2459, 3511-3513
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/prefill.py` around lines 1372 - 1374, The NVFP4 output-width logic in the prefill path is gated on kv_cache_sf, which can incorrectly double the shape even when the KV cache is not packed. Update the out_head_dim calculation in the prefill code paths (and the matching decode-side logic if needed) to use the V-cache dtype, specifically checking whether v is packed NVFP4/uint8, so output width doubling only happens for packed KV storage and not for unrelated kv_cache_sf presence.
1975-1981: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMirror this device fix in the ragged custom-mask path.
This paged path now satisfies
segment_packbits’ device requirement, but the raggedplan()path still passes CPU-derivedmask_indptrdirectly whencustom_maskis on GPU, leaving the same crash for ragged multimodal/custom-mask callers.Suggested matching fix for the ragged path
packed_custom_mask, mask_indptr = segment_packbits( custom_mask.contiguous().view(-1), - mask_indptr, + mask_indptr.to(custom_mask.device), bitorder="little", )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/prefill.py` around lines 1975 - 1981, The ragged custom-mask path in plan() still passes mask_indptr from the CPU side directly into segment_packbits, which can violate the device requirement when custom_mask is on GPU. Mirror the same fix used in the paged path: in the custom-mask handling inside plan(), ensure mask_indptr is moved to custom_mask.device before calling segment_packbits, so the ragged multimodal/custom-mask flow matches the device expectations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@flashinfer/prefill.py`:
- Around line 1372-1374: The NVFP4 output-width logic in the prefill path is
gated on kv_cache_sf, which can incorrectly double the shape even when the KV
cache is not packed. Update the out_head_dim calculation in the prefill code
paths (and the matching decode-side logic if needed) to use the V-cache dtype,
specifically checking whether v is packed NVFP4/uint8, so output width doubling
only happens for packed KV storage and not for unrelated kv_cache_sf presence.
- Around line 1975-1981: The ragged custom-mask path in plan() still passes
mask_indptr from the CPU side directly into segment_packbits, which can violate
the device requirement when custom_mask is on GPU. Mirror the same fix used in
the paged path: in the custom-mask handling inside plan(), ensure mask_indptr is
moved to custom_mask.device before calling segment_packbits, so the ragged
multimodal/custom-mask flow matches the device expectations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: cb547752-7c96-4605-b9a6-9f1924cc0001
📒 Files selected for processing (2)
flashinfer/prefill.pytests/attention/test_nvfp4_attention_sm120.py
…ound The previous rationale claimed a split-KV chunk boundary can land in the middle of a 16-element scale block. That mechanism is wrong: NVFP4 scale factors group 16 consecutive head-dim elements of a single token, while split-KV partitions the token axis, so a split boundary never slices a scale block. Reword the gate docstring, call-site comments and test docstring to state what is actually known -- corruption was observed empirically when qo_len << kv_len (decode / prefix-cache extend) and disappears with split-KV disabled at no measured decode throughput cost -- and cite the interaction between small per-split KV chunks and the 1-byte-KV NUM_MMA_KV tile floor as an unconfirmed hypothesis rather than fact. Also drops project-specific NOTE tags from nearby comments. Comment-only; no behavior change. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv. Signed-off-by: Jetha Chan <jethachan@gmail.com>
run() fills declared-but-unprovided JIT scalars from a fixed mapping and raised a bare KeyError when a JIT module declares a scalar the mapping does not know how to derive. Raise a ValueError naming the scalar and listing the derivable set instead. Also drop the redundant max(0, ...) clamp on the provided-scalar count: prepare_jit_additional_args always returns at least one entry per declared tensor name, so the excess over the tensor-name count cannot be negative; a comment records that invariant. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv. Signed-off-by: Jetha Chan <jethachan@gmail.com>
…ll-VO For head_dim_qk >= 512 with head_dim_vo < 512 (asymmetric heads), FA2DetermineCtaTileQ always returns CTA_TILE_Q=16, so the CTA_TILE_Q=32 instantiation in the paged/ragged kernel-instantiation lists can never be dispatched; drop it and instantiate only CTA16 for that shape class. The dispatch macro case-32 arm still references the symbol, which stays unresolved in the module exactly like the pre-existing never-selected CTA_TILE_Q=32 of symmetric head_dim < 512 modules (verified: shipped modules carry it as an undefined, lazily-bound symbol). Also updates the KernelTraits::IsInvalid comment to describe the three-way CTA_TILE_Q selection. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv. Signed-off-by: Jetha Chan <jethachan@gmail.com>
…ual-stride rejection
Completes the support-or-reject-explicitly contract for unequal K/V
strides in tree: every consumer reachable from the updated entry points
must either support independently-strided K/V pools or reject them
loudly.
Support half: asymmetric (head_dim_qk != head_dim_vo) NVFP4 paged
prefill over (512,256) and (256,128) x page_size {1,16} x num_kv_heads
{2,8}, causal. K/V pools and their scale-factor tensors are separately
allocated with genuinely different stride families; bf16 sources are
quantized with the in-tree NVFP4 KV quantization kernel and the FA2
output is checked against a float32 reference attention computed on
nvfp4_kv_dequantize_paged output, so kernel and reference consume the
exact same quantized bytes (dequantization oracle, not a requantized
approximation).
Reject half: the CUDA-core decode entry point
(BatchDecodeWithPagedKVCacheRun) addresses both K and V through a
single set of (K) strides, so its restored ICHECK must fire on K/V
pools whose stride families differ instead of silently misaddressing V.
A positive control with identically padded (equal-stride,
non-contiguous) pools runs and matches the reference, proving the
negative case fails because of the stride inequality and not the padded
allocation.
Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.
Signed-off-by: Jetha Chan <jethachan@gmail.com>
Regression test for the FA2DetermineCtaTileQ shared-memory probe at head dims that reach it today: plan()/JIT do not validate head dims, so (qk, vo) = (448, 256) under pos_encoding_mode NONE is accepted, and at 2-byte KV its short-q 1x4-layout cost (104448 bytes) exceeds the 101376-byte per-block opt-in limit of 99KB parts. The test computes the expected tile from the device's actual opt-in limit (so the assertion is exact on every architecture), asserts the planned cta_tile_q via PrefillPlanInfo (the same technique as test_fp8_prefill.py), and for 2-byte KV runs the kernel against an exact float32 reference: on 99KB parts this proves the probe fires and the CTA64 fallback keeps the configuration dispatchable where the CTA16 dispatch would exceed the per-block limit, and on larger-smem parts it proves the CTA16 selection runs. For 1-byte KV the assertion is plan-level: the FA2 1-byte KV producers require head_dim to be a multiple of 128 elements (the 128-bit-per-lane load loop steps NUM_MMA_D by 8, and the k128B swizzle needs an 8-aligned upcast stride), so no currently-runnable 1-byte configuration reaches the flipped CTA64->CTA16 region -- the pin locks the documented planner behavior for when one does. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv. Signed-off-by: Jetha Chan <jethachan@gmail.com>
…trides
The asymmetric NVFP4 stride test does not execute the produce_v fix in
page_produce_kv_on_the_fly: that producer runs only under
USE_KV_SHARED_SMEM, which excludes FP4 and requires HEAD_DIM_QK ==
HEAD_DIM_VO, so the NVFP4 asymmetric path takes the prefetched
thr_local_kv_offset_{k,v} arrays instead.
Add the configuration that does execute it: 16-bit KV at
(qk, vo) = (512, 512), where USE_KV_SHARED_SMEM holds for both CTA
tiles the planner can pick (static reasoning from prefill.cuh):
CTA_TILE_Q=16 for short q (NUM_WARPS_KV=4; NUM_MMA_D_VO=32 % 4 == 0)
and CTA_TILE_Q=32 for long q (kLargeHeadWarpSplit: NUM_WARPS_KV=2;
32 % 2 == 0), so USE_VO_SPLIT -- and with fp16's equal head dims,
USE_KV_SHARED_SMEM -- is true either way; the qo_len parametrization
covers both tiles and the kv_layout parametrization covers NHD and
HND. K and V pools are views of differently padded parent tensors
(identical logical shapes, unequal stride families, mirroring the
decode negative test's construction), so
get_paged_kv_offset_for_logical_row<produce_v=true> must route V rows
through the V strides: with the fix reverted, the V loads walk K's
stride family and the output diverges from the exact float32
reference, which is how this test was validated to catch the bug it
pins. The configuration is SM80+, so it runs on the standard CI
runners.
Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.
Signed-off-by: Jetha Chan <jethachan@gmail.com>
The (448, 256) CtaTileQ smem-probe test parametrizes over kv_dtype in
{float16, float8_e4m3fn}. On pre-SM100 GPUs the FP8 (1-byte) parametrization
errors before reaching the tile assertion: _fa2_head_dim_nvcc_flags restricts
non-NVFP4 1-byte large-head modules to major versions [10, 11, 12], so the JIT
spec-gen inside plan() raises "No supported CUDA architectures found for major
versions [10, 11, 12]". skip_if_head_dim_unsupported only gates the 16-bit path,
so it misses this.
Add a dtype-aware skip mirroring the module gate, and narrow the docstring
wording "exact on every architecture" -> "exact on every supported
architecture". The fp16 parametrization is unaffected (2-byte fallback, SM80+).
Addresses review feedback on flashinfer-ai#3684 from @qsang-nv.
Signed-off-by: Jetha Chan <jethachan@gmail.com>
gen_customize_batch_prefill_module now requires the scale-factor tensors (maybe_k_cache_sf / maybe_v_cache_sf) as additional inputs whenever the KV dtype resolves to NVFP4, raising ValueError otherwise. The host-side test_customize_batch_prefill_nvfp4_large_head_uses_prefill_flags still called the generator with empty additional-tensor lists, so it tripped that ValueError before reaching either flag assertion and failed on every arch (it never touches the GPU). Pass the two uint8_t SF tensors, mirroring the production caller, so generation completes and the assertions run: _fa2_prefill_head_dim_nvcc_flags emits sm_86 (allow_nvfp4_sm8_large_head), and the plain _fa2_head_dim_nvcc_flags still restricts to [10,11,12] and raises. Signed-off-by: Jetha Chan <jethachan@gmail.com>
…ib.suppress) Signed-off-by: Jetha Chan <jethachan@gmail.com> Co-authored-by: TechPrototyper <tech@smartlogics.net>
|
@qsang-nv rebased onto current |
|
@qsang-nv the CI run (30594351621) has been stuck ~1.5 days on |
|
Update — run 30594351621 completed: H100 ran the full suite green (attention/nvfp4/gemm/moe); the only red is |
|
/bot run |
|
[FAILED] Pipeline #60789494 — 13/18 executed test jobs passed Compared with nightly #60712014. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 5/6 passed
Failure detailsNew relative to nightly (attribution uncertain)
Timeouts, infrastructure, or incomplete jobs
|
|
Since the internal CI's RTX 5090 rows were reported as infrastructure failures, here is an independent datapoint from consumer hardware, at the current PR head. Environment. RTX 5090 (sm_120), driver 595.71.05, CUDA 13.0, torch 2.11.0+cu130, containerized Linux. flashinfer checked out at Results — zero failures, zero errors across all three suites:
Skips are the expected non-sm120 paths (SM100/hopper-only kernels, unsupported dtypes/head-dims). No test hit the 1800 s timeout. Two environment notes for reproducibility:
junit XMLs for all three suites available on request. |
|
@flashinfer-bot run |
|
@TechPrototyper thank you for this. I'd been meaning to rent a 5090 on RunPod and run exactly this myself and just haven't had the time — you saved me a real chunk of work. The Your run was at Environment. RTX 5090 (sm_120), driver 580.126.09, CUDA 13.0 (nvcc 13.0.88), torch 2.9.1+cu130, editable/JIT install at 19,202 tests executed, zero failures, zero errors:
The last row is the two attention test files this PR modifies, run without the fp4 filter, since the PR also edits shared headers (
Attention skips are arch gating rather than the feature being passed over — 15,584 @qsang-nv — the junit XMLs available if useful. |
6 participants
Summary
Add an asymmetric VO-split NVFP4 paged-prefill path to the FA2 attention kernels so consumer/SoC Blackwell (SM120 / SM121) can serve models whose global-attention layers are 512-wide — specifically Gemma 4 — with a 4-bit (NVFP4) KV cache.
Upstream NVFP4 paged prefill is symmetric
head_dim ≤ 256only. Gemma 4's full-attention layers arehead_dim_qk = 512, head_dim_vo = 256(value reuses key), which hits three walls: the equal-K/V-stride host check, the VO-splitstatic_assert(!is_fp4), and the output-width hardcode. This PR adds the FP4 V-load + explicit scale-factor strides + asymmetric K/V offsets so the kernel serves the asymmetric(qk=512, vo=256)shape directly, and keysCTA_TILE_Qonhead_dim_qkso the dispatch is feasible at 512.What's in it
[all-data | all-SF]contiguous scale-factor layout (the trtllm-gen per-page swizzle doesn't commute with head-dim slicing), and the asymmetric K/V byte/scale offsets.CTA_TILE_Qkeyed onhead_dim_qkfor the asymmetric(QK=512, VO=256)case (the dispatcher otherwise picks an infeasible tile /max_mma_kv: 0).plan(): movemask_indptrto the custom mask's device beforesegment_packbits— a standalone correctness fix (the paged + ragged wrappers passed a CPUmask_indptr).Validation
qk=512 / vo=256paged NVFP4 prefillmax_abs_err ≈ 0.0047(the FP4 e2m1 rounding floor itself is ~0.0156); output correctly sized byvo=256; symmetricqk=vo=256/128cases unregressed; multi-tile seq 512/1024 + causal all sane.bf16. (vLLM side: companion PR linked below.)Notes
mask_indptrdevice fix into its own PR if preferred.NVFP4 split-KV correctness (added)
While validating this kernel under SGLang's radix cache (and as a heads-up for vLLM's
automatic prefix caching), I found a real correctness bug in the NVFP4 split-KV
path -- distinct from the VO-split work, and shared by every consumer of this kernel.
Symptom. With a 4-bit KV cache and prefix reuse on, retrieval cliffs: a Gemma-4
E2B/E4B needle is retrieved perfectly up to ~550 tokens of reused-prefix context,
then collapses (degenerate repetition /
<pad>) from ~680 tokens on.bf16KV has nosuch cliff. The tell:
--disable-radix-cache(any cold, contiguous prefill) makes itvanish.
Root cause. Not page-table fragmentation, not calibration (both ruled out by an
instrumented A/B). It's the extend geometry -- a short query window attending a
long cached KV (
qo_len << kv_len) -- which makes the FA2 scheduler picksplit-KV (flash-decoding). NVFP4 KV is packed
uint8with a per-16-element FP8block scale;
kv_chunk_sizeis not aligned to those blocks, so a split boundary canland mid-block, and the small per-split chunk also trips the 1-byte-KV
NUM_MMA_KVtile floor. A position-aligned dense-cache trace shows the final query token's layer-0
attention output drift ~0.1-1% under split-KV vs the cold full-prefill for an
identical query -- small, but it compounds over ~30 layers until small models lose
the needle. A full-prefill (
qo_len == kv_len) never splits, which is why dense-prefilltests stay green while prefix caching / long-context decode break.
Fix (this push). Gate split-KV off when
kv_data_typeis NVFP4 (uint8/float4_e2m1fn_x2) in both the paged and ragged prefillplan(); FP8 / 16-bit KV areuntouched. ~5 lines, provably correct, at a small decode/short-query parallelism cost
for 4-bit KV only. Adds
test_nvfp4_split_kv_gate_dtype_logic(hardware-freedtype-classification guard).
Verified. Gemma-4 E2B/E4B radix-on retrieval holds to 1448 tokens (==
bf16) withthe gate vs cliffing at ~600 without it -- on both sm120 (RTX PRO 6000) and
sm121 (GB10 / DGX Spark). The full Gemma-4 multimodal matrix (text/image/audio,
nvfp4 vs bf16) is green across E2B/E4B/12B/26B-A4B/31B with radix caching on.
Follow-up. A perf-preserving version keeps flash-decoding by making FP4 split points
16-token-aligned (respecting the scale blocks) and clearing the
NUM_MMA_KVfloor perchunk -- left as a follow-up so the correctness fix can land now.
Summary by CodeRabbit
Release Notes
b12xfirst.lsevalues in NVFP4 SM120 attention and corrected KV attention tiling logic to use Q/K dimensions.12.1a).Build-size note (per review)
DISPATCH_GQA_GROUP_SIZEgains agroup_size == 6case (needed for Gemma-4-class GQA ratios), which adds one instantiation per macro user - roughly +20% compile time/binary size for the translation units that expand this macro.