Repository navigation
fix(cake_kimi_k3_mla): bound the lazy E4M3 reference drop on the row-tile kernels (Inf/NaN rows with attention-sink keys); regenerate SM100 / SM103 - #6126
Merged
yyihuang merged 1 commit intoOct 6, 2026
Conversation
…tile kernels (sink keys overflowed the FP32 rescale into Inf/NaN rows); regenerate both architectures Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
yyihuang
requested review from
Anerudhan,
Aneureka,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
dierksen,
feih-nv,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww and
yzh119
as code owners
October 6, 2026 12:38
Collaborator
Author
|
/bot run |
Collaborator
Author
|
@flashinfer-bot run tests/mla/test_cake_kimi_k3_mla.py |
Contributor
ℹ️ Recent review info⚙️ Run configuration
📒 Files selected for processing (25)
💤 Files with no reviewable changes (1)
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review. 📝 WalkthroughWalkthroughThe pull request updates Kimi K3 MLA program mappings and generated CUDA kernels. It changes paged-attention softmax state handling, adds architecture-specific paths to the wide kernel, updates launch bindings, and adds sink-key-row decode tests. ChangesKimi K3 MLA
Priority: ➖ Normal Estimated code review effort: 4 (Complex) | ~60 minutes Change: Bug fix Sequence Diagram(s)sequenceDiagram
participant KERNELS
participant Run
participant CudaKernel
KERNELS->>Run: select generated program
Run->>CudaKernel: launch with configured grid, block, stream, and shared memory
CudaKernel-->>Run: return kernel result
Suggested reviewers: Merge Risk: ⚪ Minimal · up to The sink-key fix has no identified merge-blocking issue. The reported GPU test results remain author-reported rather than independently verified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 4.55% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 22 functions across 22 files. (1 skipped: 1 unsupported.)
✨ 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
Oct 6, 2026
This was referenced Oct 6, 2026
Open
9 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.
Summary
Fixes an Inf/NaN defect in the
backend="cake"Kimi-K3 MLA FP8 paged attention decode route (#5552, #5602, #5871) and regenerates the programs of both architectures from the fixed producer. The host side (flashinfer/mla/cake_kimi_k3_mla.py, planner, grids, kernel arguments) is unchanged; this PR moves only the generated program sources, the jit registry and the test file.Defect. Serving Kimi-K3 with FP8 KV cache produced rows of
inf/nanoutput for some requests in a batch while the same requests alone, or with a different batch composition, were fine. Real Kimi-K3 KV caches contain an attention-sink key (token 0) whose score sits ~115-125 octaves above every other key for most heads. The swapped-AB row-tile kernels (RT 16/32/48/64/96, the H12 decode and MTP routes) keep one lazy E4M3 reference per softmax warpgroup and re-reference every column of the warpgroup to the exceeding column's own tile maximum. A column without a sink (a head whose query does not attend the sink, or a padding column that reads the next request's query rows inside the batch) exceeds on a later tile, the whole warpgroup moves its reference down by ~118 octaves, the FP32 rescalealpha = 2^118overflows the running sum and the O^T accumulator, and the merged row becomes Inf/NaN. This is why the failure depended on the batch and reproduced withnum_split=1.Fix (kernel only). Each softmax column keeps its running maximum in SMEM (the slot previously used for a dead column-sum write, so the SMEM carve-up and ring depths are unchanged) and the reference may never drop more than 32 octaves below the column's running maximum. Moves below 32 octaves are bit-identical to the previous programs; a sink column therefore keeps its reference near the sink and the non-sink columns of its warpgroup are only clamped where their mass is already below 2^-32 of the row. The wide two-CTA route (H96 decode, prefill) uses a monotone running maximum and is not affected; its programs are regenerated from the same producer and are byte-identical apart from the shared headers.
Tests
test_decode_sink_key_rowsintests/mla/test_cake_kimi_k3_mla.py: sink key (+2.0 on all 576 channels at token 0) with sink-attending query heads at +1.25 (~150 octaves above the random keys), H12 batch 1 and the 8-request 11-14-page band that reproduced the failure, H96 two-request and long-request cases, checked against an FP32 reference at FP8-class tolerance (atol/rtol 0.1) plus relative L2 <= 0.05. The previous programs fail these cases with Inf/NaN rows; the regenerated programs pass.tests/mla/test_cake_kimi_k3_mla.py: 38 passed (34 existing + the 4 new sink cases) on B200, 38 passed on B300 on the generated programs of this PR.tests/attention/test_trtllm_gen_mla.py -k trtllm_mla_blackwell(dispatch regression): 27 passed on B200; 27 passed on B300.prefill_h96_b1_q2048_kv131072on both GPUs is the exporter's loaded-SM-clock gate under the sustained power cap (1245 / 1965 MHz on B200, 1170 / 2032 MHz on B300, below the 0.70 fraction; the same row and gate as every previous head of this family), reproduced by a closed-prior retry round on B300 (1192 / 2032 MHz) and on B200 (1207 / 1965 MHz); on B200mirror_h96_b1_q1_kv76800(the shortest wide row, ~24 us) fails the paired ABBAdirectional_disagreementgate in both the first run and the closed-prior retry while bitwise parity holds (source/export 1.0053, then 0.9816 with the sign flipped), the arm-position offset of the paired harness on short rows already disclosed for this row in perf(cake_mla): kimi_k3 shape, regenerate both architectures for the r52 kernel (deferred O rescale, wider long-request threshold) #5602;mirror_h96_b2_q1_kv77699hit the same gate in the first run (1.0096) and passed in the retry. The B200 and B300 exporter runs delivered byte-identical files (the two delivered patches have the same SHA-256; 21 program sources + README + jit registry). The wide decode kernel's two architecture lowerings now fold into one source under__CUDA_ARCH__guards like every other kind, so the registry carries onemain_wideprogram instead of two.Performance
No host or schedule change; the bound costs one
max, one compare and one select per column per tile on the softmax warps and never fires on the benchmark rows. Export receipts above show source and generated programs within 0.2 % on every measured row.Unsupported (rejected explicitly)
Unchanged from #5552 / #5602.
Export receipts (29 rows per architecture, source vs generated: same-session ABBA CUPTI cold-L2 medians, CUDA-graph replay; B200 column from the closed-prior retry round, B300 column from the first run whose single refusal the retry reproduced)
🤖 Generated with Claude Code
Summary by CodeRabbit
Improvements
Tests
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.