Repository navigation
perf(cake_mxfp8_megamoe_ep16): sm_107a CuTe-DSL fused programs read the peer tables from shared memory - #8971
Merged
yyihuang merged 1 commit intoNov 6, 2026
Conversation
…he peer tables from shared memory The sm_107a fused programs (CTA-0 and return-all-CTA) read five by-value peer-pointer tables at a runtime rank. CuTe-DSL lowers each such read to a 31-deep select chain over the table slots, which ptxas turns into long runs of predicated constant loads on the publication critical path. As on sm_103a (flashinfer-ai#8735), each CTA now copies the 16 slots of these tables into 640 bytes of shared memory during setup, before the existing CTA barrier, and the route push, row push, rank flags and FC2 route-term lookups read the slot from shared memory. Pointer values and outputs are unchanged (bitwise identical to the CUDA backend). The sm_100a and sm_103a programs, the top-k reduce programs and the CUDA backend are unchanged.
yyihuang
requested review from
Anerudhan,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
flashinfer-bot and
yzh119
as code owners
November 6, 2026 13:36
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_mxfp8_megamoe_ep16.py |
Contributor
|
Important Review skippedReview was skipped because the changed files are too large to review. 📏 Files whose diffs exceed the per-file size limit: 2
Retrying the review will skip these files again. If they are generated content such as fixtures or snapshots, exclude them with
 |
yzh119
approved these changes
Nov 6, 2026
yzh119
left a comment
yzh119
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. Worth a quick check that the smem allocation covers 160768 + 640 bytes, and that a barrier sits between the table fill and the first read.
yyihuang
enabled auto-merge (squash)
November 6, 2026 13:41
yyihuang
pushed a commit
to yyihuang/flashinfer
that referenced
this pull request
Nov 6, 2026
…ashinfer-ai#8972, mega_moe_v3) into cake-lm-head-nondet-train Resolve tests/test_helpers/cake_package.py: keep this branch's LM-head det/nondet PACKAGES rows and the dense GEMM det/nondet rows from flashinfer-ai#8926.
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.
cake_mxfp8_megamoe_ep16is the fused MoE kernel with BF16 activations and MXFP8 weights for 16-way expert parallelism (EP16). It has a CUDA backend and a CuTe-DSL backend. Each backend has three programs per architecture: the fused kernel, the fused kernel with the return-all-CTA config, and the top-k reduce. This PR changes only the two sm_107a CuTe-DSL fused programs (CTA-0 and return-all-CTA).Compared with the CuTe-DSL backend on current
main(the previous Cake version), EP16 on VR200 runs 1.010–1.013× faster in both launch modes.Table. NVIDIA Vera Rubin NVL72 (VR200, sm_107a, 212 SMs per GPU), 16 GPUs (4 nodes × 4 GPUs in one NVLink rack, IMEX), driver 620.62, CUDA 13.0, PyTorch 2.14.1+cu130, CuTe-DSL (nvidia-cutlass-dsl) 4.8.1. Latency is one eager EP16 forward with a cold (flushed) L2, timed with CUDA events. A round value is the maximum over ranks of the per-rank median of 40 iterations. Each table value is the median over 32 rounds: 16 rounds load the baseline first, 16 load this PR first, and sessions are recreated each round. The baseline is the CuTe-DSL backend on current
main, and Speedup = baseline / this PR. In "Pre-sleep" mode, a GPU sleep of 1e6 cycles precedes each launch. In "Back-to-back" mode, launches run with no sleep. T is tokens per rank.balancedrouting spreads routes evenly across experts, andhotset50sends half of the routes to a hot set of experts. p is an exact one-sided Mann-Whitney U test on the per-round results, with "this PR faster" as the alternative. With 32 vs 32 rounds, 5.46e-19 is the smallest possible value; it means this PR was faster in every round.The CUDA backend, the sm_100a and sm_103a CuTe-DSL programs, and the top-k reduce programs are unchanged; the family ships on sm_100a, sm_103a and sm_107a.