fix(attention): correct SM120 NVFP4 qk_correction layout, row-sum reduction, and lse - #3838
Merged
bkryu merged 3 commits intoJul 10, 2026
Conversation
…reduction Two defects in the SM120 NVFP4 attention path (added in flashinfer-ai#3640) versus the SageAttention3 kernel it ports: 1. qk_correction addressing. The kernel's TMA descriptor is built from tile_to_shape(SmemLayoutAtomDS{}, ...) with a stride-0 atom row (mainloop.cuh), so it addresses the correction tensor as compact [batch, heads, seq_len/128, seq_len] and never uses the passed strides. The Python side materialized an expanded [batch, heads, seq_len, seq_len] tensor via repeat_interleave (absent in SageAttention3's preprocess_qkv), so the kernel read block-0 rows for every Q block and crossed head/batch boundaries. Pass the compact tensor instead; this also removes the O(seq_len^2) fp32 materialization that dominated end-to-end time (7.3 ms of 9.5 ms at seq_len 16K). 2. Row-sum reduction width. One accumulator row spans 4 threads in the m16n8 acc layout; RowReductionThr was 8 (SageAttention3 uses 4), so finalize() folded the neighboring row's sum into every row_sum and halved the output (uniform-score V=ones probe returns 0.516 instead of 1.0; iid least-squares alpha vs exact reference is 0.49). On RTX PRO 6000 (SM120, CUDA 13): uniform-score V=ones probe 0.516 -> 1.031 (V-quantization exact value); iid cos 0.974 -> 0.982 with alpha 0.489 -> 0.980; block-structured Q cos 0.408 -> 0.984. End-to-end (quantize + fwd) 2.590 -> 0.621 ms at B4 H8 S4096 D128 and 9.528 -> 2.431 ms at B1 H8 S16384 D128, versus 1.107 / 5.594 ms for BF16 SDPA flash on the same shapes. Adds regression tests for both defects (structured-Q correction accuracy and output magnitude). AI-assisted.
waynehacking8
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 4, 2026 15:16
Contributor
📝 WalkthroughWalkthroughThis PR updates SM120 attention to use a compact qk_correction layout, adjusts softmax row reduction threading, writes LSE values in the kernel epilogue, and adds regression tests for layout, magnitude, and LSE behavior. ChangesSM120 attention layout, softmax, and LSE changes
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant attention_kernel_ws
participant LSEWriter
participant epilogue_params
participant mLSE
attention_kernel_ws->>LSEWriter: write_lse(shape_O, shape_LSE, stride_LSE, softmax_fused, tiled_mma_pv, row_offset, bidh, bidb)
LSEWriter->>LSEWriter: compute row and LSE value
LSEWriter->>epilogue_params: use ptr_LSE and stride_LSE
LSEWriter->>mLSE: store LSE at (row, bidh, bidb)
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a compact layout for the qk_correction tensor (one row per 128-token Q block) to optimize memory and alignment with the kernel's TMA layout, updating the CUDA bindings, Python preprocessing, validation, and documentation. It also fixes a bug in the softmax row reduction threshold (reducing RowReductionThr from 8 to 4) to prevent incorrect output scaling, and adds regression tests for both changes. The review feedback recommends casting the inputs of the qk_correction matrix multiplication to float32 before the operation to avoid potential numerical instability or overflow/underflow.
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.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/attention/test_nvfp4_attention_sm120.py (1)
183-245: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider adding a negative-path test for the shape validation itself.
These tests validate the correct compact shape end-to-end, but there's no direct test asserting
nvfp4_attention_sm120_fwd/_check_inputsrejects an incorrectly-shapedqk_correction(e.g., the old expanded[B,H,S,S]layout). Since a shape-layout mismatch was the actual root cause fixed here, a dedicated negative test would guard against regressing the validation logic itself.🤖 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 `@tests/attention/test_nvfp4_attention_sm120.py` around lines 183 - 245, Add a negative-path test around nvfp4_attention_sm120_fwd/_check_inputs that explicitly passes an incorrectly shaped qk_correction (for example the old expanded [B, H, S, S] layout) and asserts it is rejected. Reuse the existing structured-Q setup from test_nvfp4_attention_sm120_structured_q_correction to build valid q/k/v and qk_correction, then mutate only qk_correction to the wrong shape and verify the validation fails with the expected shape-check behavior.
🤖 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 `@csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cu`:
- Around line 268-274: `params.seqlen_s` is still being set for the old expanded
correction layout, but `qk_correction` is now compacted to one row per 128-token
Q block. Update the setup in `nvfp4_attention_sm120_binding` so
`params.seqlen_s` matches the compact row count used by the `delta_s_ptr`
metadata (the `seq_len / 128` dimension when `per_block_mean` is enabled,
otherwise 1). Keep the existing `TVM_FFI_ICHECK_EQ` shape checks for
`qk_correction` aligned with this metadata.
---
Nitpick comments:
In `@tests/attention/test_nvfp4_attention_sm120.py`:
- Around line 183-245: Add a negative-path test around
nvfp4_attention_sm120_fwd/_check_inputs that explicitly passes an incorrectly
shaped qk_correction (for example the old expanded [B, H, S, S] layout) and
asserts it is rejected. Reuse the existing structured-Q setup from
test_nvfp4_attention_sm120_structured_q_correction to build valid q/k/v and
qk_correction, then mutate only qk_correction to the wrong shape and verify the
validation fails with the expected shape-check behavior.
🪄 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: 509a7a86-e216-4552-861c-4eff679f7c6b
📒 Files selected for processing (4)
csrc/nvfp4_attention_sm120/nvfp4_attention_sm120_binding.cuflashinfer/nvfp4_attention_sm120.pyinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/consumer/softmax.cuhtests/attention/test_nvfp4_attention_sm120.py
…ection matmul LSEWriter had no call site, so fwd returned an uninitialized lse buffer (the existing tests only assert non-NaN, which held or failed depending on allocator reuse). Call it from the consumer loop after the softmax finalize, using the per-warp-group PV mma for the row mapping plus an explicit row offset, and remove the fp8_scalexfp4_scale_log2 factor that row_sum carries for the FP4 P quantization so lse is the plain ln-sum-exp of the scaled scores. The dead run() forwarding wrapper and its stale m16n8 fragment asserts are updated/removed. Also compute the qk_correction matmul in fp32 (review feedback): a float16 matmul output would overflow at 65504 even though the accumulation runs in fp32. Validated on RTX PRO 6000: uniform scores give lse = ln(seq_len) exactly; iid/structured lse mean abs error vs exact logsumexp is ~0.01-0.02 (max 0.30 on causal edge rows). New test covers causal and non-causal lse. AI-assisted.
waynehacking8
changed the title
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh (1)
94-134: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the dead LSE helpers
write_lse_infinity,compute_lse_log2,log2_to_ln, andln_to_log2are only defined here;LSEWriter::write_lseis the only caller, and its inline math already diverges fromcompute_lse_log2because of theSoftmaxFused::fp8_scalexfp4_scale_log2 * ln_2term.🤖 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/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh` around lines 94 - 134, The LSE helper functions here are dead code and should be removed: write_lse_infinity, compute_lse_log2, log2_to_ln, and ln_to_log2 are only defined in this epilogue and are not used by LSEWriter::write_lse. Clean up the unused helpers and any related includes/constants so the remaining write_lse path keeps its existing inline math and there is no stale duplicate logic.
🤖 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.
Nitpick comments:
In
`@include/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuh`:
- Around line 94-134: The LSE helper functions here are dead code and should be
removed: write_lse_infinity, compute_lse_log2, log2_to_ln, and ln_to_log2 are
only defined in this epilogue and are not used by LSEWriter::write_lse. Clean up
the unused helpers and any related includes/constants so the remaining write_lse
path keeps its existing inline math and there is no stale duplicate logic.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 76c78280-e4e3-4632-8a57-9f79bcbeab2e
📒 Files selected for processing (4)
flashinfer/nvfp4_attention_sm120.pyinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/compute/epilogue/lse_writer.cuhinclude/flashinfer/attention/sm120/nvfp4_attention_sm120/kernel/attention_kernel.htests/attention/test_nvfp4_attention_sm120.py
🚧 Files skipped from review as they are similar to previous changes (1)
- flashinfer/nvfp4_attention_sm120.py
…LSE helpers Review feedback: add a negative-path test that the old expanded [B, H, S, S] qk_correction is rejected by the shape validation, and remove the now-unused write_lse_infinity / compute_lse_log2 / log2_to_ln / ln_to_log2 helpers (store_zero fills lse itself and write_lse computes its value inline). AI-assisted.
Collaborator
|
/bot run tests/attention |
Collaborator
Collaborator
|
[FAILED] Pipeline #57005597: 7/20 passed |
bkryu
approved these changes
Jul 9, 2026
5 tasks
saltyminty
pushed a commit
that referenced
this pull request
Jul 13, 2026
…3897) <!-- .github/pull_request_template.md --> ## 📌 Description Enables the SM120 NVFP4 attention kernel (#3640) on SM121 (GB10 / DGX Spark). The kernel code needs no changes — it compiles and runs correctly for `sm_121a`; this PR only widens the SM120-only gates - `pytest tests/attention/test_nvfp4_attention_sm120.py`: **7/7 pass** on this branch; **13/13 pass** with #3838's fixes and regression tests applied on top. - Kernel benchmarks vs FA2 BF16 ragged prefill at identical shapes (non-causal, attention-only kernel time): | Config (H8) | FA2 BF16 | NVFP4 | Speedup | |---|---|---|---| | B4 S4096 D128 | 2.97 ms (93 TF/s) | 1.28 ms (215 TF/s) | 2.32× | | B2 S8192 D128 | 5.81 ms (95 TF/s) | 2.48 ms (221 TF/s) | 2.34× | | B1 S32768 D128 | 46.09 ms (95 TF/s) | 19.51 ms (225 TF/s) | 2.36× | | B1 S16384 **D64** | 5.89 ms (93 TF/s) | 3.81 ms (144 TF/s) | 1.55× | Speedup is stable across head counts (2.2–2.6× for H=1…32 at D128); throughput saturates from B·H ≥ 16 at S4096. Speedup is around 1.5-1.6x for head dim 64. <details> <summary>Commands to reproduce the perf numbers</summary> NVFP4 (shape lists zip together; one row printed per config): ```bash # D=128 rows + head-dim sweep python benchmarks/bench_nvfp4_attention_sm120.py \ --batch-size 4 2 2 1 1 --num-heads 8 --head-dim 128 \ --seq-len 4096 4096 8192 16384 32768 \ --no-causal --warmup 3 --repeat 10 # D=64 rows python benchmarks/bench_nvfp4_attention_sm120.py \ --batch-size 4 2 2 1 1 --num-heads 8 --head-dim 64 \ --seq-len 4096 4096 8192 16384 32768 \ --no-causal --warmup 3 --repeat 10 # head-count sweep (B*H occupancy) python benchmarks/bench_nvfp4_attention_sm120.py \ --batch-size 1 --num-heads 1 2 4 8 16 32 --head-dim 128 \ --seq-len 4096 --no-causal --warmup 3 --repeat 10 ``` The "NVFP4" column is the `attention_only` number (CUDA-graph replay, pure kernel time, quantization excluded); `end_to_end` additionally includes `quantize_qkv` each iteration. FA2 BF16 baseline (identical shapes; CUPTI kernel timing is on by default; uniform full-length sequences — do not pass `--random_actual_seq_len`): ```bash for cfg in "4 4096" "2 4096" "2 8192" "1 16384" "1 32768"; do set -- $cfg python benchmarks/flashinfer_benchmark.py \ --routine BatchPrefillWithRaggedKVCacheWrapper --backends fa2 \ --batch_size $1 --s_qo $2 --s_kv $2 \ --num_qo_heads 8 --num_kv_heads 8 \ --head_dim_qk 128 --head_dim_vo 128 \ --q_dtype bfloat16 --kv_dtype bfloat16 --refcheck done # D=64 baseline: same loop with --head_dim_qk 64 --head_dim_vo 64 # head sweep baseline: --batch_size 1 --s_qo 4096 --s_kv 4096, loop --num_qo_heads/--num_kv_heads over 1 2 4 8 16 32 ``` Environment: NVIDIA GB10 (SM121), CUDA 13.0, torch 2.11. FA2 median of 30 iters (CUPTI); NVFP4 median of 10 iters (CUDA-graph). Comparing FA2 kernel time against NVFP4 `attention_only` is apples-to-apples — both exclude host launch overhead. </details> <!-- What does this PR do? Briefly describe the changes and why they’re needed. --> ## 🔍 Related Issues <!-- Link any related issues here --> - #3809 - #3838 - PR 3838 is an orthogonal correctness fix ## 🚀 Pull Request Checklist Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete. ### ✅ Pre-commit Checks - [x] I have installed `pre-commit` by running `pip install pre-commit` (or used your preferred method). - [x] I have installed the hooks with `pre-commit install`. - [x] I have run the hooks manually with `pre-commit run --all-files` and fixed any reported issues. > If you are unsure about how to set up `pre-commit`, see [the pre-commit documentation](https://pre-commit.com/). ## 🧪 Tests - [x] Tests have been added or updated as needed. - [x] All tests are passing (`unittest`, etc.). ## Reviewer Notes <!-- Optional: anything you'd like reviewers to focus on, concerns, etc. --> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Expanded NVFP4 attention support to include GPU compute capability 12.1 in addition to 12.0. * Broadened which NVFP4 SM120 modules are generated and enabled, and improved JIT build flag selection across compatible CUDA environments. * **Bug Fixes** * Updated compute-capability validation and the related error/skip messaging to reflect 12.0 and 12.1 support. * Adjusted test gating to run when either supported compute capability is available, reducing unnecessary skips. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
5 tasks
5 tasks
saltyminty
pushed a commit
that referenced
this pull request
Aug 25, 2026
<!-- .github/pull_request_template.md --> ## 📌 Description This PR optimizes the CUTLASS/CUDA SM120 NVFP4 attention forward kernel, with a focus on head-dimension-128 workloads used by Cosmos inference. The largest performance contribution is **symmetric N64/N64 score-slot reuse** inside the existing N128 attention tile: 1. Compute the complete N128 QK score tile and establish the row maximum and online-softmax rescale across all 128 columns. 2. Softmax-quantize the first N64 score slot and consume it in the PV MMA. 3. Reuse that retired score-register slot immediately for the next tile's QK scores while the second N64 slot is consumed. 4. Repeat symmetrically for the second slot. This avoids keeping a second full score fragment live and exposes more QK/PV overlap while preserving the full-N128 softmax reduction. Additional kernel changes: - Split the main loop into a prologue, branch-free steady state, and final drain. - Remove the obsolete math-order barrier and reduce warp synchronization. - Compile masking out of the noncausal specialization. For causal traversal, only the first tile can intersect the diagonal; subsequent tiles are fully valid. - Derive mask row/column coordinates directly instead of retaining an identity tensor. - Tune consumer register allocation and persistent CTA scheduling/work distribution for both short and long causal/noncausal workloads. - Add a compile-time LSE specialization so the output-only path does not allocate or write LSE. ### LSE API behavior `return_lse=False` is now the default and returns only the attention output tensor: ~~~python out = flashinfer.nvfp4_attention_sm120_fwd(..., return_lse=False) ~~~ `return_lse=True` preserves the LSE-capable path and returns `(out, lse)`: ~~~python out, lse = flashinfer.nvfp4_attention_sm120_fwd(..., return_lse=True) ~~~ Disabling LSE only removes the LSE allocation/writeback; it does not remove any computation required for the attention output. The tests compare the output-only and LSE-enabled specializations, and the existing output/LSE reference checks continue to pass. ### Performance Local measurements used one NVIDIA RTX PRO 6000 Blackwell Server Edition GPU (SM120, 188 SMs), BF16 input/output, `D=128`, `per_block_mean=True`, and `return_lse=False`. The table reports CUDA-Graph attention-only kernel timing (median of 100 iterations after 10 warmups); QKV quantization is excluded. The maximum supported 2430 MHz graphics clock was requested and monitored during the runs. | Shape | Mode | [#3640](#3640) baseline (reported) | CuTe DSL reference (reported) | This PR (measured) | Latency vs #3640 | Latency vs CuTe DSL | |---|---:|---:|---:|---:|---:|---:| | B4, H8, S4096, D128 | Noncausal | 0.299 ms / 920.5 TFLOP/s | 0.262 ms / 1050.2 TFLOP/s | **0.262 ms / 1050.1 TFLOP/s** | **-12.4%** | same at reported precision | | B1, H8, S32768, D128 | Noncausal | 4.970 ms / 884.9 TFLOP/s | 4.398 ms / 1000.0 TFLOP/s | **4.079 ms / 1078.2 TFLOP/s** | **-17.9%** | **-7.3%** | | B4, H8, S4096, D128 | Causal | 0.223 ms / 616.6 TFLOP/s | 0.180 ms / 764.1 TFLOP/s | **0.165 ms / 830.6 TFLOP/s** | **-26.0%** | **-8.3%** | | B1, H8, S32768, D128 | Causal | 2.958 ms / 743.3 TFLOP/s | 2.516 ms / 874.0 TFLOP/s | **2.207 ms / 996.5 TFLOP/s** | **-25.4%** | **-12.3%** | The #3640 and CuTe DSL columns reproduce the values shared in the RTX Blackwell kernel performance design document; they were not remeasured in the same run as this PR. The relative percentages therefore use the published rounded latencies and should be treated as cross-run comparisons. ## Contributors - [Atharva Joshi (@atharvajoshi10)](https://github.com/atharvajoshi10) — contributed the CuTe DSL kernel optimization and design document, including the symmetric N64/N64 score-slot reuse strategy, together with the RTX PRO 6000 benchmark results that motivated this CUTLASS implementation. ## 🔍 Related Issues - Follow-up to #3640. - Builds on the SM120 NVFP4 attention/LSE implementation in #3838. ## 🚀 Pull Request Checklist ### ✅ Pre-commit Checks - [x] I installed `pre-commit` in an isolated environment. - [x] I ran `pre-commit run --all-files`; all hooks passed. ## 🧪 Tests - [x] Tests were added for both `return_lse=False` and `return_lse=True`. - [x] `python -m pytest -q tests/attention/test_nvfp4_attention_sm120.py` — **16 passed**. - [x] Output-only and LSE-enabled attention outputs match with `rtol=0, atol=5e-4`; LSE-enabled reference checks pass. ## Reviewer Notes Please pay particular attention to: - The intentional API default change: callers that need the previous `(out, lse)` result should pass `return_lse=True`. - The full-N128 max/rescale followed by per-N64 softmax quantization and score-slot reuse. - The persistent scheduler changes for triangular causal work distribution. The upstream `pre-commit` check is green. The full PR test matrix is currently skipped by the repository's unauthorized-PR gate and requires maintainer approval before it can run. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added grouped-query and multi-query attention support with separate query and key/value head counts. * Added optional LSE output and support for unpadded key/value sequence lengths. * Added configurable output buffers and output data types. * Enhanced tracing and benchmarking for the new attention options. * **Bug Fixes** * Improved validation for head ratios, tensor shapes, sequence lengths, and masking. * Improved scheduling and handling of empty or irregular sequence shapes. * **Tests** * Added coverage for GQA/MQA, LSE behavior, unpadded sequences, output dtypes, tracing, and invalid inputs. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Atharva Joshi <atjoshi@nvidia.com>
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
3 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.
Description
While benchmarking the SM120 NVFP4 attention path from #3640 on an RTX PRO 6000 for #3809, I found three defects versus the SageAttention3 kernel it ports. All are invisible to the current tests (iid inputs make the correction negligible, cosine/mean-abs-err thresholds absorb a uniform scale error, and lse is only checked for NaN), but they matter a lot on real inputs.
1.
qk_correctionis addressed as a compact tensor, but Python passes an expanded one.The mainloop builds the DS TMA descriptor from
tile_to_shape(SmemLayoutAtomDS{}, ...)where the atom isLayout<Shape<kBlockM, kBlockN>, Stride<_0, _1>>(kernel/traits.h:165,compute/mainloop.cuh:193-200). The resulting gmem layout has stride 0 within each 128-row block, i.e. the kernel addressesptr_dsas a compact[batch, heads, seq_len/128, seq_len]tensor and never uses the strides passed from the binding. SageAttention3'spreprocess_qkvproduces exactly that compactdelta_s(oneqm @ k^Trow per 128-token block). Our port added arepeat_interleave(128, dim=2)that materializes[batch, heads, seq_len, seq_len], so at runtime the kernel read block 0's correction rows for every Q block and crossed head/batch boundaries forbidh > 0/bidb > 0.I verified the addressing empirically before changing anything: with B=1 H=1, scrambling every row
r >= seq_len/128of the expanded tensor leaves the output bitwise identical (those rows are never read), and a marker placed in gmem rowrmoves exactly the output rows of Q blockr, forr = 0..seq_len/128-1.The fix passes the compact tensor through (
quantize_qkvoutput shape changes to[B, H, S/128, S], or[B, H, 1, S]forper_block_mean=False) and updates the binding shape check. This also removes the O(seq_len^2) fp32 materialization, which was 99% of the preprocessing cost (7.3 ms of the 9.5 ms end-to-end at S=16K, and an 8.6 GB allocation). Per review feedback the correction matmul now also runs in fp32, since a float16 matmul output would overflow at 65504.2. Row-sum reduction folds in the neighboring row.
One accumulator row spans 4 threads in the acc layout (2 columns per thread per 8-column group).
SoftmaxFused::RowReductionThrwas 8 -- SageAttention3 has 4 (softmax_fused.h:35upstream) -- so the__shfl_xor(4)step infinalize()added the neighboring row's sum into everyrow_sum, halving the output. A probe with uniform scores and V = all-ones (exact output must be 1.0 everywhere, independent of P quantization) returns 0.5156 = half the V-dequant value, exactly.3.
lseis never written.LSEWriter(compute/epilogue/lse_writer.cuh) had no call site, sofwdreturned an uninitialized buffer -- the existing non-NaN assertions pass or fail depending on allocator reuse (SageAttention3 has its LSE store commented out entirely; the writer here was half-adapted dead code with m16n8 fragment asserts that don't compile against this kernel's mma). It's now called from the consumer loop after the softmax finalize, using the per-warp-group PV mma for the row mapping plus an explicit row offset, and the value drops thefp8_scalexfp4_scale_log2factor thatrow_sumcarries for the FP4 P quantization, solseis the plain ln-sum-exp of the scaled scores. Uniform scores now givelse = ln(seq_len)exactly; iid/structured lse mean abs error vs exactlogsumexpis 0.01-0.02.Accuracy A/B on RTX PRO 6000 (CUDA 13, torch 2.11, this repo at c53229e)
Reference is exact fp32 SDPA on the same inputs;
alphais the least-squares scale<out,ref>/<ref,ref>(1.0 = unbiased), so it exposes the halving that cosine cannot see.The structured-Q rows are the case the SageAttention smoothing machinery exists for; on main the correction is applied to the wrong blocks, so accuracy collapses exactly where the algorithm is supposed to help.
End-to-end perf (quantize_qkv + fwd, in-repo benchmark, non-causal)
Kernel-only time is unchanged (~0.42 / ~1.9 ms, 660 / 570 TFLOPs/s); the fixed e2e includes the fp32 correction matmul and the lse store. The 16K shape shows some run-to-run spread on my box (2.5-3.1 ms) -- 3-4x faster than main either way. End-to-end goes from 2.3x slower than BF16 flash attention to roughly 1.5-2.2x faster, which is the regime diffusion workloads (fresh Q/K/V every step) actually run in.
Note:
quantize_qkv's output shape forqk_correctionchanges; the pair of public APIs stays self-consistent, so code using them together is unaffected.Related Issues
#3809 (this makes the existing SM120 NVFP4 path usable and fast end-to-end on diffusion shapes; kernel-level follow-ups tracked there), #3640.
Pull Request Checklist
Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.Tests
unittest, etc.).pytest -q tests/attention/test_nvfp4_attention_sm120.py: 13 passed (7 existing + 6 new) on RTX PRO 6000. The new regression tests fail on main:test_nvfp4_attention_sm120_structured_q_correction(cos 0.41/0.45 vs the 0.95 floor),test_nvfp4_attention_sm120_output_magnitude(0.516 vs 1.0 +/- 0.05), andtest_nvfp4_attention_sm120_lse(uninitialized buffer).Reviewer Notes
The decisive probes are easy to re-run: (a) scramble rows
>= seq_len/128of the expanded correction on main and observe bitwise-identical output; (b) uniform scores with V = all-ones must return 1.0 and lse = ln(seq_len). Happy to split the three fixes into separate PRs if you prefer.Summary by CodeRabbit
New Features
Bug Fixes