Skip to content

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

Merged
yyihuang merged 5 commits into
flashinfer-ai:mainfrom
yyihuang:cake-713-dsv4-round3
Sep 28, 2026

Conversation

@yyihuang

@yyihuang yyihuang commented Sep 28, 2026 •

Copy link
Copy Markdown
Collaborator

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 of trtllm_batch_decode_sparse_mla_dsv4
(SM100a / SM103a), regenerated from the same producer pipeline as #5610:

  • FP8/H128 persistent prefill program (128+ query-token rows): the steady KV tile already runs at the tensor-pipe floor; the per-item
    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.
  • Sub-128-token FP8/H128 uniform-gather program: compiled without --register-usage-level=10 (kept on the lane-gather program):
    -5.8 % SASS, GB300 decode rows 007/010/031/034 1.24x -> 1.26x.
  • BF16/H128 striped prefill program: the sparse-index block of the next work item is staged during the current item (two-slot ring) and
    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.
  • BF16/H32 retained-KV decode program (rows 73/74/76/77): hoisted scale loads, one epilogue release fence, eight warps issue the
    gathers: GB300 1.08-1.10x -> 1.30-1.34x, B200 1.09-1.12x -> 1.32-1.33x.
  • BF16/FP8 H8/H16 decode programs (16 rows): PV reads the V quarter from the staged K tile instead of re-gathering it: GB300 FP8 rows
    1.36x -> 1.47x (geomean), BF16 1.23x -> 1.31x; B200 1.39x -> 1.52x, 1.26x -> 1.35x.
  • Routing: width-260 BF16/H128 topk128x rows run the four-owner two-stage program (GB300 1.23x -> 1.44x, B200 1.11x -> 1.29x);
    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.
  • Every other program is a byte-identical regeneration (renamed for the new producer revision).

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 calls
per arm, 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.300x, min 1.098x (#5610: 1.243x, min 1.054x) 94/94, geomean 1.280x, min 1.027x (#5610: 1.232x, min 1.041x)
40 hardening rows faster than the default path 40/40, geomean 1.330x, min 1.071x (#5610: 1.30x, min 1.058x) 40/40, geomean 1.307x, min 1.065x (#5610: 1.25x, min 1.037x)
exported program vs source build (agreement gate, 134 rows) 134/134, geomean 1.0018 134/134, geomean 1.0026
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
tests/mla/test_cake_dsv4*.py (CPU, this host) 212 passed, 84 skipped 212 passed, 84 skipped

Baselines and their source PRs

🤖 Generated with Claude Code

Summary by CodeRabbit

  • New Features
    • Updated CAKE DSv4 routing across supported GPU variants, including revised handling for dense BF16/H64 queries and selected BF16/H128 input widths.
    • Expanded coverage for sparse-index staging and attention processing across value-head groups.
  • Performance
    • Optimized output packing and stores, and adjusted kernel staging and memory allocation for updated execution paths.
  • Compatibility
    • Removed obsolete kernel variants and registrations; workloads that previously used those routes may now follow different supported routes.

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.

averyhNV and others added 5 commits September 27, 2026 18:43
…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>
@yyihuang

Copy link
Copy Markdown
Collaborator Author

/bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py

@yyihuang

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

Mirroring Failed

Failed to mirror PR to GitLab. Check logs for details.

@coderabbitai

coderabbitai Bot commented Sep 28, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

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 configuration

Configuration used: defaults

Review profile: CHILL

Plan: Advanced

Run ID: 0f36d517-a6a1-47ff-908e-cf79d4a6388b

📥 Commits

Reviewing files that changed from the base of the PR and between be6a019 and 76d586f.

📒 Files selected for processing (53)
  • csrc/cake_dsv4/sm_100a/cake_dsv4_02dcf18075202d4e41b7_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_02dcf18075202d4e41b7_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_31a5b7b148bf73128281_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_31a5b7b148bf73128281_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_373680a69d10cbe081ff_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_373680a69d10cbe081ff_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_556382f3cbe2b21a7c90_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_556382f3cbe2b21a7c90_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_7884ef54aba94349ba0b_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_7884ef54aba94349ba0b_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_9ba23e3d8ac2ccc5e723_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_9ba23e3d8ac2ccc5e723_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_a7e76a07fad26b1285d9_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_a7e76a07fad26b1285d9_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_b21faf78133754869fa6_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_b21faf78133754869fa6_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_b70757c2e21f685e6ecd_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_b70757c2e21f685e6ecd_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_c7e7d66f92cb51dff469_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_c7e7d66f92cb51dff469_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_family_2091bf44aeea91f0d627_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_family_7d305c261561d61898e4_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_family_99530fece09ed4244850_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_fd1d5cb94793406b50f7_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_0dbf5a57379323ec656e_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_0dbf5a57379323ec656e_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_1d36f73de1730da32e19_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_1d36f73de1730da32e19_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_2e5afca2df7fbbc114c5_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_2e5afca2df7fbbc114c5_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_3be84e230bc2c2790f22_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_3be84e230bc2c2790f22_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_5f97b81e1638ea9bd249_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_5f97b81e1638ea9bd249_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_5fdd9977213d6dcdb5ec_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_5fdd9977213d6dcdb5ec_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_74a066ed790a5f79060f_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_74a066ed790a5f79060f_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_7d00f02c7c345ea5542c_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_862a64e07a59d7930f92_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_862a64e07a59d7930f92_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_bb29ac96275e69cf7673_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_bb29ac96275e69cf7673_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_c0bb86b63c5d6405a5ae_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_c0bb86b63c5d6405a5ae_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_ccd0d7cc7e830dbb1b46_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_ccd0d7cc7e830dbb1b46_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_family_643b127a24e434244cbb_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_family_be0f6094245e783eeb40_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_family_c141f529c51d80c4364d_binding.cu
  • flashinfer/jit/cake_dsv4.py
  • flashinfer/mla/cake_dsv4.py
  • tests/mla/test_cake_dsv4.py
💤 Files with no reviewable changes (16)
  • csrc/cake_dsv4/sm_103a/cake_dsv4_1d36f73de1730da32e19_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_74a066ed790a5f79060f_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_c7e7d66f92cb51dff469_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_9ba23e3d8ac2ccc5e723_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_b70757c2e21f685e6ecd_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_ccd0d7cc7e830dbb1b46_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_5fdd9977213d6dcdb5ec_kernel.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_ccd0d7cc7e830dbb1b46_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_9ba23e3d8ac2ccc5e723_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_b70757c2e21f685e6ecd_binding.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_family_2091bf44aeea91f0d627_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_5fdd9977213d6dcdb5ec_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_7d00f02c7c345ea5542c_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_family_643b127a24e434244cbb_binding.cu
  • csrc/cake_dsv4/sm_103a/cake_dsv4_1d36f73de1730da32e19_kernel.cu
  • csrc/cake_dsv4/sm_100a/cake_dsv4_c7e7d66f92cb51dff469_kernel.cu

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 6 remain after this review.


📝 Walkthrough

Walkthrough

The 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.

Changes

CAKE DSv4 kernel updates

Layer / File(s) Summary
Five-stage KV pipeline and value-quarter processing
csrc/cake_dsv4/sm_100a/*_kernel.cu, csrc/cake_dsv4/sm_100a/*_binding.cu, csrc/cake_dsv4/sm_103a/*_kernel.cu, csrc/cake_dsv4/sm_103a/*_binding.cu
Several kernels expand their KV pipeline to five stages, revise shared-memory and barrier layouts, and change value-quarter selection in MMA and K gathers. Launchers request the matching shared-memory sizes.
Index staging and buffering
csrc/cake_dsv4/sm_100a/*_kernel.cu, csrc/cake_dsv4/sm_100a/*_binding.cu, csrc/cake_dsv4/sm_103a/*_kernel.cu, csrc/cake_dsv4/sm_103a/*_binding.cu
Index-loading paths stage values before shared-memory writes and add double-buffered index slots with slot-specific barriers.
Softmax, correction gathers, and merging
csrc/cake_dsv4/sm_100a/cake_dsv4_b21faf78133754869fa6_*, csrc/cake_dsv4/sm_103a/cake_dsv4_5f97b81e1638ea9bd249_*
Softmax paths preload sink values. Correction warps gather valid KV index groups, loader gather scheduling changes, and split-arrival counters reset before merge work.
Vectorized BF16 output stores
csrc/cake_dsv4/sm_100a/cake_dsv4_{7884ef54aba94349ba0b,a7e76a07fad26b1285d9}_*, csrc/cake_dsv4/sm_103a/cake_dsv4_{3be84e230bc2c2790f22,c0bb86b63c5d6405a5ae}_*
Output paths process 16 values in one pass and store eight packed BF16 words with a vectorized store.

JIT and dispatch updates

Layer / File(s) Summary
JIT registrations and obsolete kernels
flashinfer/jit/cake_dsv4.py, csrc/cake_dsv4/sm_100a/*, csrc/cake_dsv4/sm_103a/*
Generated program identities, source references, dependencies, and selected compile flags change. Obsolete split-3, fixed-query, and dedicated reducer registrations and their generated files are removed.
MLA dispatch and tests
flashinfer/mla/cake_dsv4.py, tests/mla/test_cake_dsv4.py
Dense BF16/H64 queries use the shared routing rules. Selected BF16/H128 top-k widths retain the default one-split route. Routing and compile-flag tests are updated.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~60 minutes

Suggested reviewers: aleozlx, yzh119, bkryu

Merge Risk: ⚪ Minimal · up to 76d58

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 Review

Security architecture risk: 🔵 Low · up to 76d58

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
No architecture-level concerns identified.

Security review details

Security Blast Radius

  • inferred — The directly affected security scope is caller-supplied sparse-index metadata and KV/output buffers processed by the selected GPU decode programs. The supplied change evidence identifies no new service, privilege, or deployment boundary.

Trust Boundaries and Controls

  • observed — The inspected family binding retains host-side launch validation and forwards to the renamed internal kernel. Its exported route_of selects a route without launching the kernel or accessing device buffers.

Resilience and Maintainability Implications

  • inferred — The visible slot waits, publication fence, and empty signals support the normal producer-consumer lifecycle, but the available source and summaries do not establish GPU execution behavior after interruption or failure, or equivalence with the previous kernel.
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 19.23% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 52 functions across 35 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the round-three DSv4 sparse-MLA performance work, target architectures, and main implementation areas. It is long but remains specific and relevant.
Description check ✅ Passed The description provides a detailed summary, performance rationale, baseline context, routing changes, and extensive validation results. It does not include every template heading, such as Related Iss…
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
  • Fix all pre-merge checks with AI
✨ Finishing Touches 💡 1
🛠️ Fix failing CI checks 💡
  • Commit to this branch
  • Create a new PR
🧪 Generate unit tests (beta)
  • Create a new PR

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

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@yyihuang

Copy link
Copy Markdown
Collaborator Author

/bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

GitLab MR !1755 has been created, and the CI pipeline #70181549 is currently running. I'll report back once the pipeline job completes.

@yyihuang

Copy link
Copy Markdown
Collaborator Author

/bot run tests/mla/test_cake_dsv4.py tests/mla/test_cake_dsv4_hardening.py

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

GitLab MR !1755 has been created, and the CI pipeline #70187738 is currently running. I'll report back once the pipeline job completes.

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

[SUCCESS] Pipeline #70187738: 25/25 executed test jobs passed

@yyihuang
yyihuang merged commit ebe0efb into flashinfer-ai:main Sep 28, 2026
27 of 32 checks passed
yyihuang added a commit that referenced this pull request Sep 28, 2026
…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>
yyihuang added a commit that referenced this pull request Sep 29, 2026
…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>
yyihuang added a commit that referenced this pull request Sep 30, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants