Repository navigation
perf(cake_sparse_mla): round-3 DSv4 sparse-MLA programs for SM100/SM103 (persistent-item epilogue, index ring, retained-KV decode body, SwapsAb V alias, routing) - #5630
Conversation
…to the four-owner two-stage program Width 256 < topk <= 388 now runs the split4 owner kernel with a fully masked fourth tile instead of the three-owner pair: same-session paired CUPTI on GB300 decode width 260 14.78/14.75 us -> 12.67/12.70 us (1.23x -> 1.44x vs trtllm-gen), B200 15.74 -> 13.60/13.57 us (1.11x -> 1.29x); width 388 rows unchanged. Mirrors the Cake seed rule; the dispatch test pins the new routing. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…he portfolio producer Dense Q with equal q_lens is the same row memory as ragged Q; the portfolio producer plus two-warp reducer beats the fixed-Q cluster-2 guard program on the sparse fixed-Q decode rows (GB300 1.03x -> 1.21x vs trtllm-gen). SWA-only, 24+ token and prefill fixed-Q rows keep the guard program. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…rules With the sparse fixed-Q decode rows on the portfolio producer no exported shape reaches the fixed-Q programs, so the dense (not ragged) BF16/H64 rows take the ragged rules for every case (SWA-only, 24+ tokens, prefill) and the fixed-Q programs leave the registry. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…grams (round 3) Round-3 regeneration of the Cake DSv4 sparse-MLA programs for sm_100a and sm_103a. Seven programs change against the round-2 set, the other 18 are byte-identical to their tracked files: - fp8_h128_prefill_source_persistent (+ _uniform): persistent-body item switch with next-item Q/index prefetch under the current item's PV tail; the uniform (sub-128-token decode) variant now builds without the ptxas register-usage pin, which read 1-2 % slower on the uniform decode rows. - bf16_h128_prefill_v42 (+ _snake): index ring for the multi-tile prefill items. - bf16_h32_topk128x_early_v47: v48 body. - bf16_h8_h16_source_exact / fp8_h8_h16_source_exact: V-operand alias in the SwapsAb decode body, KV ring 4 -> 5 stages. - sm_100a additionally gains the split_reduce program of the split4 rule. Dropped (no shape resolves to them any more): bf16_h128_topk128x_split3_sm100, bf16_h64_fixed_q, bf16_h64_fixed_q_reduce; the 25 stale round-2 program files per arch are removed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…riant carries the ptxas register-usage flag The uniform (sub-128-token decode) variant of the FP8/H128 persistent body ships without -Xptxas=--register-usage-level=10: its softmax chain does not pace the item and the pinned build reads 1-2 % slower on the uniform decode rows. The public-flags test now asserts the pin on the prefill variant and its absence on the uniform variant instead of requiring it on both. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
|
/bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py |
|
@flashinfer-bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py |
|
Mirroring Failed Failed to mirror PR to GitLab. Check logs for details. |
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (53)
💤 Files with no reviewable changes (16)
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 6 remain after this review. 📝 WalkthroughWalkthroughThe change updates generated CAKE DSv4 CUDA kernels and JIT registrations for SM 100a and SM 103a. It also changes MLA routing rules, removes obsolete split-reduction paths, and updates routing and compile-flag tests. ChangesCAKE DSv4 kernel updates
JIT and dispatch updates
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~60 minutes Suggested reviewers: Merge Risk: ⚪ Minimal · up to This change regenerates the CAKE DSv4 decode kernels for SM100a and SM103a, reroutes several BF16 shapes, and removes programs that exported shapes no longer use. No concrete defect was found in the kernel layouts, bindings, registrations, or routing rules. The change appears ready to merge once the author-reported GPU test runs are confirmed. Security Architecture ReviewSecurity architecture risk: 🔵 Low · up to The changes are confined to GPU decode routing and kernels. No introduced security vulnerability was established, but the new index-buffer synchronization and changed workload routes warrant review. Retained concerns Security review detailsSecurity Blast Radius
Trust Boundaries and Controls
Resilience and Maintainability Implications
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 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 |
|
/bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py |
|
/bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py |
|
[SUCCESS] Pipeline #70187738: 25/25 executed test jobs passed |
…0/SM103 (stacked on #5630) (#5644) ## Summary Round-4 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5630 (this PR is stacked on #5630; host and tests unchanged): - **BF16/FP8 H8/H16 decode programs (16 rows):** the FP8 body issues its K gathers back-to-back after hoisting the page-offset loads and runs an 8-slot K ring; both bodies get a deeper K/V alias ring so a 3-tile item's second tile no longer waits for the first PV. GB300 FP8 rows 1.40x -> **1.70x** (geomean), B200 1.52x -> **1.73x**; 3-tile BF16 rows -0.2..-0.5 us (GB300 decode-000062 1.50x -> 1.64x). - **BF16/H128 persistent prefill programs (128+ query-token rows and the K-reuse rows):** FA4-style lazy softmax rescale (a row keeps its running max until the true max grows by more than 2^8; the deferred factor is applied then). prefill-style-000088: GB300 1.34x -> **1.47x**, B200 1.22x -> **1.37x**; 000092 1.35x -> **1.41x** / 1.22x -> **1.33x**. - **BF16 one-tile SWA decode programs (18 rows):** the sparse-index load is issued before the Q TMA (statement order only; bit-identical output): -0.1..-0.26 us per row. - Every other program is a byte-identical regeneration (renamed for the new producer revision). ## Validation Export r13 (regenerated programs as published here; host and tests at #5630's `76d586ff6`), same-session paired CUPTI (active-union median, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 | 134/134 | | **94 canonical rows faster than the default path** | 94/94, geomean **1.343x**, min 1.130x (#5630: 1.300x, min 1.098x) | 94/94, geomean **1.319x**, min 1.048x (#5630: 1.280x, min 1.027x) | | 40 hardening rows faster than the default path | 40/40, geomean **1.353x**, min 1.085x (#5630: 1.330x, min 1.071x) | 40/40, geomean **1.310x**, min 1.069x (#5630: 1.307x, min 1.065x) | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, fresh clone of this commit) | 296 passed | 296 passed | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16), hardened by #5591 (8589d49), round 2 #5610 (4b8167e), round 3 #5630 (76d586f), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance Improvements** * Updated attention processing to handle more key/value data stages per pipeline, potentially improving throughput. * Sparse-index lookups now occur earlier in the loading process, helping streamline attention data preparation. * **Numerical Improvements** * Refined softmax accumulation to avoid unnecessary correction for small score changes while preserving correction for larger changes. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
…0/SM103 (exact numerics; stacked on #5644) (#5686) ## Summary Round-5 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5644, plus one host-side binding line. Every regenerated program computes bit-identical output to its #5644 predecessor: this round only re-times data movement and epilogues, with no change to the arithmetic. - **BF16 H8/H16 decode programs (16 rows):** eight warps issue the K gathers (both arches); the FP8 H8/H16 programs do the same on SM100a. - **BF16/H128 persistent prefill programs (SM103a):** the O epilogue loads the next tile's TMEM accumulator ahead of the store (−1.5..−4.3 %). - **BF16/H64 split-K decode programs (9 rows):** owner-CTA form with hoisted partial select, zero-fill and 256-bit stores: −1.7..−2.9 us per row (GB300 class geomean 1.25x -> **1.51x**, B200 1.18x -> **1.45x**). - **FP8/H128 persistent prefill programs:** TMA-store O epilogue (−0.5..−0.6 us on the 9-tile rows, −7.3 % on prefill-style-000093). - **BF16/H32 retained-KV decode programs (4 rows):** the softmax publishes P before the row-sum exchange (separate stats barrier), −0.3..−1.2 %. - Every other program is a byte-identical regeneration (renamed for the new producer revision). - **Host (`flashinfer/mla/cake_dsv4.py`, first commit of this branch):** the FP8 persistent bodies now store O through a TMA descriptor (`tmap_o`); `_TMA_SOURCE_ALIASES` binds it to the same `[tokens, heads, 512]` output rows the plain `O` argument sees. Without it every FP8 launch on those routes raised "argument 'tmap_o' (tma_buffer) has no host value" (28 failures in the hardening tests on the regenerated programs alone). `test_registered_arg_plans_use_known_names` now also requires every TMA descriptor to alias a tensor the host actually places in its value table, so a vocabulary-only match cannot pass again. ## Validation Exports of the regenerated programs as published here (B200: export r15 against #5644's `6d479bea1`; GB300: export r16 against this branch's host-fix commit `fd2131236`), same-session paired CUPTI (active-union median, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 correct; 133/134 export gates (row 085: baseline-arm order sensitivity, see doc § Round-over-round) | 134/134 | | **94 canonical rows faster than the default path** | 94/94 > 1, geomean **1.370x**, min 1.129x (#5644: 1.343x, min 1.130x) | 94/94, geomean **1.356x**, min 1.103x (#5644: 1.319x, min 1.048x) | | 40 hardening rows faster than the default path | 40/40 > 1, geomean **1.385x**, min 1.096x (#5644: 1.353x, min 1.085x) | 40/40 > 1, geomean **1.333x**, min 1.045x (#5644: 1.310x, min 1.069x) | | output bit-identical to the #5644 programs on every changed row | yes | yes | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, fresh clone of this commit) | 296 passed, 0 failed | 296 passed, 0 failed | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16), hardened by #5591 (8589d49), round 2 #5610 (4b8167e), round 3 #5630 (ebe0efb), round 4 #5644 (6d479be), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added output tensor-map support for selected CAKE DSV4 attention paths, enabling outputs to be staged and transferred through optimized tensor operations. * Updated selected kernels to process larger output groups and adjust work distribution for supported GPU architectures. * Refreshed kernel-family registration and dispatch across supported architectures. * **Bug Fixes** * Improved handling of invalid key/value indices and synchronization of output statistics and transfers. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
…0/SM103 (exact numerics; stacked on #5686) (#5700) ## Summary Round-6 performance of the `backend="cake"` DeepSeek-V4 sparse-MLA decode path of `trtllm_batch_decode_sparse_mla_dsv4` (SM100a / SM103a), regenerated from the same producer pipeline as #5686. No host change. Every regenerated program computes bit-identical output to its #5686 predecessor: this round only re-times data movement, store widths, pipeline depths and the kernel ABI, with no change to the arithmetic. - **BF16/H128 split-K programs, classes E / F / B / C (18 rows):** reducer family — 256-bit partial stores, batched partial loads, and a leaner 5-split reducer on SM103a; class E additionally gets a per-arch owner (SM100a: 4-stage K ring; SM103a: new 8-slot K/PV ring body). E rows −2.7..−3.7 us (GB300) / −2.3..−2.5 us (B200); F rows −0.4..−0.5 us; B S=3 rows −0.65..−0.70 us. - **FP8/H64 single-CTA SWA programs, class K1 (11 rows):** 256-bit O stores in the lane-per-row epilogue, −0.06..−0.27 us. - **BF16 one-tile SWA programs, classes A / D / G (17 rows) + the H32 retained-KV and H8/H16 guards decode programs (SM103a):** tensor maps passed by value (`__grid_constant__`) instead of through a pointer, −0.26..−1.03 us per row (−4.6..−9.0 % on A/D/G). - **FP8/H64 wide one-partition programs, class K4 (2 rows):** third 32 KiB K stage, −0.06..−0.16 us. - **FP8/H128 persistent programs, classes K2 / L (33 rows, SM103a):** epilogue TMEM load-ahead of 4, −0.6..−3.3 % on the H128 rows. - **BF16/H128 striped prefill programs, class M (7 rows):** V rows issued ahead (LDS order only), −1.5..−4.7 % (B200), −0.3..−1.6 % (GB300). - **BF16 H8/H16 decode programs, class H (8 rows, SM100a):** the never-read padded gather rows of the last item are skipped, −0.45..−1.02 us on the padded rows. - Every other program is a byte-identical regeneration (renamed for the new producer revision). ## Validation Exports of the regenerated programs as published here (B200: export r16; GB300: export r17, both from Cake tree `939659882ba` against a fresh clone of #5686's `10c80d75b`), same-session paired CUPTI (active-union median, cold L2) against the default trtllm-gen path at the same revision: | | GB300 (sm_103a) | B200 (sm_100a) | |---|---|---| | 94 canonical + 40 hardening shapes correct (BF16 atol=rtol=1e-2, FP8 0.1; padded / separate-table / offset axes) | 134/134 correct; 134/134 export gates | 134/134 correct; 134/134 export gates (decode-000002 and hardening-000036 in a second measurement round, see doc § Gate table) | | **94 canonical rows faster than the default path** | 94/94 > 1, geomean **1.411x**, min 1.116x (#5686: 1.370x, min 1.129x) | 94/94 > 1, geomean **1.352x**, min 1.075x (#5686: 1.356x, min 1.103x) | | 40 hardening rows faster than the default path | 40/40 > 1, geomean **1.390x**, min 1.092x (#5686: 1.385x, min 1.096x) | 40/40 > 1, geomean **1.354x**, min 1.085x (#5686: 1.333x, min 1.045x) | | output bit-identical to the #5686 programs on every changed row | yes (92 changed rows) | yes (48 changed rows) | | compute-sanitizer synccheck + memcheck on the changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, fresh clone of this commit) | 296 passed, 0 failed | 296 passed, 0 failed | ## Baselines and their source PRs - **trtllm-gen default path** = `trtllm_batch_decode_sparse_mla_dsv4(...)` without `backend="cake"` (the TRTLLM-GEN DSv4 sparse-MLA kernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling). Origin PRs: #3269 (9c76c99, 2026-05-21); SM120 kernels #3395 (f954694, 2026-06-15); DSv4.1 unification #5197 (eb5f05b, 2026-09-18). - **Cake backend under test** = `backend="cake"` from #4573 (13db2cf, 2026-09-16), hardened by #5591 (8589d49), round 2 #5610 (4b8167e), round 3 #5630 (ebe0efb), round 4 #5644 (6d479be), round 5 #5686 (10c80d7), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py`. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added optimized attention-kernel variants for newer NVIDIA GPU architectures, including improved handling of split outputs and KV data. * Removed the need for callers to provide a separate tensor-map workspace for several kernel variants. * **Performance Improvements** * Optimized data loading, pipelining, and output conversion across multiple GPU variants. * Added a specialized path for combining three partial outputs. * **Bug Fixes** * Improved bounds handling for gathered KV rows, including out-of-range indices. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Developed by the CAKE team, the kernels in this PR outperform the relevant SOTA baselines at the time of merge on the vast majority of performance benchmark shapes, matching them otherwise.
Summary
Round-3 performance of the
backend="cake"DeepSeek-V4 sparse-MLA decode path oftrtllm_batch_decode_sparse_mla_dsv4(SM100a / SM103a), regenerated from the same producer pipeline as #5610:
epilogue drained the 64x256 fp32 accumulator with 16-byte stores (4096 half-sector transactions per CTA per item). It now uses 32-byte
stores. prefill-style-000089 / 000093: GB300 1.24x / 1.23x -> 1.33x / 1.31x, B200 1.17x / 1.18x -> 1.25x / 1.25x vs the default path.
--register-usage-level=10(kept on the lane-gather program):-5.8 % SASS, GB300 decode rows 007/010/031/034 1.24x -> 1.26x.
loaded batched. prefill-style-000088 / 000092: GB300 1.23x -> 1.34x / 1.35x, B200 1.13x -> 1.21x / 1.22x; hardening rows 21-37 all
faster.
gathers: GB300 1.08-1.10x -> 1.30-1.34x, B200 1.09-1.12x -> 1.32-1.33x.
1.36x -> 1.47x (geomean), BF16 1.23x -> 1.31x; B200 1.39x -> 1.52x, 1.26x -> 1.35x.
dense (not ragged) BF16/H64 rows follow the ragged rules (fixed-Q guard row: GB300 1.03x -> 1.21x, B200 1.03x -> 1.11x); the
three-owner topk128x program and the fixed-Q programs are no longer reachable and leave the registry.
Validation
Export r12 (this branch: host and tests at
84b04ff02, regenerated programs as published here), same-session paired CUPTI (active-union median, 500 warmup / 3000 callsper arm, cold L2) against the default trtllm-gen path at the same revision:
tests/mla/test_cake_dsv4*.py(GPU, fresh clone of this commit)tests/mla/test_cake_dsv4*.py(CPU, this host)Baselines and their source PRs
trtllm_batch_decode_sparse_mla_dsv4(...)withoutbackend="cake"(the TRTLLM-GEN DSv4 sparse-MLAkernels plus the framework launches it needs for equal semantics: padded-Q masking, separate-table merge, workspace/counter handling).
Origin PRs: Add DeepSeek V4 sparse MLA TRTLLM-GEN kernels #3269 (9c76c99, 2026-05-21); SM120 kernels feat(attention): add SM120 sparse MLA kernels #3395 (f954694, 2026-06-15); DSv4.1 unification refactor(sparse-mla): unify SM120 execution, calibration and DSv4.1 support #5197 (eb5f05b, 2026-09-18).
backend="cake"from feat(cake_sparse_mla): support SM100 and SM103 #4573 (13db2cf, 2026-09-16), hardened by fix(cake_sparse_mla): harden the DSv4 sparse-MLA Cake backend for #4671 (padded Q, separate SWA/compressed tables, caller-owned workspace, length offsets) and regenerate the SM100/SM103 programs #5591 (8589d49), round 2 perf(cake_sparse_mla): warp-uniform FP8/H128 TMA gathers and 32-byte epilogue stores in the DSv4 sparse-MLA Cake programs (SM100/SM103) #5610(4b8167e, 2026-09-27), regenerated again by this PR.
tests/mla/test_cake_dsv4.py/test_cake_dsv4_hardening.py.🤖 Generated with Claude Code
Summary by CodeRabbit
References
The shared reference corpus for CAKE kernel development includes the following projects, documentation, and existing CAKE work:
These are corpus-level references. PR-specific implementation details, changes, and benchmark references are documented above.