Repository navigation
perf(cake_mla): kimi_k3, regenerate both architectures for the r54 kernel (redundant p_full barrier removed) - #5627
Merged
yyihuang merged 1 commit intoSep 28, 2026
Conversation
…nel (redundant p_full barrier removed) Producer Cake fd3372c3 (main + CAKE-712 r54): the wide two-CTA kernel drops the p_full mbarrier, whose ordering o_ready already carried (same eight producers, arrived after the P store fence at the same site, waited at the same point by the MMA warp). Paired A/B vs the r52 programs is neutral within the session noise floor on B300 and B200; the r54 ceiling round otherwise closes its batch-2 candidates (sibling pacing) as losses and records convergence for the wide-decode and prefill groups. Both architectures' exports agree byte-for-byte (csrc + registry); FlashInfer tests 30 passed on each GPU, dispatch 12+6 skipped (B300) / 18 (B200). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
yyihuang
requested review from
Anerudhan,
Aneureka,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww and
yzh119
as code owners
September 28, 2026 03:24
Collaborator
Author
|
/bot run tests/mla/test_cake_kimi_k3_mla.py |
Collaborator
Author
|
@flashinfer-bot run tests/mla/test_cake_kimi_k3_mla.py |
Contributor
|
Navigate logical layers of code changes, visualize relationships, and explore their blast radius. No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (5)
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 6 remain after this review. 📝 WalkthroughWalkthroughThe SM100A and SM103A FP8 paged-attention kernels have new symbols and reduced shared-memory barrier layouts. Their launch bindings and JIT module records now use the updated symbols, and ChangesFP8 Paged Attention
Priority: ➖ Normal Estimated code review effort: 3 (Moderate) | ~20 minutes Change: Refactor Suggested reviewers: Merge Risk: ⚪ Minimal · up to No actionable issue remains before merge; proceed with normal validation for the updated GPU kernels. Security Architecture ReviewSecurity architecture risk: 🔵 Low · up to The existing architecture-specific route and public launch interface appear preserved, and no new security issue is demonstrated. The removed synchronization still warrants review because its equivalence has not been established for every execution state. Retained concerns Security review detailsSecurity Blast Radius
Trust Boundaries and Controls
Resilience and Maintainability Implications
Hardening Proposals
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Collaborator
yzh119
approved these changes
Sep 28, 2026
Collaborator
|
[SUCCESS] Pipeline #70152258: 25/25 executed test jobs passed |
yyihuang
changed the title
yyihuang
added a commit
that referenced
this pull request
Sep 28, 2026
…as, deferred cluster wait and a 128-wide TP8 tail tile (SM100a/SM103a), regenerated programs (#5647) 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 7 of the Kimi-K3 Stable LatentMoE front / tail projections in `flashinfer/experimental/kimi_k3_latent_moe` (#5575, #5584, tracker #4254): a paired lever ledger on B200 and B300 (A/B in interleaved CUDA-graph groups, cold-L2 CUPTI timing, an A/A control per row, a second independent process on every row under 10 us) adopts five per-row rules in the Cake tree, four of which change the exported programs; every other candidate on the round's list was measured or attributed and closed. - **Tensor-map binding (unchanged in this package).** Every program binds its tensor maps by value (`const __grid_constant__ CUtensorMap` kernel parameters), as the previous rounds did; the package has no TMA-descriptor workspace. The Cake production launcher now binds by value too on every instance except four -- the TP8 decode-tail instances at N_PAD 16 / 32 / 64 and the TP1 T=1024 prefill tail GEMM -- where paired measurements on three node pairs per GPU show the by-value form node-dependent (T=16: +3.8 % on one pair, -4..-8 % on two others; T=32: -3.5 % on one B300 GPU; T=64: -2..-4 %; the T=1024 GEMM -1..-3 % on B200), so those keep the pointer ABI there. The corresponding programs here stay by value; a caller-owned descriptor workspace that would let them bind by pointer is a follow-up, not part of this PR. The kernel symbols and route keys are unchanged by this. - **Packed bf16x2 global-y norm row path** (decode tail, T >= 32 at TP8 and every TP1 row): the latent row stays in `bf16x2` words through the RMSNorm (square via `bf16x2 -> f32` adds, one `rsqrt`, `bf16x2` multiply by the weight); the same roundings in the same places, so `y` and `out` are bitwise identical to the previous chain. Tail TP8 T=32 / 64 / 128: +13.7 / +6.9 / +5.3 % (B200), +9.9 / +5.9 / +4.7 % (B300). - **Landing-zone alias for the T=128 cluster instances** (`_la`): the DSMEM landing zone of the K-split combine is aliased on the drained A ring instead of owning 64 KiB of shared memory, which buys two ring stages (5 -> 7); rank 0 hands "ring drained" to rank 1 through an 8-byte `st.async` word. Only the N_PAD-128 instances that stream >= 32 chunk units per CTA take it: tail TP1 T=128 +8.3 / +7.0 %, front TP8 T=128 +7.9 / +6.4 % (B200 / B300); the short TP8 tail rows lose with the handshake and keep their zone. - **Deferred cluster wait on the T=16 staged-rows tail instance** (`_ww`): the instance that already skips the tensor-map prefetch lets its TMA / MMA warps skip the cluster rendezvous (the epilogue warps alone wait for the pair) and start the ring fill right after the CTA-local setup: +2.7 / +2.1 % (B200 / B300, two processes) and +2.0 / +1.7 % under the pointer ABI on a third node pair. - **128-wide pair tile for the TP8 T=512 prefill tail** (`tail_gemm:tp8e0f0s9n128`): 56 column tiles x 2 row-tile pairs = 224 CTAs (1.5 waves of whole tiles instead of 112 CTAs on 148 SMs) with a 9-deep ring: +8.1..+8.9 % in two processes per GPU; T=256 and T >= 1024 keep the 256-wide tile (they lose 2-12 % with the narrow tile), so the rule is per row. Host planner (`cake_backend.py`): `land_alias_auto(n_pad, k_units)`, the `wait_warps` plan field, `tail_gemm_config(M, tp) -> (num_stages, block_n)` and the `n<block_n>` kernel-key component (`tail_gemm:tp8e0f0s9n128`); `_tail_workspace` is sized by the tile width; the late-trigger norm instance leaves the production route set. **Follow-up commit `891f69adb`** (on top of the delivery commit `8ce7ef04e`): the module-level annotation of the `_TAIL_WS` cache in `cake_backend.py` is typed as the 4-tuple key `(device, sk_tiles, max_seg, block_n)` that the round-7 `_tail_workspace` uses; the upstream `pre-commit` mypy hook rejected the 3-tuple annotation on the delivery head. Annotation only: no runtime change, no generated program touched; the receipts and the pytest runs below are bound to the programs of `8ce7ef04e`, which this commit leaves byte-identical. ## Evidence (Cake export protocol, producer `e8937a5205a`, target `b4a2da646a3`) Every row runs the source (Cake production launcher) and the exported programs in counterbalanced CUDA-graph groups; correctness = both arms against the FP32/BF16 torch reference, bitwise source == export. | arch | GPU | rows sealed | clock-unqualified (disclosed) | correct (bitwise source == export) | source/export | steps | |---|---|---|---|---|---|---| | sm_100a | B200 | 50 / 60 | 9 clock-unqualified + `tail_tp8_m16` (source / export ratio 0.930 below the 0.94 gate on the retry; correct, bitwise) | 50 / 50 | 0.9734 - 1.0755 | `5a00b5fa1ca4b7fc9b733cf7` (protocol) + `0aca3d85a1ce0e87bde7d441` (export) + `25bac3a89b03cca512a70471` (retry round + FlashInfer tests) | | sm_103a | B300 | 50 / 60 | 9 clock-unqualified + `front_tp1_m1024` (endpoint-drift gate on both measurements; correct, bitwise, ratio 0.9992) | 50 / 50 | 0.9646 - 1.0639 | `c9996135d315f7e29e22e3a3` (export) + `ed57976db5185227697f584d` (retry round + FlashInfer tests) | Clock-unqualified rows: the sustained dense tcgen05 GEMM rows run at the power cap and the interleaved timing session sampled the loaded SM clock below 0.85 x max on both bounded rounds (retry sampled / max MHz): sm_100a `front_tp1_m2048` 1620, `front_tp1_m4096` 1432, `front_tp1_m8192` 1350, `front_tp1_m16384` 1305, `front_tp8_m8192` 1507, `front_tp8_m16384` 1387, `tail_tp1_m4096` 1522, `tail_tp1_m8192` 1425, `tail_tp1_m16384` 1350 (/ 1965); sm_103a `front_tp1_m2048` 1590, `front_tp1_m4096` 1395, `front_tp1_m8192` 1260, `front_tp1_m16384` 1245, `front_tp8_m8192` 1417, `front_tp8_m16384` 1282, `tail_tp1_m4096` 1455, `tail_tp1_m8192` 1312, `tail_tp1_m16384` 1252 (/ 2032) -- the same set the previous round disclosed; these rows are correct and above 1.00 in the source-kernel contract and are re-measured with the exported programs in the speedup table below. The `tail_tp8_m16` ratio on B200: the Cake source arm launches that instance with device-memory tensor-map descriptors (the per-instance pointer ABI of the producer tree), while this package binds every program's descriptors by value; on that node the by-value form is 7 % slower, so the source runs faster than the export. The exported program is correct and bitwise; the row is disclosed rather than the gate loosened. The delivery (registries + generated sources) is byte-identical from the B200 and the B300 export (sha256 `b72137682d2d29f8`). <details><summary>sm_100a per-row receipts</summary> | row (B200, sm_100a) | route | source | export | source / export | correctness (both arms) | bitwise src == exp | verdict | |---|---|---|---|---:|---|---|---| | front_tp1_m1 | decode | 50.27 us | 49.98 us | 1.0058 | pass | yes | sealed | | front_tp1_m2 | decode | 50.56 us | 50.62 us | 0.9987 | pass | yes | sealed | | front_tp1_m4 | decode | 50.75 us | 50.24 us | 1.0102 | pass | yes | sealed | | front_tp1_m8 | decode | 50.34 us | 50.40 us | 0.9987 | pass | yes | sealed | | front_tp1_m16 | decode | 50.62 us | 50.50 us | 1.0025 | pass | yes | sealed | | front_tp1_m32 | decode | 51.01 us | 51.04 us | 0.9994 | pass | yes | sealed | | front_tp1_m64 | decode | 53.63 us | 53.70 us | 0.9988 | pass | yes | sealed | | front_tp1_m128 | decode | 59.07 us | 59.30 us | 0.9962 | pass | yes | sealed | | front_tp1_m256 | prefill | 60.77 us | 60.90 us | 0.9979 | pass | yes | sealed | | front_tp1_m512 | prefill | 86.78 us | 86.46 us | 1.0037 | pass | yes | sealed | | front_tp1_m1024 | prefill | 166.05 us | 166.02 us | 1.0002 | pass | yes | sealed | | front_tp1_m2048 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1545/1965 MHz first, 1620/1965 MHz retry) | | front_tp1_m4096 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1380/1965 MHz first, 1432/1965 MHz retry) | | front_tp1_m8192 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1387/1965 MHz first, 1350/1965 MHz retry) | | front_tp1_m16384 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1305/1965 MHz) | | front_tp8_m1 | decode | 21.54 us | 21.66 us | 0.9941 | pass | yes | sealed | | front_tp8_m2 | decode | 21.66 us | 21.76 us | 0.9956 | pass | yes | sealed | | front_tp8_m4 | decode | 21.54 us | 21.66 us | 0.9941 | pass | yes | sealed | | front_tp8_m8 | decode | 21.73 us | 21.86 us | 0.9941 | pass | yes | sealed | | front_tp8_m16 | decode | 22.02 us | 22.02 us | 1.0000 | pass | yes | sealed | | front_tp8_m32 | decode | 23.20 us | 23.39 us | 0.9918 | pass | yes | sealed | | front_tp8_m64 | decode | 25.60 us | 25.66 us | 0.9975 | pass | yes | sealed | | front_tp8_m128 | decode | 31.20 us | 31.17 us | 1.0011 | pass | yes | sealed | | front_tp8_m256 | prefill | 44.42 us | 44.54 us | 0.9971 | pass | yes | sealed | | front_tp8_m512 | prefill | 46.88 us | 47.01 us | 0.9973 | pass | yes | sealed | | front_tp8_m1024 | prefill | 82.14 us | 81.44 us | 1.0086 | pass | yes | sealed | | front_tp8_m2048 | prefill | 118.02 us | 118.08 us | 0.9995 | pass | yes | sealed | | front_tp8_m4096 | prefill | 243.78 us | 243.81 us | 0.9999 | pass | yes | sealed | | front_tp8_m8192 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1477/1965 MHz first, 1507/1965 MHz retry) | | front_tp8_m16384 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1387/1965 MHz) | | tail_tp1_m1 | decode | 31.52 us | 31.42 us | 1.0030 | pass | yes | sealed | | tail_tp1_m2 | decode | 31.52 us | 31.46 us | 1.0020 | pass | yes | sealed | | tail_tp1_m4 | decode | 31.62 us | 31.58 us | 1.0010 | pass | yes | sealed | | tail_tp1_m8 | decode | 31.81 us | 31.84 us | 0.9990 | pass | yes | sealed | | tail_tp1_m16 | decode | 33.79 us | 33.98 us | 0.9943 | pass | yes | sealed | | tail_tp1_m32 | decode | 34.27 us | 34.43 us | 0.9954 | pass | yes | sealed | | tail_tp1_m64 | decode | 36.32 us | 36.42 us | 0.9974 | pass | yes | sealed | | tail_tp1_m128 | decode | 42.85 us | 39.84 us | 1.0755 | pass | yes | sealed | | tail_tp1_m256 | prefill | 46.53 us | 46.56 us | 0.9993 | pass | yes | sealed | | tail_tp1_m512 | prefill | 56.93 us | 56.93 us | 1.0000 | pass | yes | sealed | | tail_tp1_m1024 | prefill | 107.55 us | 108.10 us | 0.9950 | pass | yes | sealed | | tail_tp1_m2048 | prefill | 203.39 us | 203.10 us | 1.0014 | pass | yes | sealed | | tail_tp1_m4096 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1417/1965 MHz first, 1522/1965 MHz retry) | | tail_tp1_m8192 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1432/1965 MHz first, 1425/1965 MHz retry) | | tail_tp1_m16384 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1350/1965 MHz) | | tail_tp8_m1 | decode | 7.87 us | 7.90 us | 0.9958 | pass | yes | sealed | | tail_tp8_m2 | decode | 7.97 us | 8.00 us | 0.9960 | pass | yes | sealed | | tail_tp8_m4 | decode | 7.97 us | 8.03 us | 0.9920 | pass | yes | sealed | | tail_tp8_m8 | decode | 8.06 us | 8.13 us | 0.9921 | pass | yes | sealed | | tail_tp8_m16 | decode | 9.73 us | 10.46 us | 0.9297 | pass | yes | disclosed: source / export ratio 0.9297 below the 0.94 gate (correct + bitwise; the Cake source arm launches this instance with the pointer tensor-map ABI, the package binds by value) | | tail_tp8_m32 | decode | 10.94 us | 11.14 us | 0.9828 | pass | yes | sealed | | tail_tp8_m64 | decode | 11.71 us | 12.03 us | 0.9734 | pass | yes | sealed | | tail_tp8_m128 | decode | 14.50 us | 14.59 us | 0.9936 | pass | yes | sealed | | tail_tp8_m256 | prefill | 15.65 us | 15.65 us | 1.0000 | pass | yes | sealed | | tail_tp8_m512 | prefill | 18.53 us | 18.62 us | 0.9948 | pass | yes | sealed | | tail_tp8_m1024 | prefill | 26.53 us | 26.56 us | 0.9988 | pass | yes | sealed | | tail_tp8_m2048 | prefill | 40.03 us | 40.22 us | 0.9952 | pass | yes | sealed | | tail_tp8_m4096 | prefill | 65.73 us | 66.14 us | 0.9937 | pass | yes | sealed | | tail_tp8_m8192 | prefill | 116.67 us | 116.70 us | 0.9997 | pass | yes | sealed | | tail_tp8_m16384 | prefill | 227.63 us | 227.84 us | 0.9991 | pass | yes | sealed | </details> <details><summary>sm_103a per-row receipts</summary> | row (B300, sm_103a) | route | source | export | source / export | correctness (both arms) | bitwise src == exp | verdict | |---|---|---|---|---:|---|---|---| | front_tp1_m1 | decode | 50.24 us | 50.08 us | 1.0032 | pass | yes | sealed | | front_tp1_m2 | decode | 50.37 us | 50.46 us | 0.9981 | pass | yes | sealed | | front_tp1_m4 | decode | 50.40 us | 50.46 us | 0.9987 | pass | yes | sealed | | front_tp1_m8 | decode | 50.59 us | 50.40 us | 1.0038 | pass | yes | sealed | | front_tp1_m16 | decode | 50.82 us | 50.62 us | 1.0038 | pass | yes | sealed | | front_tp1_m32 | decode | 51.30 us | 51.36 us | 0.9988 | pass | yes | sealed | | front_tp1_m64 | decode | 54.27 us | 54.34 us | 0.9988 | pass | yes | sealed | | front_tp1_m128 | decode | 59.26 us | 59.33 us | 0.9989 | pass | yes | sealed | | front_tp1_m256 | prefill | 60.61 us | 60.80 us | 0.9969 | pass | yes | sealed | | front_tp1_m512 | prefill | 83.94 us | 84.51 us | 0.9932 | pass | yes | sealed | | front_tp1_m1024 | prefill | 154.40 us | 154.53 us | 0.9992 | pass | yes | disclosed: endpoint-drift gate (correct + bitwise, ratio 0.9992) | | front_tp1_m2048 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1605/2032 MHz first, 1590/2032 MHz retry) | | front_tp1_m4096 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1297/2032 MHz first, 1395/2032 MHz retry) | | front_tp1_m8192 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1275/2032 MHz first, 1260/2032 MHz retry) | | front_tp1_m16384 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1252/2032 MHz first, 1245/2032 MHz retry) | | front_tp8_m1 | decode | 21.70 us | 21.66 us | 1.0015 | pass | yes | sealed | | front_tp8_m2 | decode | 21.73 us | 21.79 us | 0.9971 | pass | yes | sealed | | front_tp8_m4 | decode | 21.73 us | 21.70 us | 1.0014 | pass | yes | sealed | | front_tp8_m8 | decode | 21.89 us | 21.86 us | 1.0015 | pass | yes | sealed | | front_tp8_m16 | decode | 22.14 us | 22.18 us | 0.9986 | pass | yes | sealed | | front_tp8_m32 | decode | 23.17 us | 23.30 us | 0.9945 | pass | yes | sealed | | front_tp8_m64 | decode | 25.73 us | 25.86 us | 0.9950 | pass | yes | sealed | | front_tp8_m128 | decode | 31.36 us | 31.39 us | 0.9989 | pass | yes | sealed | | front_tp8_m256 | prefill | 42.24 us | 42.15 us | 1.0023 | pass | yes | sealed | | front_tp8_m512 | prefill | 44.90 us | 44.74 us | 1.0036 | pass | yes | sealed | | front_tp8_m1024 | prefill | 77.12 us | 77.19 us | 0.9992 | pass | yes | sealed | | front_tp8_m2048 | prefill | 111.78 us | 111.90 us | 0.9989 | pass | yes | sealed | | front_tp8_m4096 | prefill | 231.86 us | 231.97 us | 0.9995 | pass | yes | sealed | | front_tp8_m8192 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1612/2032 MHz first, 1417/2032 MHz retry) | | front_tp8_m16384 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1282/2032 MHz) | | tail_tp1_m1 | decode | 31.97 us | 31.78 us | 1.0060 | pass | yes | sealed | | tail_tp1_m2 | decode | 31.97 us | 31.71 us | 1.0081 | pass | yes | sealed | | tail_tp1_m4 | decode | 32.16 us | 31.87 us | 1.0090 | pass | yes | sealed | | tail_tp1_m8 | decode | 32.19 us | 32.22 us | 0.9990 | pass | yes | sealed | | tail_tp1_m16 | decode | 33.82 us | 33.82 us | 1.0000 | pass | yes | sealed | | tail_tp1_m32 | decode | 34.59 us | 34.56 us | 1.0009 | pass | yes | sealed | | tail_tp1_m64 | decode | 36.67 us | 36.67 us | 1.0000 | pass | yes | sealed | | tail_tp1_m128 | decode | 42.62 us | 40.06 us | 1.0639 | pass | yes | sealed | | tail_tp1_m256 | prefill | 46.30 us | 46.11 us | 1.0042 | pass | yes | sealed | | tail_tp1_m512 | prefill | 54.66 us | 54.69 us | 0.9994 | pass | yes | sealed | | tail_tp1_m1024 | prefill | 104.29 us | 104.61 us | 0.9969 | pass | yes | sealed | | tail_tp1_m2048 | prefill | 191.92 us | 192.74 us | 0.9958 | pass | yes | sealed | | tail_tp1_m4096 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1440/2032 MHz first, 1455/2032 MHz retry) | | tail_tp1_m8192 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1447/2032 MHz first, 1312/2032 MHz retry) | | tail_tp1_m16384 | prefill | — | — | — | — | — | disclosed: clock-unqualified (1260/2032 MHz first, 1252/2032 MHz retry) | | tail_tp8_m1 | decode | 8.03 us | 8.03 us | 0.9999 | pass | yes | sealed | | tail_tp8_m2 | decode | 8.06 us | 8.03 us | 1.0040 | pass | yes | sealed | | tail_tp8_m4 | decode | 8.06 us | 8.10 us | 0.9960 | pass | yes | sealed | | tail_tp8_m8 | decode | 8.13 us | 8.13 us | 0.9999 | pass | yes | sealed | | tail_tp8_m16 | decode | 9.60 us | 9.95 us | 0.9646 | pass | yes | sealed | | tail_tp8_m32 | decode | 11.11 us | 11.20 us | 0.9915 | pass | yes | sealed | | tail_tp8_m64 | decode | 12.06 us | 12.16 us | 0.9921 | pass | yes | sealed | | tail_tp8_m128 | decode | 14.43 us | 14.56 us | 0.9912 | pass | yes | sealed | | tail_tp8_m256 | prefill | 15.36 us | 15.23 us | 1.0084 | pass | yes | sealed | | tail_tp8_m512 | prefill | 17.63 us | 17.79 us | 0.9910 | pass | yes | sealed | | tail_tp8_m1024 | prefill | 24.90 us | 25.02 us | 0.9949 | pass | yes | sealed | | tail_tp8_m2048 | prefill | 38.08 us | 38.59 us | 0.9867 | pass | yes | sealed | | tail_tp8_m4096 | prefill | 63.17 us | 63.30 us | 0.9980 | pass | yes | sealed | | tail_tp8_m8192 | prefill | 112.23 us | 112.19 us | 1.0003 | pass | yes | sealed | | tail_tp8_m16384 | prefill | 212.99 us | 213.19 us | 0.9991 | pass | yes | sealed | </details> ### Measured speedups (Cake contract, paired cold-L2 CUPTI graph timing vs the fastest stock chain of each row) Kernel arm = the exported programs of this PR (CUDA-graph form, three alternating groups, cold-L2 CUPTI timing, correctness of both arms against the torch reference, graph == eager); denominator = the fastest stock chain of the row (see Baselines). **B200: 60 / 60 rows > 1.00, min 1.028 (`tail_tp1_m16384`), geomean 1.582. B300: 60 / 60 rows > 1.00, min 1.020 (`tail_tp1_m2048`), geomean 1.589.** The seven tail TP8 rows below 14.5 us were measured in two independent processes per GPU (identical medians). | row | B200: exported program / fastest stock chain us = speedup | B300: same | |---|---|---| | `front_tp1_m1` | 49.98 / 86.24 = 1.725 | 50.21 / 87.07 = 1.734 | | `front_tp1_m2` | 50.08 / 86.46 = 1.726 | 50.37 / 87.27 = 1.732 | | `front_tp1_m4` | 50.69 / 87.74 = 1.731 | 50.82 / 88.42 = 1.740 | | `front_tp1_m8` | 50.08 / 86.40 = 1.725 | 50.24 / 86.85 = 1.729 | | `front_tp1_m16` | 50.43 / 88.67 = 1.758 | 50.59 / 88.99 = 1.759 | | `front_tp1_m32` | 51.23 / 90.62 = 1.769 | 51.39 / 90.02 = 1.752 | | `front_tp1_m64` | 54.37 / 91.01 = 1.674 | 54.46 / 90.47 = 1.661 | | `front_tp1_m128` | 58.75 / 94.88 = 1.615 | 58.98 / 94.56 = 1.603 | | `front_tp1_m256` | 61.34 / 116.54 = 1.900 | 60.93 / 115.30 = 1.892 | | `front_tp1_m512` | 86.94 / 164.57 = 1.893 | 84.77 / 158.62 = 1.871 | | `front_tp1_m1024` | 164.51 / 289.02 = 1.757 | 152.64 / 266.37 = 1.745 | | `front_tp1_m2048` | 330.98 / 595.42 = 1.799 | 307.97 / 559.81 = 1.818 | | `front_tp1_m4096` | 671.49 / 1201.13 = 1.789 | 626.79 / 1183.02 = 1.887 | | `front_tp1_m8192` | 1330.02 / 2371.39 = 1.783 | 1243.51 / 2320.69 = 1.866 | | `front_tp1_m16384` | 2592.51 / 4732.41 = 1.825 | 2443.17 / 4592.68 = 1.880 | | `front_tp8_m1` | 21.79 / 57.50 = 2.639 | 21.86 / 57.92 = 2.650 | | `front_tp8_m2` | 21.89 / 57.15 = 2.611 | 21.86 / 57.79 = 2.644 | | `front_tp8_m4` | 21.98 / 56.54 = 2.572 | 21.89 / 57.22 = 2.614 | | `front_tp8_m8` | 21.92 / 56.99 = 2.600 | 22.02 / 57.73 = 2.622 | | `front_tp8_m16` | 22.37 / 57.47 = 2.569 | 22.30 / 57.73 = 2.588 | | `front_tp8_m32` | 23.39 / 59.68 = 2.551 | 23.52 / 59.23 = 2.518 | | `front_tp8_m64` | 25.82 / 60.80 = 2.354 | 25.98 / 60.71 = 2.336 | | `front_tp8_m128` | 31.45 / 63.36 = 2.014 | 31.81 / 62.53 = 1.966 | | `front_tp8_m256` | 44.70 / 72.06 = 1.612 | 42.30 / 70.91 = 1.676 | | `front_tp8_m512` | 47.30 / 78.43 = 1.658 | 45.12 / 76.19 = 1.689 | | `front_tp8_m1024` | 81.06 / 106.51 = 1.314 | 77.31 / 102.75 = 1.329 | | `front_tp8_m2048` | 117.98 / 170.88 = 1.448 | 111.07 / 162.59 = 1.464 | | `front_tp8_m4096` | 250.53 / 321.38 = 1.283 | 237.83 / 305.91 = 1.286 | | `front_tp8_m8192` | 489.94 / 653.09 = 1.333 | 466.72 / 642.55 = 1.377 | | `front_tp8_m16384` | 938.78 / 1316.96 = 1.403 | 877.40 / 1270.92 = 1.448 | | `tail_tp1_m1` | 31.36 / 39.20 = 1.250 | 31.62 / 39.62 = 1.253 | | `tail_tp1_m2` | 31.36 / 40.32 = 1.286 | 31.65 / 40.35 = 1.275 | | `tail_tp1_m4` | 31.58 / 40.00 = 1.266 | 31.78 / 40.26 = 1.267 | | `tail_tp1_m8` | 31.68 / 39.42 = 1.244 | 31.90 / 39.75 = 1.246 | | `tail_tp1_m16` | 33.98 / 40.10 = 1.180 | 33.92 / 40.35 = 1.190 | | `tail_tp1_m32` | 34.46 / 39.94 = 1.159 | 34.59 / 40.45 = 1.169 | | `tail_tp1_m64` | 36.48 / 40.13 = 1.100 | 36.70 / 40.80 = 1.112 | | `tail_tp1_m128` | 39.87 / 44.54 = 1.117 | 40.10 / 44.96 = 1.121 | | `tail_tp1_m256` | 46.50 / 48.80 = 1.050 | 46.18 / 48.93 = 1.060 | | `tail_tp1_m512` | 56.80 / 72.03 = 1.268 | 54.43 / 69.63 = 1.279 | | `tail_tp1_m1024` | 107.90 / 121.28 = 1.124 | 103.71 / 115.71 = 1.116 | | `tail_tp1_m2048` | 207.97 / 215.87 = 1.038 | 197.73 / 201.63 = 1.020 | | `tail_tp1_m4096` | 410.93 / 423.23 = 1.030 | 398.32 / 415.99 = 1.044 | | `tail_tp1_m8192` | 833.68 / 878.27 = 1.054 | 812.20 / 867.82 = 1.069 | | `tail_tp1_m16384` | 1715.30 / 1762.59 = 1.028 | 1658.83 / 1734.24 = 1.046 | | `tail_tp8_m1` | 7.97 / 16.10 = 2.020 | 8.06 / 16.06 = 1.992 | | `tail_tp8_m2` | 7.94 / 17.12 = 2.157 | 8.06 / 17.02 = 2.111 | | `tail_tp8_m4` | 7.97 / 17.22 = 2.160 | 8.06 / 17.34 = 2.151 | | `tail_tp8_m8` | 8.13 / 16.80 = 2.067 | 8.16 / 16.51 = 2.023 | | `tail_tp8_m16` | 10.08 / 18.24 = 1.810 | 10.37 / 17.70 = 1.707 | | `tail_tp8_m32` | 10.53 / 17.73 = 1.684 | 10.94 / 17.66 = 1.614 | | `tail_tp8_m64` | 12.10 / 19.42 = 1.606 | 12.32 / 19.14 = 1.553 | | `tail_tp8_m128` | 14.72 / 19.46 = 1.322 | 14.53 / 19.33 = 1.330 | | `tail_tp8_m256` | 15.78 / 21.86 = 1.385 | 15.23 / 21.66 = 1.422 | | `tail_tp8_m512` | 18.72 / 26.78 = 1.431 | 17.66 / 26.11 = 1.478 | | `tail_tp8_m1024` | 26.53 / 36.54 = 1.378 | 24.99 / 35.30 = 1.412 | | `tail_tp8_m2048` | 39.90 / 50.27 = 1.260 | 38.11 / 48.80 = 1.280 | | `tail_tp8_m4096` | 65.86 / 95.01 = 1.443 | 63.36 / 92.10 = 1.454 | | `tail_tp8_m8192` | 120.19 / 183.94 = 1.530 | 111.71 / 174.91 = 1.566 | | `tail_tp8_m16384` | 244.26 / 356.29 = 1.459 | 226.59 / 339.46 = 1.498 | ### Baselines and their source PRs - **Stock chain (Cake contract denominator, per row the fastest of):** torch / cuBLAS `torch.mm` / `torch.addmm` (FP32-upcast or `out_dtype=float32` router GEMM, BF16 latent / shared GEMMs, `addmm` for the tail sum) and every FlashInfer `mm_bf16` backend that computes the product (`auto`, `cudnn`, `cublaslt`, `tgv`, `tinygemm`, `cutlass`, `cute-dsl`), `flashinfer.rmsnorm` for the tail norm, torch SiTU; eager and CUDA-graph twins, the candidate timed in the winning form. FlashInfer `mm_bf16`: `flashinfer/gemm/gemm_base.py`, origin #2070 (2062dec, 2026-01-09), last touched by #5593 (97b3bd8, 2026-09-27). FlashInfer `rmsnorm`: `flashinfer/norm/__init__.py`, cute-dsl port #2428 (8bf921a, 2026-03-12), last touched by #4753 (5a8e62a, 2026-09-01). Torch: the `nvcr.io/nvidia/pytorch:26.01-py3` container of the measurement steps. - **Previous round of this operator (the programs this PR replaces):** #5584 (1eb503c, 2026-09-27); the exported-program contract of that round is the per-row comparison baseline below. - **Operator semantics reference:** `nvidia/Kimi-K3-NVFP4` `modeling_kimi_linear.py` (`KimiSparseMoeBlock`, `KimiMLP`, `KimiRMSNorm`, `SituAndMul`, `KimiMoEGate`); serving layout from vLLM `models/kimi_k3/nvidia/{model.py, latent_moe_runner.py, low_latency_gemm.py}` and SGLang `srt/models/kimi_k3.py`. - **Test oracle:** the torch reference in `tests/experimental/test_cake_kimi_k3_latent_moe.py` (transcribed from the Cake reference module). - **Merge-base:** upstream `main` 3142c80 (#5627, 2026-09-27); no upstream change to the package or the baselines between the merge-base and this branch. ## Test ``` pytest tests/experimental/test_cake_kimi_k3_latent_moe.py ``` `pytest tests/experimental/test_cake_kimi_k3_latent_moe.py`: **39 passed**, 24 warnings on B200 (sm_100a) and on B300 (sm_103a), run on the delivered tree in the export clone of each GPU (plan / route tests + GPU correctness for both stages x TP {1, 8} x the row set, bit-identical re-launch and CUDA-graph replay, `y` byte-exact). Producer-side gates on the same kernel tree: e2e GPU slice, unit tests, compute-sanitizer synccheck + memcheck 0 errors, bench-regression +6.5 % / +7.3 % over the floor, 62 / 62 rows correct in the paired contract and in four launch-sequence stress passes per GPU, 60 / 60 rows faster than the stock chain on each GPU (min 1.024). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Performance and Reliability** * Improved Kimi K3 latent MoE execution across supported GPU architectures, including more efficient decode and prefill planning and better handling of output boundaries. * Updated normalization and output calculations for more efficient processing. * **Bug Fixes** * Corrected cases where output stores could extend beyond the valid hidden-width range. <!-- 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>
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 (csrcand the registry).What r54 changes against the r52 programs (
698a96197):p_fullmbarrier is removed from the wide two-CTA kernel. Since r52 the softmax / correction warps arriveo_ready(tile parity, 4 warps x 2 CTAs, cluster arrive to CTA 0) right after the P store'sfence.proxy.asyncat the same site where they arrivedp_full, and the MMA warp waited both barriers at the same point before PV(t) and before the drain;p_fulltherefore carried no ordering of its own. PV(t) is now gated byo_readyalone: 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).Performance (same-session CUPTI cold-L2, CUDA-graph replay, 3 alternating rounds; Cake vs the fastest existing FlashInfer route)
Every row on both GPUs is faster than the best existing FlashInfer route (minimum 1.016x,
mirror_h96_b2on B300; the same thinnest row as in #5602). Columns: median of three alternating rounds;best FIis whichever of trtllm-gen / cute-dsl was faster on that row and GPU (trtllm-gen has no H96 MLA decode kernel, soautoresolves to cute-dsl there).Baselines and their source PRs
The
vs best FIcolumn compares against the two existing FlashInfer routes of the same public entry points, resolved at this branch's upstream merge-basef4d45edf1(2026-09-27, #5596):backend="trtllm-gen", also whatbackend="auto"selects for H12):trtllm_batch_decode_with_kv_cache_mla/trtllm_prefill_with_kv_cache_mlainflashinfer/mla/_core.pylaunching the prebuilt TRT-LLM-gen FMHA cubins pinned inflashinfer/artifacts.py(TRTLLM_GEN_FMHA = 2d6a5a02…/fmha/trtllm-gen/; pin set by fix: SageAttention support block size doesn't divide sequence; support K-smoothing #4654857bc3a11, 2026-08-26). Test oracle:tests/attention/test_trtllm_gen_mla.py.backend="cute-dsl", whatautoselects for H96): the CuTe-DSL MLA decode family underflashinfer/cute_dsl/attention/—mla_decode_fp8.py,mla_config.py,wrappers/batch_mla.py([CuTe DSL] Add modular FMHA prefill and MLA decode attention kernels #2805,00ac5053d, 2026-04-13; H64/H96 configs Support Kimi K2.5 H64 CuTe DSL MLA decode #32352d0e0efac2026-05-11; SM107 reland SM 107 Reland + Merge Back from v0.6.16 Release Branch #42805b1af97242026-07-30),mla_dispatch.py+monolithic/mla_decode.py(Restore monolithic CuTe-DSL MLA decode alongside modular, gated by cute_dsl_impl= #3296e91ac8f152026-05-21; LSE base semantics feat(mla): uniform LSE base semantics via return_lse_base_on_e #4650d7a7447cb2026-09-21), reached throughflashinfer/mla/_core.py(cute-dsl MLA op Add cute dsl mla decode op #274331b63bc3d2026-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 thecute-dslMLA kernels (git log 7b8f6129c..f4d45edf1 -- flashinfer/cute_dsl/attentionis empty), theTRTLLM_GEN_FMHApin is unchanged, andflashinfer/mla/_core.pychanged 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 perf(cake_mla): kimi_k3 shape, regenerate both architectures for the r52 kernel (deferred O rescale, wider long-request threshold) #5602).source/export= CAKE Source latency divided by CAKE Exported Kernel latency, per row and GPU):export receipts, 29 rows x 2 GPUs
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)
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_b1against atol 0.01 / rtol 0.02). The kernel change removes a barrier only; numerics are bitwise those of r52.🤖 Generated with Claude Code
Summary by CodeRabbit
References
The shared reference corpus for CAKE kernel development includes the following projects, documentation, and existing CAKE work:
These are corpus-level references. PR-specific implementation details, changes, and benchmark references are documented above.