Conversation
This was referenced Jul 26, 2026
Closed
…120) Implements the e2m1 (NVFP4) KV-cache read path for the FA2 paged-prefill and decode kernels on consumer Blackwell (SM120): - FP4 E2M1 decode vector casts and KV scale-stride plumbing through the kernels - pass JIT scalar args for NVFP4 paged prefill; disambiguate FP4 KV JIT module names so distinct dtype configs don't collide in the JIT cache - size the attention output from V rather than Q (asymmetric QK vs VO head dims) - allow GQA group size 6 in decode dispatch - guard the FP4 paged-prefill JIT flags and alias the CUTLASS DSL OperandMajorMode for the FP4 path Signed-off-by: Jetha Chan <jethachan@gmail.com> (cherry picked from commit 77758ee)
FA2DetermineCtaTileQ picked a 1x4 warp layout without accounting for the real shared-memory budget, yielding max_mma_kv=0 at HEAD_DIM_QK=512 under real GQA with 1-byte (fp8/NVFP4) KV. Make the warp-layout choice smem-aware, and cleanly reject KV tiles that do not fit real shared memory instead of failing later. Signed-off-by: Jetha Chan <jethachan@gmail.com> (cherry picked from commit 23ce85a)
…ackbits Fixes a device mismatch when building packed custom masks for bidirectional (multimodal-prefix) attention. Signed-off-by: Jetha Chan <jethachan@gmail.com> (cherry picked from commit ee9301f)
Signed-off-by: Jetha Chan <jethachan@gmail.com> (cherry picked from commit 131bd2a)
…eads
The VO-split rederive left three dispatch gates keyed on head_dim_vo while
FA2DetermineCtaTileQ selects CTA_TILE_Q from the QK dim. For asymmetric
Gemma-4 heads (qk=512, vo=256) VO-split does not engage (NUM_MMA_D_VO==16),
so the full-D path applies and only CTA_TILE_Q=16 fits the 256-register
o_frag wall. Align all four sites:
- utils.cuh FA2DetermineCtaTileQ: VO>=512 -> {16,32}; QK>=512 & VO<512 -> 16
- prefill.cuh KernelTraits::IsInvalid first clause: key on HEAD_DIM_QK
- batch_prefill_{paged,ragged}_kernel_inst.jinja: instantiate {16,32} when head_dim_qk>=512
Validated NVFP4 paged prefill (sm120, RTX): max_abs_err qk256/vo256=0.0048,
qk128=0.0046, qk512/vo256=0.0047 (seq128) / 0.0044 (seq512) / 0.0042
(seq1024) / 0.0166 causal. No NaN; output sized by VO.
Signed-off-by: Jetha Chan <jethachan@gmail.com>
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. Signed-off-by: Jetha Chan <jethachan@gmail.com>
NVFP4 paged KV is packed uint8 with a per-16-element FP8 block scale. Split-KV (flash-decoding) chunks the KV range by kv_chunk_size, which is not aligned to the 16-element scale blocks; a chunk boundary landing mid-block -- plus the small per-split chunk tripping the 1-byte-KV NUM_MMA_KV tile floor -- corrupts the dequantized reads. It only surfaces when a short query attends a long KV (qo_len << kv_len), i.e. decode and prefix-cache extend, so dense full-prefill tests miss it while prefix caching / long-context decode break. Gate split-KV off when kv_data_type is NVFP4 (uint8 / float4_e2m1fn_x2) in the paged and ragged prefill plan(). FP8 / 16-bit KV are unaffected. Verified on Gemma-4 E2B/E4B: radix-on retrieval cliffs at ~600 tokens of reused-prefix context without the gate, holds to 1448 (== bf16) with it -- on both sm120 (RTX PRO 6000) and sm121 (GB10 / DGX Spark). Follow-up: make the FP4 split path scale-block-aware (16-token-aligned chunk boundaries + NUM_MMA_KV floor) to restore flash-decoding parallelism for long-context NVFP4 decode. Signed-off-by: Jetha Chan <jethachan@gmail.com>
Address CodeRabbit review on flashinfer-ai#3684 (5 findings): prefill.py: - Gate NVFP4 out_head_dim doubling on `v.dtype == torch.uint8` (the packed storage), not just `kv_cache_sf is not None`, at all three prefill sites. A stray scale-factor tensor on a non-uint8 cache no longer silently doubles the output width. Mirrors the decode-side guard. - Ragged custom-mask plan(): move mask_indptr to custom_mask.device before segment_packbits, mirroring the paged-path fix. mask_indptr inherits qo_indptr's device (often CPU) while custom_mask is on GPU. prefill.cuh: - init_rope_freq: bound the fill loop on NUM_MMA_D_QK/2, not NUM_MMA_D_VO/2. rope_freq is sized [NUM_MMA_D_QK/2][4] and the rotary appliers index it up to NUM_MMA_D_QK/2; the asymmetric VO-split dispatch (qk=512, vo=256) made NUM_MMA_D_QK != NUM_MMA_D_VO reachable, leaving the upper half uninitialized for in-kernel-RoPE callers. - page_produce_kv_sf deswizzle branch: guard the unconditional smem store on flat_byte < SF_TOTAL_BYTES (the rounded-up NUM_SF_ITERS leaves over-range lanes). Live cp_async path was already predicated. - NUM_MMA_KV occupancy budget (all three dispatchers): add the FP4 K/V scale-factor smem term, which scales with CTA_TILE_KV. Prevents the budget from over-selecting a tile whose SharedStorage then exceeds the smem limit. Signed-off-by: Jetha Chan <jethachan@gmail.com>
The fwd entry point passes the caller's uninitialized lse (torch.empty) straight to run_mha_fwd. The kernel writes `out` for every query row but does not guarantee writing every (batch, head, seq) entry of lse, so lse can read back whatever the allocator handed out — observed as flaky NaNs that pass on a clean allocation and fail when a prior test dirtied the pages. Zero lse before launch, mirroring the existing seq_len==0 branch. Found smoke-testing flashinfer-ai#3684 on sm120 (RTX PRO 6000): the nvfp4 sm120 accuracy test's lse NaN-check flaked only when run after the paged-prefill nvfp4 test. Signed-off-by: Jetha Chan <jethachan@gmail.com>
The old assertion expected a literal "constexpr bool REQUIRE_FP4_KV_CACHE = true;" that no code path ever emitted; check the #error guard and static_assert the config template actually renders for FP4 KV. Signed-off-by: Jetha Chan <jethachan@gmail.com>
BatchDecodeWithPagedKVCacheRun accepts independent K/V stride arrays since the paged_kv_t independent-stride constructor was introduced, but the FA2 decode kernel (decode.cuh) still computes a single protective_get_kv_offset per row -- from the K strides -- and uses it to address both k_data and v_data, so asymmetric strides would silently read V at K offsets. Restore the host-side per-dimension stride-equality check that guarded this before, with an error message naming the limitation. Restoring the check (rather than making the decode kernels V-stride-aware) is the minimal honest fix: no supported path needs asymmetric decode today -- asymmetric NVFP4 (VO-split) decode rides the prefill wrapper, and symmetric caches have equal K/V strides by construction. If asymmetric decode is needed later, decode.cuh should grow separate K/V offsets the way the prefill path already has. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv. Signed-off-by: Jetha Chan <jethachan@gmail.com>
…oducer page_produce_kv_on_the_fly computed both K and V global-memory offsets via get_paged_kv_offset_for_logical_row, which called protective_get_kv_offset -- the K strides -- so the shared-KV (on-the-fly) path addressed V rows with K offsets once K and V may carry independent strides. Thread produce_v through get_paged_kv_offset_for_logical_row and select protective_get_v_offset vs protective_get_k_offset with if constexpr, mirroring what the prefetched-offset path already does with its separate K/V offset arrays. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv and @lesj0610. Signed-off-by: Jetha Chan <jethachan@gmail.com>
Revert the extension of the b12x auto-preference from SM120-only to all SM12x, restoring the upstream heuristic: SM121 (GB10) intentionally keeps cutlass/cudnn first because they are faster there in most cases, while b12x remains available on SM121 as an explicit backend. If flipping the SM121 default to b12x is warranted, it should come as a separately justified PR with benchmarks rather than ride along here. Also removes the heuristic test that asserted the flipped ordering. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv. Signed-off-by: Jetha Chan <jethachan@gmail.com>
Three corrections to the shared-memory feasibility probe in FA2DetermineCtaTileQ: * Probe cudaDevAttrMaxSharedMemoryPerBlockOptin instead of the per-SM attribute. The kernel dispatcher bounds max_smem_per_threadblock by the per-block opt-in limit, so the planner now checks the same limit the dispatch-side "even the smallest KV tile exceeds shared memory" guard enforces. On parts where the two attributes differ (per-SM 102400 vs opt-in 101376 on SM86/89/120-class devices), configurations in the gap -- e.g. (qk, vo) = (432, 256) at 2-byte KV, 101888 bytes -- previously passed the probe and then failed dispatch; they now take the CTA64 fallback instead. * Size the KV step by the actual KV dtype width instead of assuming 2 bytes. sizeof(DTypeKV) is threaded from the batch-prefill plan and workspace-size entry points (where DTypeKV is in scope) through PrefillPlan / PrefillPlanWorkspaceSize / PrefillPlanImpl / PrefillSplitQOKVIndptr as a defaulted trailing parameter, and passed directly in single-prefill dispatch; callers that do not know the KV dtype keep the previous conservative 2-byte assumption. * Correct the comment. The old (512, 256) example was unreachable here (head_dim_qk >= 512 returns CTA16 before the probe), but the probe is not merely forward-looking: plan()/JIT do not validate head dims, so within this branch head_dim_qk may be any multiple of 16 below 512 under pos_encoding_mode NONE, and (qk, vo) = (448, 256) at 2-byte KV (104448 bytes) exceeds the 101376-byte opt-in limit of 99KB parts -- the probe fires today and the CTA64 fallback keeps the configuration dispatchable. For the same dims at 1-byte KV the corrected sizing flips the selection (59392 bytes fits, so CTA16 replaces the CTA64 the 2-byte assumption forced): the kv_dtype_bytes change alters tile selection for reachable configurations rather than being a no-op on current hardware. Both behaviors are pinned by test_batch_prefill_paged_cta_tile_q_smem_probe_qk448_vo256. Addresses review feedback on flashinfer-ai#3684 from @qsang-nv, including the round-2 correction of this commit's earlier "no reachable configuration triggers the fallback" claim, which was wrong for the reasons above. Signed-off-by: Jetha Chan <jethachan@gmail.com>
…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>
Previously, _nvfp4_kv_requires_disabled_split_kv() returned True for NVFP4 KV cache, disabling split-KV (flash-decoding) as an empirical workaround for corrupted outputs when short queries attend long KV ranges (decode / prefix-cache extend). This change returns False unconditionally, re-enabling split-KV for NVFP4. This is critical for MTP (Multi-Token Prediction) verification performance on SM120, where short query_len=3 must attend to long KV sequences and needs split-KV parallelism across KV heads. The original corruption has been resolved by upstream fixes to the NVFP4 paged KV layout (4-D HND format with correct strides) and scale-factor handling.
ch2lab
force-pushed
the
fix/re-enable-split-kv-nvfp4
branch
from
July 28, 2026 10:08
5ea4078 to
277325f
Compare
jethac
force-pushed
the
nvfp4-vosplit-rederive
branch
from
July 31, 2026 00:43
e79f21b to
fa89401
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
_nvfp4_kv_requires_disabled_split_kv()previously returnedTruefor NVFP4 KV cache, disabling split-KV (flash-decoding) as an empirical workaround for corrupted outputs when short queries attend long KV ranges (qo_len << kv_len, i.e. decode and prefix-cache extend).This PR returns
Falseunconditionally, re-enabling split-KV for NVFP4.Motivation
This is critical for MTP (Multi-Token Prediction) verification performance on SM120 (Blackwell RTX 5090). In the MTP scenario:
query_len=3(speculative tokens) must attend to long KV sequencesWhy it's safe now
The original corruption has been resolved by upstream fixes to the NVFP4 paged KV layout:
(B, 2*N_kv, N, F)with correct strides (replacing the broken 5-D(B,2,N,H,F)layout that had stride inheritance issues)o_dtypecorrected to BF16 for FA2 TC path (was incorrectly FP8)Context
This is part of the SM120 NVFP4 KV cache effort tracked in:
Test Plan