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 by yyihuang · Pull Request #5591 · flashinfer-ai/flashinfer

yyihuang · 2026-09-27T05:37:54Z

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

Hardens the backend="cake" DeepSeek-V4 sparse-MLA decode path of trtllm_batch_decode_sparse_mla_dsv4
(SM100a / SM103a) for the asks in #4671 and re-exports the generated programs:

padded query rows beyond the metadata token count are ignored (no metadata read past num_tokens, -1 never
dereferenced, padded output rows untouched);
separate SWA and compressed index tables (extra_sparse_indices, extra_sparse_topk_lens) are consumed in-kernel,
the combined table keeps working;
caller-owned workspace: get_cake_dsv4_workspace_bytes(...), no device allocation on any call path (MTP, CUDA graph
capture), split counters self-reset;
raw lengths plus sparse_topk_lens_offset;
many-token (MTP) rows route to the bodies that beat trtllm-gen (16-512 tokens): BF16/H128 SWA and SWA+topk4x
rows use the persistent prefill body from 64 tokens (_BF16_H128_PREFILL_MIN_TOKENS), BF16/H128 topk128x width
260/388 rows above 16 tokens use the row-first single-owner program (_BF16_TOPK128X_SPLIT_MAX_TOKENS);
FP8/H64 rows admitted to the persistent body with >= 128 tokens use the H64-specific single-CTA M64 program
(fp8_h64_prefill_source_persistent_m64, _FP8_H64_M64_MIN_TOKENS); BF16/H32 topk4x rows share the retained-KV
topk128x program (bf16_h32_topk128x_early_v47, last-arriver split merge);
BF16/H128 SWA+topk4x rows whose persistent-grid tail clusters are the majority launch the boustrophedon work-feed program
of the prefill body (bf16_h128_prefill_v42_snake, _bf16_h128_prefill_uses_snake_feed);
the boustrophedon program bf16_h128_prefill_v42_snake is regenerated from a K-reuse body (the PV operand is served from the
gathered K tile, 25 % fewer L2 bytes per KV tile); the striped program is byte-identical to the previous export;
the FP8/H128 persistent programs rescale the O accumulator lazily (the running max only advances when a tile raises it by more
than 8 in the exp2 domain, so most tiles skip the TMEM rescale on the serial PV chain; P stays inside e4m3, O / sum unchanged);
the BF16/H128 persistent prefill program drains O with 32-byte stores (half the LSU instruction count of the epilogue tail);
FP8 H8/H16 canonical 12-token rows (batch 3, ragged q_len 5, widths 192/256 topk4x or 260 topk128x) use the FP8 port of the
1-CTA SwapsAb body (fp8_h8_h16_source_exact, heads on the MMA N side), 1.24-1.49x vs the default path on both arches;
12-token decode routing follows the Cake seeds: FP8 low-head widths up to 384 -> one-partition producer,
FP8/H64 rows with >= 3 complete sparse tiles -> the persistent FP8 body, BF16/H128 topk128x width 260/388 -> the
two-stage split3/split4 programs on both targets.

Also hardens the default trtllm-gen host path (#4671 P1/P4, see flashinfer/mla/_core.py).

Validation

Export r9 (Cake 421f1a5671a, this branch at f63041c + the regenerated programs), 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.24x, min 1.006x	94/94, geomean 1.23x, min 1.013x
40 hardening rows faster than the default path	40/40, geomean 1.28x, min 1.016x	40/40, geomean 1.25x, min 1.029x
compute-sanitizer synccheck + memcheck on the changed kernels	0 errors	0 errors
`tests/mla/test_cake_dsv4*.py` (GPU, on the exported programs)	292 passed	292 passed
`tests/mla/test_cake_dsv4*.py` (CPU, this host)	208 passed	208 passed

The default trtllm-gen path itself is hardened for #4671 P1/P4 in flashinfer/mla/_core.py.

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: Add DeepSeek V4 sparse MLA TRTLLM-GEN kernels #3269 (9c76c99, 2026-05-21, flashinfer/decode.py, flashinfer/mla/_core.py,
DSv4 TRTLLM-GEN kernels); 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).
Cake backend under test = backend="cake" from feat(cake_sparse_mla): support SM100 and SM103 #4573 (13db2cf, 2026-09-16, flashinfer/mla/cake_dsv4.py,
flashinfer/jit/cake_dsv4.py, csrc/cake_dsv4/), regenerated by this PR.
Test oracle = the PyTorch reference of tests/mla/test_cake_dsv4.py / test_cake_dsv4_hardening.py (DSv4 RopeQuant
exposure feat(mla): expose DSv4 TRTLLM-GEN RopeQuant #4918, 27d5b02, 2026-09-07).
Branch merge-base with upstream main at measurement: d3af75a. The branch has since merged upstream main
c0ea2a1; of the upstream commits in between, the only ones touching flashinfer/mla/_core.py (feat(cake_mla): add generated Blackwell TRT-LLM MLA decode backend for SM100 and SM103 #4557, feat(cake_kimi_k3_mla): Cake-generated Kimi-K3 MLA FP8 paged attention backend for SM100/SM103 (backend="cake") #5552, feat(cake_mla_varq_dcp_decode): Cake MLA variable-query decode with DCP for SM100 / SM103 #5577) change
_trtllm_batch_decode_with_kv_cache_mla_impl (other backend="cake" MLA decode paths) and leave the DSv4 default path and
trtllm_batch_decode_sparse_mla_dsv4 unchanged, so the comparison above still describes the merged head.

CI

tests/mla/test_cake_dsv4.py and tests/mla/test_cake_dsv4_hardening.py (CPU and GPU parametrizations).

🤖 Generated with Claude Code

References

The shared reference corpus for CAKE kernel development includes the following projects, documentation, and existing CAKE work:

GPU programming and instructions: NVIDIA CUDA Programming Guide, NVIDIA PTX ISA, and gau-nernst's tcgen05 tutorial and CUDA kernel examples.
Kernel programming libraries and compilers: NVIDIA CUTLASS / CuTe, Triton, TileLang, NVIDIA cuTile Python, and ThunderKittens (ThunderKittens 2.0 techniques).
Attention and inference: FlashAttention (including Hopper and CuTe implementations), FlashInfer, including its TRT-LLM kernel integration, Flash Linear Attention, SageAttention, FlashAttention-FP4, and FastVideo.
GEMM and MoE: DeepGEMM, SonicMoE, Alpha-MoE, and Mixture of Kittens.
Clustering and nearest-neighbor kernels: Flash K-Means and FlashLib.
Existing CAKE implementations and PRs: CAKE-generated kernel progress tracker and PR index (#4254).

These are corpus-level references. PR-specific implementation details, changes, and benchmark references are documented above.

…ata ABI (flashinfer-ai#4671) Rewrite flashinfer/mla/cake_dsv4.py around the regenerated kernel ABI: - Bind kernel arguments by name through the registration arg_plan / program signature (tma_buffer/buffer -> tensors, parameter -> ints, workspace -> descriptor slab, grid -> launch grid). Every variant receives the uniform metadata vocabulary (swa_indices, compressed_indices, sparse_topk_lens, swa_index_stride, compressed_index_stride, sparse_topk_lens_offset, sparse_topk, num_query_tokens); only the names in its plan are passed. Pre-hardening `sparse_indices` binds only for a combined table with offset 0; `completion_base` is rejected. - Resolve metadata like the Cake helper without copies: combined table -> two row-strided views, separate SWA [T,128] + extra [T,topk_c] tables, compressed-only lens => offset 128, explicit sparse_topk_lens_offset. num_query_tokens = metadata rows; query/out are prefix views when they carry more rows. Non-unit column stride / non-int32 tables raise. - Carve ONE caller-owned workspace deterministically: [TMA descriptor slab 1 KiB | split-merge counters 256 KiB | partial_O | partial_lse]. Publish get_cake_dsv4_workspace_bytes(), cake_dsv4_workspace_layout() and cake_dsv4_workspace_reset(). No device allocation on the call path; counter priming is refused during CUDA graph capture with a clear error. Remove the host-side _partition_arrivals generation state and the per-call _padded_sparse_indices copy. - _core.py: backend="cake" accepts extra_sparse_indices, extra_sparse_topk_lens and the new sparse_topk_lens_offset kwarg, and tolerates metadata with fewer rows than the query (padded batches). Docstring documents the cake contract. trtllm-gen/sparse/cute-dsl paths are unchanged. - Tests: CPU binding/workspace/metadata/capture tests in tests/mla/test_cake_dsv4.py (117 passing) and GPU hardening tests in tests/mla/test_cake_dsv4_hardening.py (padded rows, separate tables, lens offset, all -1 rows, CUDA graph replay + memory_allocated) for H in {32,64,128}, BF16/FP8, q_len {1,4}. - csrc/cake_dsv4/README.md: API additions and workspace contract. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…r#4671 P1/P4) - query/out rows beyond the metadata tables (padded DP/MTP batches) are sliced to prefix views on the host; the cubin never reads padded metadata and never writes padded output rows (dense padding must be whole requests). - new multi_ctas_kv_counter_buffer kwarg: a caller-owned, zero-initialized buffer validated against get_trtllm_gen_multi_ctas_kv_counter_bytes; the per-call allocation remains only when nothing is passed. - GPU tests for both behaviours (padded rows dense/ragged, counter reuse over launches, size validation, backend guard, CUDA-graph replay after mutation). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…te tables With extra_sparse_topk_lens (separate SWA/compressed tables) the public API validated the compressed lengths through _check_dsv4_sparse_mla_inputs and then forwarded that checked tensor as sparse_topk_lens *and* the caller's extra_sparse_topk_lens, so run_cake_dsv4 rejected every separate-table call ("pass either sparse_topk_lens ... or extra_sparse_topk_lens ..., not both"). Forward the combined table only when the caller supplied one. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…BF16 H128 SWA rows (cake backend) - fp8_h128: split the KV range only as far as one wave of 2-CTA clusters needs (min(5, (sm_count // 2) // tokens)); the fixed 5-way split ran 4-5 waves of single-tile work items on 16-64 token MTP rows. - BF16 H128 SWA-only rows with 128+ metadata tokens use the full-V H128 family (bf16_h128_topk128x), ~2x faster than the dedicated SWA producer on 128-512 token rows; shorter rows keep the dedicated producer. Both rules mirror the Cake dispatcher. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Mirror of the Cake producer rule measured on GB300/B200: <= 3 tiles unsplit, 4-5 tiles within 32 clusters, 6+ tiles up to five splits within one wave. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…ent bodies Mirror of the Cake dispatcher: FP8/H128 rows with 16+ tokens use the persistent prefill body, BF16/H64 compressed rows with 24+ tokens and width <= 640 use the KV-reuse prefill body; canonical 12-token rows are unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Mirror of the Cake dispatcher: the persistent fp8_h128_prefill_source_persistent body beats the best forced split of the cluster producer on every FP8/H128 row (GB300 canonical 12/12 above trtllm-gen, MTP rows 1.10-1.75x), so the fp8_h128 split route, _fp8_h128_num_splits and the sm_count binding are removed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…P8/H128 route The generated bindings check every pointer buffer for contiguity. The host handed the compressed table of a combined sparse_indices table over as a column-sliced view, which the binding rejected on every combined-table call (46 GPU hardening tests). SparseMetadata.kernel_kwargs() now exposes both tables as contiguous spans with the same base pointer and row stride, no copy. fp8_h128_prefill_source_persistent reads request boundaries from cum_seq_lens_q only; dense queries reached it with no offsets. run_cake_dsv4 synthesizes arange(batch+1) * max_q_len for ragged-only routes from a cached process-lifetime tensor (no per-call allocation, stable pointer for CUDA graphs). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…ragged requests The fp8_lowhead_split / fp8_lowhead_h64_split producers index their ``O`` argument as [tokens, heads, splits, 512] (the Cake dispatcher hands them its partial buffer) and the reducer writes the caller's rows. The host bound ``O`` to the caller's output, so the producer overran it by one output's worth and the reducer merged never-written workspace partials: wrong FP8 H32 (both archs) / H64 (SM100) results and memory corruption of neighbouring tensors in the hardening tests. trtllm-gen path: with ragged queries the kernel takes its token count from cum_seq_lens_q, so a padded batch (metadata rows < query rows) still wrote the padded output rows. Cut cum_seq_lens_q/seq_lens at the request boundary that equals the metadata row count; padding must drop whole requests. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…roducer Mirrors the Cake seed's FP8_ONE_PARTITION_MAX_TILES = 3: three sparse tiles are owned by one producer partition (final O written directly) instead of the 2 + 1 split plus reducer. Route tests updated; the width-640 case keeps the split path. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…to the persistent body The persistent FP8 body (same kernel as FP8/H128; heads >= num_heads are predicated and the second CTA's OOB Q rows zero-filled) measured GB300 1.23-1.26x / B200 1.15x vs trtllm-gen on the width-640/388 FP8/H64 decode rows where the per-token cluster body sat at 0.95-1.04x (CAKE-624 W5). The rule mirrors the Cake seed (H64_PERSISTENT_MIN_TILES = 3 complete tiles, H64_PERSISTENT_MAX_TOKENS = 12). Width-260 and many-token rows keep the H64 producers. The route test now derives the canonical ragged token count from the contract's linspace rule instead of batch * max_q_len. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

… sm_103a too The two-stage split programs (three or four disjoint full-V KV owners plus one LSE reducer) now ship on both Blackwell targets: GB300 width-260/388 rows measured 1.18-1.26x vs trtllm-gen against 0.83-1.05x for the single-owner kernel (CAKE-624 W2). The sm_100a dispatch is unchanged. A recorder test pins the program and split count per width on both targets. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Mirrors the Cake seed rule (CAKE-624 W5 + W10): 12-token rows with >= 3 complete sparse tiles, every many-token compressed row (>= 2 tiles) and SWA-only rows from 128 tokens take the persistent FP8 body already shipped for FP8/H128 (num_heads is a runtime argument there). The per-token cluster producers and the SWA producer sat at 0.19-0.87x vs trtllm-gen on the 64-512 token hardening rows; the persistent body measures 1.02-1.92x on the 16-128 token rows on both targets. Evaluated before the SWA-only test so one-tile many-token rows reach it. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…single owner Mirrors the Cake seed's BF16_TOPK128X_SPLIT_MAX_TOKENS (CAKE-624 W12): the split3/split4 two-stage programs keep the 12/16-token rows; width 260-388 rows with more query tokens select bf16_h128_topk128x_row_first (one full-V owner per token whose -1 sparse rows gather the tile's first index). hardening-000025/31 0.45-0.95x -> 1.13-1.65x vs trtllm-gen on both targets. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…stent prefill body Mirrors the Cake dispatcher's BF16_H128_PREFILL_MIN_TOKENS (CAKE-624 W11): the 64-512 token MTP rows leave the per-token owners (0.39-0.91x vs trtllm-gen) for the KV-reuse prefill body (1.00-1.18x); the former SWA -> full-V rule at 128 tokens (0.55-0.73x) is removed. The exported bf16_h128_prefill_v42 program is regenerated with the seed's width guard relaxed to 128..1152 and the QK-warp exit-wait parity fix. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

The retained-KV body with the last-arriver split merge beats the dedicated topk4x body on the H32 topk4x rows on both SM100a and SM103a, so both H32 sparse routes select bf16_h32_topk128x_early_v47. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…CTA M64 program Mirrors the Cake seed's H64_M64_MIN_TOKENS: FP8/H64 rows admitted to the persistent body with at least 128 query tokens select fp8_h64_prefill_source_persistent_m64 (one CTA per token, grid = tokens, ragged query offsets synthesized like the FP8/H128 persistent route); the other admitted rows keep fp8_h128_prefill_source_persistent. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…tch-derived routes Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…r program Mirrors the Cake seed's LANE_GATHER_MIN_TOKENS: the persistent FP8 body has two exported programs, fp8_h128_prefill_source_persistent_uniform (elected-lane uniform K/V gathers, below 128 query tokens) and fp8_h128_prefill_source_persistent (lane-issued gathers, 128 tokens and up). Same launch block, grid = 2 * tokens, ragged query offsets synthesized. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Export of the hardened DSv4 family on B200 against this branch: 134 shapes (94 canonical + 40 hardening), all gates passed, trtllm-gen default path geomean 1.220x over the 94 canonical rows, source/export agreement 1.009x. Program files no longer referenced by the regenerated registry are removed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Export of the hardened DSv4 family on GB300 against this branch: 134 shapes, correctness and source/export agreement on every row, trtllm-gen default path geomean 1.237x over the 94 canonical rows (prefill-style 089/093 at 0.990x and 0.993x are below parity and documented in the Cake design doc). Program files no longer referenced by the regenerated registry are removed. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…8 prefill body for tail-majority rows Mirrors the Cake seed's bf16_h128_prefill_uses_snake_feed (CAKE-624 W17): with C = min(T, SMs // 2) clusters, the striped feed leaves a short heavy prefix to the few regular clusters when the tail clusters are the majority, so rows such as hardening-000037 (512 tokens, width 640) launch bf16_h128_prefill_v42_snake instead of bf16_h128_prefill_v42. One-tile items and regular-majority rows keep the striped program. Predicate and program-selection tests. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Adds the boustrophedon work-feed program of the BF16/H128 prefill body (bf16_h128_prefill_v42_snake, family 63f33aae) generated from Cake 20baf19a41e on B200; removes 14 sm_100a family bindings that no registry entry has referenced since the SM100/SM103 regeneration. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Adds the boustrophedon work-feed program of the BF16/H128 prefill body (bf16_h128_prefill_v42_snake, family f7339c08) generated from Cake 20baf19a41e on GB300; removes the sm_103a source files that no registry entry references (the pre-generation SM103 implementation files and family bindings from flashinfer-ai#4573, superseded by the generated programs). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…exact program Host route fp8_h8_h16_source_exact (same shape lock as the BF16 source-exact route) and its launch shape; route tests updated, two off-profile cases keep the one-partition producer. Programs are regenerated in the export commit. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Adds the FP8 H8/H16 SwapsAb source-exact program (fp8_h8_h16_source_exact, family 28ebcdf2) and regenerates the boustrophedon work-feed program of the BF16/H128 prefill body from the K-reuse body (bf16_h128_prefill_v42_snake, family d9a2424a; the striped program is byte-identical), generated from Cake 69bc58ee5ed on B200; removes the superseded snake family 63f33aae and the one-partition program files no registry entry references. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Adds the FP8 H8/H16 SwapsAb source-exact program (fp8_h8_h16_source_exact, family 6950f2bb) and regenerates the boustrophedon work-feed program of the BF16/H128 prefill body from the K-reuse body (bf16_h128_prefill_v42_snake, family 1d469eb5; the striped program is byte-identical), generated from Cake 69bc58ee5ed on GB300; removes the superseded snake family f7339c08 and the one-partition program files no registry entry references. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Regenerates the BF16/H128 persistent prefill program (32-byte epilogue O stores) and the FP8/H128 persistent programs (lazy O rescale, variant 4620) from Cake 421f1a5671a on B200; removes the superseded program files no registry entry references. Paired vs the default trtllm-gen path on B200: 94/94 canonical rows > 1.0 (geomean 1.226x). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Regenerates the BF16/H128 persistent prefill program (32-byte epilogue O stores) and the FP8/H128 persistent programs (lazy O rescale, variant 4620) from Cake 421f1a5671a on GB300; removes the superseded program files no registry entry references. Paired vs the default trtllm-gen path on GB300: 94/94 canonical rows > 1.0 (geomean 1.237x; prefill-style-000089 1.018x, 000093 1.006x). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

flashinfer-bot · 2026-09-27T05:38:46Z

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

flashinfer-bot · 2026-09-27T07:50:44Z

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

…grams (warp-uniform FP8/H128 gathers, 32-byte epilogue stores in the K-reuse prefill program) Regenerated from Cake db0609184cd against 8589d49 with the same export pipeline as flashinfer-ai#5591 (134 shapes per architecture: 94 canonical + 40 hardening; correctness, source/export agreement and the trtllm-gen baseline gates all pass on B200 and GB300). - FP8/H128 persistent prefill program: the K/V loader warps issue their TMA gather4 from one elected lane with warp-uniform coordinates (shfl broadcast) instead of one lane-divergent gather per lane; the persistent seed compiles with --register-usage-level=10 so the exported build keeps the production softmax schedule. prefill-style rows 089/093: GB300 1.02x -> 1.24x/1.21x, B200 1.05x -> 1.19x/1.18x vs the trtllm-gen default path. - BF16/H128 K-reuse (snake) prefill program: 32-byte epilogue O stores. hardening rows 27/37: GB300 1.175x -> 1.19x, 1.016x -> 1.107x; B200 1.172x -> 1.173x, 1.028x -> 1.100x. - All other programs are byte-identical regenerations (renamed for the new producer revision); 94/94 canonical rows > 1.0 on both architectures (geomean GB300 1.243x, B200 1.232x). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…epilogue stores in the DSv4 sparse-MLA Cake programs (SM100/SM103) (#5610) ## Summary Round-2 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 #5591: - **FP8/H128 persistent prefill program (prefill-style rows with 128+ query tokens, batch 2-48 items x 16 KV tiles):** the two K and two V loader warps issued their TMA `gather4` with lane-divergent coordinates, so ptxas serialised every gather into a 32-iteration per-lane loop and the loaders spent ~1.0 us of every 1.12-us KV tile issuing. They now issue from one elected lane with warp-uniform coordinates (`shfl.sync` broadcast, one unrolled 16-quartet site per warp with the KV descriptor selected per tile), the same form the default trtllm-gen kernels use. Same barriers, pipeline depth and data movement; only the issuing lane changes. The loop is now paced by the tensor pipe (16 QK + 8 PV MMAs per tile, ~0.8 us). prefill-style-000089 / 000093: GB300 1.02x / 1.01x -> **1.24x / 1.21x**, B200 1.05x / 1.06x -> **1.19x / 1.18x** vs the default path. The persistent seed compiles with `-Xptxas=--register-usage-level=10` so the exported build keeps the softmax schedule of the source build (without it ptxas re-orders the exp2/convert chains across the P-publication fence: 3-5 % slower on these rows). - **BF16/H128 K-reuse (snake) prefill program (`bf16_h128_prefill_v42_snake`, hardening rows 27 / 37):** the epilogue drains O with 32-byte stores like the striped program already did. hardening-000037: GB300 1.016x -> **1.107x**, B200 1.028x -> **1.100x**; hardening-000027: 1.175x -> 1.190x, 1.172x -> 1.173x. - Every other program is a byte-identical regeneration (renamed for the new producer revision). - New CPU test: registered programs build with public nvcc/ptxas options only, and the FP8/H128 persistent variants pin the register-usage level. The one-KV-tile SWA-only work items of the BF16/H128 persistent prefill body (hardening rows 23 / 29 / 35) were re-attributed and left unchanged: the item is a per-SM chain bound by the SM's L2<->SM port (256 KB per CTA per item at the measured ~66 GB/s = the measured 3.9 / 4.1 us period); the byte levers (K-reuse body, four issuing warps, DSM exchange of the V halves) measured negative or closed by arithmetic. Those rows stay at 1.03-1.16x. ## Validation Export r11 (Cake `db0609184cd`, this branch at 8589d49 + the regenerated programs), 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.243x, min 1.054x** (#5591: 1.237x, min 1.006x) | **94/94, geomean 1.232x, min 1.041x** (#5591: 1.226x, min 1.013x) | | 40 hardening rows faster than the default path | **40/40, geomean 1.30x, min 1.058x** (#5591: 1.28x, min 1.016x) | **40/40, geomean 1.25x, min 1.037x** (#5591: 1.25x, min 1.029x) | | exported program vs Cake source build (agreement gate, 134 rows) | geomean 1.002x, every row within 3 % | geomean 1.004x, every row within 3 % | | compute-sanitizer synccheck + memcheck on the two changed kernels | 0 errors | 0 errors | | `tests/mla/test_cake_dsv4*.py` (GPU, on the exported programs) | 292 passed | 292 passed | | `tests/mla/test_cake_dsv4*.py` (CPU, this host) | 210 passed | 210 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, `flashinfer/decode.py`, `flashinfer/mla/_core.py`, DSv4 TRTLLM-GEN kernels); 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) as hardened and regenerated by #5591 (8589d49, 2026-09-27), regenerated again by this PR. - **Test oracle** = the PyTorch reference of `tests/mla/test_cake_dsv4.py` / `test_cake_dsv4_hardening.py` (DSv4 RopeQuant exposure #4918, 27d5b02, 2026-09-07). - **Branch base** = upstream `main` 8589d49 (the #5591 merge); no upstream commits in between. ## CI `tests/mla/test_cake_dsv4.py` and `tests/mla/test_cake_dsv4_hardening.py` (CPU and GPU parametrizations). 🤖 Generated with [Claude Code](https://claude.com/claude-code)  ## Summary by CodeRabbit - **Bug Fixes** - Improved correctness of CAKE DSV4 attention kernels by refining shared-memory synchronization, sparse gather operations, and pipeline phase handling. - Updated output processing to use more efficient vectorized stores in selected kernels. - Corrected kernel registration and dispatch across supported GPU architectures, including updates to available variants and compile settings.  --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

…nel (redundant p_full barrier removed) (#5627) 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 Regenerates the Cake-generated `backend="cake"` Kimi-K3 MLA FP8 paged attention route (#5552; r52 programs in #5602) for the **r54 kernel** on SM100 (B200) and SM103 (B300 / GB300). The host side is unchanged from #5602 (same planner, grids and kernel arguments); this PR moves only the generated programs of both architectures, every file written by the exporter from the same producer tree, and the two architectures' exports agree byte-for-byte (`csrc` and the registry). What r54 changes against the r52 programs (`698a96197`): - **The redundant `p_full` mbarrier is removed from the wide two-CTA kernel.** Since r52 the softmax / correction warps arrive `o_ready` (tile parity, 4 warps x 2 CTAs, cluster arrive to CTA 0) right after the P store's `fence.proxy.async` at the same site where they arrived `p_full`, and the MMA warp waited both barriers at the same point before PV(t) and before the drain; `p_full` therefore carried no ordering of its own. PV(t) is now gated by `o_ready` alone: one fewer cluster arrive per warp per tile, one fewer wait per tile in the MMA warp, 45 (was 47) mbarriers. Outputs are unchanged (same FP32-reference max_abs on every checked row). - Measured against the r52 kernel of record on the same nodes (paired, arm-alternated, 3 rounds, CUDA-graph cold-L2 CUPTI timing, cute-dsl anchor): **neutral within the session noise floor on both GPUs** (B300 0.994-1.006 of the record over the 11 wide and prefill rows, noise floor 0.24 % median / 1.56 % max; B200 0.991-1.019, noise floor 0.26 % / 3.60 %). The change ships as the simpler canonical form, not as a speed-up. - The round's other candidate, pacing the row-tile clusters of one request so they share K tiles in L2, was measured as a loss in four configurations (+30 % to +156 % on mtp_h96_b8_q4) and does not ship; hardware counters show the re-read it targeted is intermittent and off the critical path of that row. ## Performance (same-session CUPTI cold-L2, CUDA-graph replay, 3 alternating rounds; Cake vs the fastest existing FlashInfer route) | row | B200 Cake us | B200 best FI (route) | B200 speedup | B300 Cake us | B300 best FI (route) | B300 speedup | |---|---:|---:|---:|---:|---:|---:| | trace_h12_b1_q1_kv76800 | 21.2 | 24.7 (cute-dsl) | 1.166x | 20.9 | 23.6 (cute-dsl) | 1.128x | | trace_h12_b2_q1_kv77699 | 29.3 | 31.0 (cute-dsl) | 1.060x | 28.9 | 30.7 (cute-dsl) | 1.060x | | trace_h12_b3_q1_kv261120 | 89.7 | 93.1 (cute-dsl) | 1.039x | 90.3 | 92.6 (cute-dsl) | 1.026x | | trace_h12_b4_q1_kv268800 | 112.9 | 123.4 (cute-dsl) | 1.093x | 113.6 | 119.9 (cute-dsl) | 1.056x | | trace_h12_b5_q1_kv276480 | 138.4 | 149.0 (trtllm-gen) | 1.076x | 139.7 | 150.4 (trtllm-gen) | 1.077x | | trace_h12_b6_q1_kv467751 | 254.0 | 272.6 (trtllm-gen) | 1.073x | 257.1 | 274.1 (trtllm-gen) | 1.066x | | trace_h12_b7_q1_kv299613 | 195.4 | 208.6 (trtllm-gen) | 1.068x | 197.8 | 210.6 (trtllm-gen) | 1.064x | | trace_h12_b8_q1_kv342305 | 246.5 | 266.0 (trtllm-gen) | 1.079x | 249.4 | 267.1 (trtllm-gen) | 1.071x | | trace_h12_b9_q1_kv348705 | 279.0 | 301.8 (trtllm-gen) | 1.082x | 280.3 | 300.3 (trtllm-gen) | 1.071x | | trace_h12_b10_q1_kv474166 | 412.0 | 446.7 (trtllm-gen) | 1.084x | 411.9 | 442.9 (trtllm-gen) | 1.075x | | trace_h12_b11_q1_kv474405 | 451.9 | 485.7 (trtllm-gen) | 1.075x | 450.5 | 482.2 (trtllm-gen) | 1.070x | | mirror_h96_b1_q1_kv76800 | 23.4 | 27.6 (cute-dsl) | 1.178x | 23.0 | 26.2 (cute-dsl) | 1.136x | | mirror_h96_b2_q1_kv77699 | 33.2 | 34.2 (cute-dsl) | 1.030x | 33.1 | 33.6 (cute-dsl) | 1.016x | | mirror_h96_b6_q1_kv467751 | 324.4 | 340.8 (cute-dsl) | 1.050x | 315.5 | 330.3 (cute-dsl) | 1.047x | | mirror_h96_b7_q1_kv299613 | 246.1 | 259.5 (cute-dsl) | 1.055x | 236.2 | 246.1 (cute-dsl) | 1.042x | | mirror_h96_b8_q1_kv342305 | 318.5 | 332.6 (cute-dsl) | 1.044x | 308.8 | 322.0 (cute-dsl) | 1.043x | | mtp_h12_b8_q2_kv342305 | 256.4 | 311.9 (cute-dsl) | 1.217x | 252.8 | 300.7 (cute-dsl) | 1.189x | | mtp_h12_b8_q4_kv342305 | 284.0 | 316.4 (cute-dsl) | 1.114x | 262.3 | 312.4 (cute-dsl) | 1.191x | | mtp_h12_b8_q8_kv342305 | 316.7 | 331.6 (cute-dsl) | 1.047x | 307.8 | 321.9 (cute-dsl) | 1.046x | | mtp_h96_b8_q2_kv342305 | 611.3 | 685.7 (cute-dsl) | 1.122x | 597.0 | 644.7 (cute-dsl) | 1.080x | | mtp_h96_b8_q4_kv342305 | 915.6 | 984.6 (cute-dsl) | 1.075x | 894.5 | 928.3 (cute-dsl) | 1.038x | | mtp_h96_b8_q8_kv342305 | 1756.1 | 2149.4 (cute-dsl) | 1.224x | 1751.5 | 2026.2 (cute-dsl) | 1.157x | | prefill_h12_b1_q2048_kv131072 | 2485.0 | 2746.1 (cute-dsl) | 1.105x | 2487.4 | 2617.8 (cute-dsl) | 1.052x | | prefill_h12_b2_q512_kv32768 | 323.7 | 389.8 (cute-dsl) | 1.204x | 319.9 | 363.3 (cute-dsl) | 1.136x | | prefill_h12_b4_q1024_kv8192 | 330.8 | 360.4 (cute-dsl) | 1.090x | 312.2 | 348.5 (cute-dsl) | 1.116x | | prefill_h96_b1_q2048_kv131072 | 19914.1 | 20985.0 (cute-dsl) | 1.054x | 20103.1 | 20862.9 (cute-dsl) | 1.038x | | prefill_h96_b2_q512_kv32768 | 2479.5 | 2619.9 (cute-dsl) | 1.057x | 2470.4 | 2613.7 (cute-dsl) | 1.058x | Every row on both GPUs is faster than the best existing FlashInfer route (minimum 1.016x, `mirror_h96_b2` on B300; the same thinnest row as in #5602). Columns: median of three alternating rounds; `best FI` is whichever of trtllm-gen / cute-dsl was faster on that row and GPU (trtllm-gen has no H96 MLA decode kernel, so `auto` resolves to cute-dsl there). ## Baselines and their source PRs The `vs best FI` column compares against the two existing FlashInfer routes of the same public entry points, resolved at this branch's upstream merge-base `f4d45edf1` (2026-09-27, #5596): - **trtllm-gen** (`backend="trtllm-gen"`, also what `backend="auto"` selects for H12): `trtllm_batch_decode_with_kv_cache_mla` / `trtllm_prefill_with_kv_cache_mla` in `flashinfer/mla/_core.py` launching the prebuilt TRT-LLM-gen FMHA cubins pinned in `flashinfer/artifacts.py` (`TRTLLM_GEN_FMHA = 2d6a5a02…/fmha/trtllm-gen/`; pin set by #4654 `857bc3a11`, 2026-08-26). Test oracle: `tests/attention/test_trtllm_gen_mla.py`. - **cute-dsl** (`backend="cute-dsl"`, what `auto` selects for H96): the CuTe-DSL MLA decode family under `flashinfer/cute_dsl/attention/` — `mla_decode_fp8.py`, `mla_config.py`, `wrappers/batch_mla.py` (#2805, `00ac5053d`, 2026-04-13; H64/H96 configs #3235 `2d0e0efac` 2026-05-11; SM107 reland #4280 `5b1af9724` 2026-07-30), `mla_dispatch.py` + `monolithic/mla_decode.py` (#3296 `e91ac8f15` 2026-05-21; LSE base semantics #4650 `d7a7447cb` 2026-09-21), reached through `flashinfer/mla/_core.py` (cute-dsl MLA op #2743 `31b63bc3d` 2026-03-27). Test oracle: `tests/attention/test_trtllm_gen_mla.py` / the CuTe-DSL MLA tests of those PRs. Since the previous merge-base `7b8f6129c` (the #5602 comparison) no upstream commit touched the `cute-dsl` MLA kernels (`git log 7b8f612..f4d45ed -- flashinfer/cute_dsl/attention` is empty), the `TRTLLM_GEN_FMHA` pin is unchanged, and `flashinfer/mla/_core.py` changed only by the sparse-MLA backend hardening of #5591. The Cake route in this PR is compared against exactly those implementations. ## Unsupported (rejected explicitly) Unchanged from #5552 / #5602: sparse top-k MLA, attention sinks, LSE output, DCP, skip-softmax, FP16 softmax, PDL as a caller option (`enable_pdl`: the attention kernel is an ordinary launch on the caller's stream; internally the route launches its split-KV merge as a programmatic dependent of the attention kernel), NVFP4 / uint8 caches, `multi_ctas_kv_counter_buffer`. Requests with > 64 packed rows whose longest KV is below 8192 tokens run the 96-row tiles. ## Tests - `tests/mla/test_cake_kimi_k3_mla.py` (30 cases): **30 passed on B200, 30 passed on B300** on the generated programs of this PR head; the dispatch slice 18 passed (B200) / 12 passed + 6 skipped (B300). Upstream pre-commit hooks pass on the changed files (the only diff they propose is a pre-existing ruff re-wrap of an untouched host line, identical to the one reported in #5602). - Export receipts (exporter's own source-vs-export timing, `source/export` = the Cake launcher over the FlashInfer program, per row and GPU): <details><summary>export receipts, 29 rows x 2 GPUs</summary> | row | route | B200 source us | B200 export us | src/exp | verdict | B300 source us | B300 export us | src/exp | verdict | |---|---|---:|---:|---:|---|---:|---:|---:|---| | mtp_h12_fixed_q5 | main_rt64_reduce_w2 | 302.0 | 302.0 | 1.0002 | PASS | 266.7 | 266.6 | 1.0004 | PASS | | mtp_h96_fixed_q8_kv8000 | main_rt96_reduce_w1 | 139.3 | 139.3 | 1.0000 | PASS | 125.5 | 125.5 | 1.0000 | PASS | | trace_h12_b1_q1_kv76800 | main_rt16_reduce_cta | 21.7 | 21.4 | 1.0135 | PASS | 21.3 | 21.0 | 1.0122 | PASS | | trace_h12_b2_q1_kv77699 | main_rt16_reduce_cta | 29.2 | 29.2 | 0.9989 | PASS | 28.9 | 28.9 | 1.0000 | PASS | | trace_h12_b3_q1_kv261120 | main_rt16_reduce_cta | 89.9 | 89.9 | 0.9996 | PASS | 90.1 | 90.4 | 0.9975 | PASS | | trace_h12_b4_q1_kv268800 | main_rt16_reduce_cta | 112.1 | 112.0 | 1.0006 | PASS | 113.3 | 113.3 | 1.0006 | PASS | | trace_h12_b5_q1_kv276480 | main_rt16_reduce_w4 | 139.2 | 139.5 | 0.9979 | PASS | 140.0 | 140.4 | 0.9973 | PASS | | trace_h12_b6_q1_kv467751 | main_rt16_reduce_w4 | 254.4 | 254.5 | 0.9995 | PASS | 254.5 | 254.7 | 0.9995 | PASS | | trace_h12_b7_q1_kv299613 | main_rt16_reduce_w4 | 196.5 | 196.4 | 1.0010 | PASS | 197.2 | 197.1 | 1.0010 | PASS | | trace_h12_b8_q1_kv342305 | main_rt16_reduce_w4 | 249.7 | 249.9 | 0.9994 | PASS | 248.4 | 248.5 | 0.9995 | PASS | | trace_h12_b9_q1_kv348705 | main_rt16_reduce_w4 | 281.9 | 282.2 | 0.9990 | PASS | 279.7 | 280.0 | 0.9991 | PASS | | trace_h12_b10_q1_kv474166 | main_rt16_reduce_w4 | 428.1 | 428.0 | 1.0002 | PASS | 412.5 | 412.5 | 1.0000 | PASS | | trace_h12_b11_q1_kv474405 | main_rt16_reduce_w4 | 473.8 | 474.0 | 0.9995 | PASS | 448.8 | 449.0 | 0.9996 | PASS | | mirror_h96_b1_q1_kv76800 | main_wide_reduce_cta | — | — | — | FAIL: failed gates: directional_disagreement | 23.0 | 22.9 | 1.0056 | PASS | | mirror_h96_b2_q1_kv77699 | main_wide_reduce_cta | — | — | — | FAIL: failed gates: directional_disagreement | 33.3 | 32.8 | 1.0146 | PASS | | mirror_h96_b6_q1_kv467751 | main_wide_reduce_w2 | 365.0 | 364.4 | 1.0018 | PASS | 318.6 | 318.3 | 1.0009 | PASS | | mirror_h96_b7_q1_kv299613 | main_wide_reduce_w2 | 269.7 | 269.5 | 1.0007 | PASS | 235.5 | 235.1 | 1.0016 | PASS | | mirror_h96_b8_q1_kv342305 | main_wide_reduce_w2 | 357.3 | 357.0 | 1.0011 | PASS | 310.9 | 310.5 | 1.0012 | PASS | | mtp_h12_b8_q2_kv342305 | main_rt32_reduce_w4 | 263.2 | 263.5 | 0.9989 | PASS | 253.0 | 253.2 | 0.9992 | PASS | | mtp_h12_b8_q4_kv342305 | main_rt48_reduce_w2 | 291.0 | 291.3 | 0.9991 | PASS | 262.3 | 262.2 | 1.0002 | PASS | | mtp_h12_b8_q8_kv342305 | main_wide_reduce_w2 | 356.7 | 356.4 | 1.0010 | PASS | 311.9 | 311.5 | 1.0013 | PASS | | mtp_h96_b8_q2_kv342305 | main_wide_reduce_w1 | 669.3 | 669.0 | 1.0006 | PASS | 600.3 | 599.9 | 1.0006 | PASS | | mtp_h96_b8_q4_kv342305 | main_wide_reduce_w1 | 1030.9 | 1030.3 | 1.0006 | PASS | 910.2 | 909.7 | 1.0005 | PASS | | mtp_h96_b8_q8_kv342305 | main_wide_reduce_w1 | 1886.4 | 1886.1 | 1.0002 | PASS | 1779.6 | 1778.7 | 1.0005 | PASS | | prefill_h12_b1_q2048_kv131072 | main_wide_reduce_w1 | 2611.1 | 2611.3 | 0.9999 | PASS | 2488.5 | 2488.1 | 1.0002 | PASS | | prefill_h12_b2_q512_kv32768 | main_wide_reduce_w1 | 336.5 | 336.3 | 1.0009 | PASS | 316.9 | 316.6 | 1.0010 | PASS | | prefill_h12_b4_q1024_kv8192 | main_wide_reduce_w1 | 339.6 | 339.2 | 1.0011 | PASS | 308.6 | 308.3 | 1.0010 | PASS | | prefill_h96_b1_q2048_kv131072 | main_wide_reduce_w1 | — | — | — | FAIL: InterleavedTimingError: reportable loaded SM clock 1222/1965 MHz is below t | — | — | — | FAIL: InterleavedTimingError: reportable loaded SM clock 1215/2032 MHz is below t | | prefill_h96_b2_q512_kv32768 | main_wide_reduce_w1 | 2634.6 | 2633.7 | 1.0004 | PASS | 2490.4 | 2490.0 | 1.0001 | PASS | sm_100a: passed 26/29, perf geomean source/export 1.0007 (n=26) sm_103a: passed 28/29, perf geomean source/export 1.0013 (n=26) </details> ## Baselines-free numerics evidence carried over from the kernel rounds Every route was measured against the FP32 reference on the producer tree of this PR: the routed kernel passes all functional rows at 0.01 / 0.02 on both GPUs (forced-wide and routed plans, 8 rows each), compute-sanitizer synccheck + memcheck and the registered synccheck gate pass, and the 20 unchecked perf rows pass on both GPUs (20 / 20 rows on each GPU, 0 mismatches, max_abs 0.00586 on `prefill_h96_b1` against atol 0.01 / rtol 0.02). The kernel change removes a barrier only; numerics are bitwise those of r52. 🤖 Generated with [Claude Code](https://claude.com/claude-code)  ## Summary by CodeRabbit * **Improvements** * Updated FP8 paged attention execution for supported SM100A and SM103A GPUs. Existing inputs, launch settings, and user-facing configuration remain unchanged. The update changes how synchronization is handled during attention processing; no user-facing API changes are included.  Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

…03 (persistent-item epilogue, index ring, retained-KV decode body, SwapsAb V alias, routing) (#5630) ## 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 - **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, 2026-09-27), 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)  ## 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.  --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

…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)  ## 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.  --------- 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)  ## 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.  --------- 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)  ## 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.  --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

averyhNV and others added 30 commits September 24, 2026 20:05

cake_dsv4: tier the FP8/H128 split count by KV tile count

a63a699

Mirror of the Cake producer rule measured on GB300/B200: <= 3 tiles unsplit, 4-5 tiles within 32 clusters, 6+ tiles up to five splits within one wave. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

cake_dsv4: apply ruff format to the hardened host and tests

73731ac

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

cake_dsv4: README lists the FP8/H64 M64 persistent route among the ba…

0003dcf

…tch-derived routes Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

flashinfer-bot added the op: attention label Sep 27, 2026

yyihuang added the run-ci label Sep 27, 2026

flashinfer-bot added run-ci and removed run-ci labels Sep 27, 2026

yyihuang mentioned this pull request Sep 27, 2026

[CAKE] Long-term CAKE-generated Kernel Progress Tracker #4254

Open

yzh119 approved these changes Sep 27, 2026

View reviewed changes

yyihuang merged commit 8589d49 into flashinfer-ai:main Sep 27, 2026
43 of 48 checks passed

yyihuang mentioned this pull request Sep 27, 2026

[DSv4 TRT-LLM Attention] Accuracy/Perf fixes and API improvement #4671

Closed

yyihuang mentioned this pull request Sep 27, 2026

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

Merged

yyihuang mentioned this pull request Sep 28, 2026

perf(cake_mla): kimi_k3, regenerate both architectures for the r54 kernel (redundant p_full barrier removed) #5627

Merged

yyihuang mentioned this pull request Sep 29, 2026

perf(cake_sparse_mla): round-5 DSv4 sparse-MLA Cake programs for SM100/SM103 (exact numerics; stacked on #5644) #5686

Merged

yyihuang mentioned this pull request Sep 29, 2026

perf(cake_sparse_mla): round-6 DSv4 sparse-MLA Cake programs for SM100/SM103 (exact numerics; stacked on #5686) #5700

Merged

This was referenced Sep 30, 2026

perf(cake_sparse_mla): round-7 DSv4 sparse-MLA Cake programs for SM100/SM103 (numerics-unlocked levers; stacked on #5700) #5728

Closed

perf(cake_sparse_mla): round-8 DSv4 sparse-MLA Cake programs for SM100/SM103 (stacked on #5728) #5742

Closed

yzh119 mentioned this pull request Oct 1, 2026

[cake deslop] cake_dsv4: arch-merge SM100/SM103 copies, drop single-route program wrappers, parametrize split/head clones, fix kernel-less routes #5794

Closed

7 tasks

This was referenced Oct 2, 2026

[Feature] Integrate Cake kernels via FlashInfer: model-by-model tracker sgl-project/sglang#42176

Open

[Feature]: Integrate Cake kernels via FlashInfer: model-by-model tracker vllm-project/vllm#59725

Open

yyihuang mentioned this pull request Oct 16, 2026

[CAKE] Attention Kernel Progress Tracker #6594

Open

3 tasks done

Navigation Menu

Repository navigation

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

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
yyihuang merged 31 commits into
flashinfer-ai:mainfrom
yyihuang:cake-624-dsv4-hardening

yyihuang commented Sep 27, 2026 •

edited

Loading

Uh oh!

flashinfer-bot commented Sep 27, 2026

Uh oh!

flashinfer-bot commented Sep 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

yyihuang commented Sep 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Baselines and their source PRs

CI

References

Uh oh!

flashinfer-bot commented Sep 27, 2026

Uh oh!

flashinfer-bot commented Sep 27, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yyihuang commented Sep 27, 2026 •

edited

Loading