[yyihuang]

Stack 10/13 · base: main

Summary

This PR adds miles_ext_kernels.cake.moe_grouped_gemm on the shared cake runtime. Its entry point is grouped_mm(x, w, offs), the ragged BF16 MoE grouped GEMM with forward and backward. It has the semantics of torch._grouped_mm(x, w.transpose(-2, -1), offs=offs). The PR vendors two packages from FlashInfer c5db8d3:

• cake_moe_grouped_det_gemm (det): every output is reduced in a fixed order.
• cake_moe_grouped_gemm (fast): it differs only in the weight gradient with the 512-wide k tile.

In miles, --moe-grouped-gemm-backend cake selects the op and requires --train-backend torchtitan (radixark/miles#4685, the top of the miles stack). When --cake-moe-grouped-gemm-variant {det,fast} is unset, it resolves to fast on the miles CLI and to det under deterministic mode. A process that never configured the runtime runs det.

Changes from the previous version of this layer:

• The 32-bit addressing limits are checked at the forward call, before any launch. They can also be checked at model build.
• The GPU tests use relative-RMS budgets that fail at a 2 % error.
• The variant is read from the runtime at each call. The variant= argument, the module-level VARIANT and the op's own configure are removed.
• prebuild also builds the input gradient's queue probe.

Stack

This is layer 10 of 13, on top of layer 9 (dense_gemm). The stack was rewritten as [Cake 0/13] layout plus 13 op layers of three commits each (vendor / op / tests). Every PR targets main, so this diff contains layers 0 to 10. The previous version of this layer was reviewed in #10 ([Cake 10/11] moe_grouped_gemm: ragged MoE grouped GEMM, det and fast).

This layer's commits:

• 220e42b vendor: moe_grouped_gemm packages
• 2edabb2 moe_grouped_gemm: the ragged BF16 MoE grouped GEMM on the cake runtime
• 701c627 tests: moe_grouped_gemm

Changes

From the layer 9 head to the layer 10 head: 112 files, +131,847 / −0. Per commit: 220e42b 101 files +130,711; 2edabb2 5 files +278; 701c627 6 files +858.

Group	Files	Lines
Vendored cake/_vendor/cake_moe_grouped_det_gemm/	51	+52,377
Vendored cake/_vendor/cake_moe_grouped_gemm/	49	+78,210
cake/_vendor/VENDORED.json (ops.moe_grouped_gemm block)	1	+124
Op module and implementation: cake/moe_grouped_gemm.py, cake/_impl/moe_grouped_gemm/{__init__,ops}.py	3	+275
Tests: tests/cake/moe_grouped_gemm/	6	+858
Registry and README (Ops row, Prebuild row)	2	+3

There are no pyproject.toml or tooling changes.

Packages

Package	Upstream	Role	Libraries (sm_100a / sm_103a / sm_107a)	Minimum CUDA
cake_moe_grouped_det_gemm	flashinfer-ai/flashinfer@c5db8d3, flashinfer/experimental/cake_moe_grouped_det_gemm (51 files, Cake revision 0716c911b456a7fd3ea2e0f668310d8b2f9a0ba9)	det: every output reduced in a fixed order	62 = 20 / 22 / 20	12.8 / 12.9 / 13.4
cake_moe_grouped_gemm	same commit, flashinfer/experimental/cake/cake_moe_grouped_gemm (49 files, Cake revision c7f1937c65c2e7f108ba359a0e0c3476ac2b67f2)	fast: the same forward and input-gradient programs; the weight gradient with the 512-wide k tile adds the split tail of its last wave with unordered FP32 atomics	54 = 18 / 20 / 16	12.8 / 12.9 / 13.4

Both packages are copied byte for byte, one level deep under cake/_vendor/. Every blob and both Cake revisions are equal to the previous layout's. There are no allow entries, because the vendored modules read no environment variable and do not import flashinfer. There are also no local files and no KCP variant.

python tools/vendor_cake.py fi --checkout ../flashinfer \
    --commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op moe_grouped_gemm \
    --package flashinfer/experimental/cake_moe_grouped_det_gemm \
    --package flashinfer/experimental/cake/cake_moe_grouped_gemm

API

OP is name="moe_grouped_gemm", variants=("det", "fast") (both are CLI choices), selected_by=(("moe_grouped_gemm_backend", "cake"),), with both packages and no extra flags. It exports check_grouped_dims, grouped_mm and max_grouped_rows. The module is torch-free.

• grouped_mm(x, w, offs) computes Y[offs[e-1]:offs[e]] = X[offs[e-1]:offs[e]] @ W[e].T.
  • Inputs are BF16 x [sum_m, K] and w [E, N, K], both with unit last-axis stride. offs [E] holds the int32 cumulative row ends on x's device. It is read on the device, with no host sync.
  • It is differentiable in x and w and returns [sum_m, N] in BF16. x.grad is BF16, and w.grad has w's dtype.
  • Empty groups get an exactly zero weight gradient. Rows past offs[-1] are left alone.
  • Numerics: FP32 accumulation of exact BF16 products, rounded once. These are cuBLAS / CUTLASS numerics, not their bits.
• check_grouped_dims(n, k, *, num_groups=None, max_rows=None) raises ValueError unless N and K are multiples of 256. With num_groups, E * N * K must fit 32-bit element offsets. With max_rows, max_rows * max(N, K) must fit too.
• max_grouped_rows(n, k) returns (2**31 - 1) // max(N, K), which is 299,593 rows at N = 4096, K = 7168.
• OP.validate:
  • --moe-grouped-gemm-backend cake requires --train-backend torchtitan. megatron, fsdp and unset are rejected.
  • --cake-moe-grouped-gemm-variant without the cake backend is rejected.
• An unsupported device raises NotImplementedError (the old layout raised RuntimeError). A toolkit too old for the architecture raises RuntimeError (CUDA 13.4 for sm_107a).
• README Ops row: GPUs sm_100a, sm_103a, sm_107a; no prebuild parameters; tests in tests/cake/moe_grouped_gemm/.

32-bit addressing

• The kernels address each output with 32-bit element offsets. sum_m * N (output), sum_m * K (input gradient) and E * N * K (weight gradient) must stay at or below 2**31 - 1.
• grouped_mm checks all three at the forward call, before any launch. A call whose input gradient or weight gradient would overflow now fails at its forward, instead of in the backward after the forward has run.
• check_grouped_dims(n, k, num_groups=, max_rows=) lets the model build refuse a configuration, and max_grouped_rows gives the row bound.

Variant routing

Situation	Result
Process never configured	det
miles CLI --moe-grouped-gemm-backend cake, flag unset, not deterministic	fast, written back to cake_moe_grouped_gemm_variant (the miles CLI default stays fast)
deterministic_mode=True in the namespace, flag unset or det	det
Explicit fast with deterministic_mode (CLI or configure)	ValueError ("not bit-reproducible")
configure() with an unknown variant	ValueError
At a call	every forward reads config.variant; its backward runs the package of its forward
fast under torch.use_deterministic_algorithms(True)	runs det, logged once (warn_runs_det), because the fast weight gradient would raise torch's non-determinism alert

fast changes only the weight gradient, and only where it picks the 512-wide k tile. That requires K to be a multiple of 512, with at least 4096 rows per group on average. There, the split tail of its last wave is added with unordered FP32 atomics, instead of in a second, ordered reduction launch. This gives last-bit differences from run to run. The forward and the input gradient run det's programs and give det's bits.

Prebuild

cake.prebuild("moe_grouped_gemm", *, device=None, max_workers=None) takes no variant argument and no model parameter, and returns the library paths. It builds:

• every program either variant may launch on the device's architecture with BF16 weights: the forward, the input gradient, and the weight gradient with k tile 256 and 512, with every stage built for the architecture;
• on sm_100a / sm_103a, the input gradient's probe of the memory domains. The probe is compiled once per compute capability through tvm-ffi's inline build. Previously, the first probed call compiled it.

If the probe does not build, a warning is logged and the input gradient splits its work queues evenly, with the same bits. After a prebuild, no call compiles.

The README Prebuild row reads: moe_grouped_gemm | none: both variants' forward, input-gradient and weight-gradient programs (both k tiles, every stage) for any expert shape, and the input gradient's queue probe on sm_100a / sm_103a.

Memory

• These workspaces are not in the runtime pool, so the pool bound (default 0) does not affect them.
• Both packages keep small workspaces per device and CUDA stream: scheduler words, wave counters, and the FP32 tail sums of the fast weight gradient.
• Every launch leaves them at rest for the next launch on that stream. They stay outside the pool because pooled buffers do not keep their content between calls.
• Two streams never share them, and a launch captured into a CUDA graph takes its own.
• release() (through miles_ext_kernels.release()) drops both packages' per-stream workspaces and the queue splits probed on them. The README Memory section lists them as out-of-pool per-stream workspaces that grow to the largest call.

Review items

These items come from the review of the previous version: #10 (comment). The measurements are the reviewer's, taken on GB300.

#	Reviewer item	Resolution	Tests
10-1	The int32 limit showed up mid-step. At K = 7168, N = 4096, sum_m = 300000 passed check_grouped_dims and the forward, and then the backward raised; 299,593 worked. The reviewer asked for a config-time check.	grouped_mm now checks sum_m * N, sum_m * K and E * N * K at the forward call, before any launch. check_grouped_dims gains num_groups / max_rows, and max_grouped_rows(N, K) gives the bound (299,593 at N = 4096, K = 7168). A model build can check E * N * K by passing num_groups. Per-call rows vary with routing, so they are checked at the forward.	CPU: test_limits.py::test_the_row_bound_is_the_32_bit_span_of_the_wider_output, ::test_the_dims_check_takes_the_most_rows_of_the_configured_run, ::test_the_dims_check_takes_the_expert_count_of_the_weight_gradient, ::test_the_dims_check_wants_multiples_of_256, ::test_a_call_past_the_row_bound_fails_at_the_forward_before_it_reaches_the_device, ::test_a_call_past_the_weight_gradient_bound_fails_at_the_forward. GPU: test_moe_grouped_gemm_gpu.py::test_the_32_bit_row_limit_fails_at_the_forward_before_any_launch. This test refuses 299,594 rows (one past the bound) at the forward. At 299,593 rows, it compares the last 2,048 rows of the output and of the input gradient, and the whole weight gradient, with closed-form FP64 references within the budget
10-2	Loose gate: the kernel's rel-Fro was about 1.7e-3, yet y × 1.02 and y × 1.05 both passed assert_close(atol=rtol=1e-2).	Per-tensor relative-RMS budgets via assert_rel_rms, against a per-group FP32 reference: {"y": 3.4e-3, "dx": 3.4e-3, "dw": 3.4e-3}. That is twice the measured value of about 1.7e-3 on sm_103a for y at K = 7168, N = 4096. With detect_scale=1.02, a 2 % error in y, dx or dw fails.	GPU: test_moe_grouped_gemm_gpu.py::test_det_matches_the_reference_and_repeats_bit_for_bit[6 cases], ::test_fast_gives_the_det_bits_where_its_reductions_are_ordered[6], ::test_the_fast_weight_gradient_adds_its_split_tail_unordered_within_the_budget
cleared	No upstream short-group weight-gradient race. 1,200 forward + backward runs, det and fast, with short groups between large ones, gave 0 failures. det was bit-exact, and fast y / dx equaled det.	Kept as regression tests.	GPU: ::test_short_groups_between_large_ones_repeat_bit_for_bit[w13/w2]
cleared	Empty input (sum_m = 0) gave an exact-zero dw.	Kept.	GPU: ::test_empty_input_gives_an_exact_zero_weight_gradient[det/fast]
cleared	Rows past offs[-1] (torchtitan padding) left valid y / dx alone and dw correct, with zero and random tails.	Kept.	GPU: ::test_rows_past_the_last_group_end_leave_the_valid_results_alone[det/fast × zero/random]

Cross-stack items from the review of #1 that touch this layer:

• The unconfigured default is now det, under one runtime rule. The miles CLI still resolves to fast.
• The prebuild contract now includes the queue probe that the first call used to compile.
• The runtime owns the arguments. The op declares OP.validate, and its own arguments.py, configure and VARIANT are removed.
• Per-stream workspaces are never shared across streams. Tests: ::test_calls_alternating_between_the_default_and_a_side_stream_give_the_same_bits[det/fast], ::test_release_drops_the_stream_workspaces_and_later_calls_give_the_same_bits.

The "known doc mismatch" in the previous body no longer applies. It said that deterministic mode never applies under torchtitan and that the default is always fast. The runtime rule is now explicit:

• unconfigured: det;
• miles CLI: fast;
• deterministic_mode: det;
• torch deterministic algorithms: det, logged once.

Tests

All tests are in tests/cake/moe_grouped_gemm/. The GPU tests use three layer shapes, given as (hidden H, expert intermediate I, local experts E) with w13 [E, 2I, H] and w2 [E, H, I]: (2048, 768, 16), (2048, 1536, 8) and (4096, 2048, 9).

• CPU, torch-free: test_arguments.py (11).

• CPU, torch on meta tensors: test_limits.py (6). These tests are listed under review item 10-1.

• GPU, test_moe_grouped_gemm_gpu.py (17). Besides the tests listed under Review items, it covers:

  • new group boundaries on the same shapes, read on the device ([det/fast]);
  • the variant being read at each call, with each backward running its forward's variant;
  • torch deterministic algorithms running det and saying so once;
  • calls the programs do not serve being refused;
  • the device needing its architecture's toolkit;
  • det bit for bit in a new process;
  • a call for a device other than the current one giving its bits ([det/fast], multi_gpu);
  • the calls importing no flashinfer.

• GPU, test_prebuild_gpu.py (4): one prebuild, then 18 ROUTES with every variant, under no_build() and with no tvm-ffi inline build either. The routes are:

  • 6 ragged layer shapes: 3 local-expert layers × w13 / w2, with group sizes cycled from 300, 0, 1, 700, 4099, 33, 511, 257 and 2048;
  • one-expert (6000 rows on one of 8 experts) and no-rows;
  • 5 uneven routings: 1 / 255 / 3000 / 33000 / 140000 rows over 8 experts;
  • padded (300 rows past offs[-1]);
  • k512-split-tail (N=4864, K=1024, sizes 16400 / 0 / 16666);
  • wide-k-many-rows (2 × 16400 rows, 4096 × 4096) and narrow-k-many-rows (2 × 9000 rows, N=4096, K=2048);
  • side-stream.

  The other tests check four things: that no route compiles after the prebuild; that a variant configured after a call runs at the next call without building; that the prebuild builds every program the planner binds; and that the prebuild is idempotent.

Validation

CPU, at each commit

The image and commands are the same as for layer 8:

• the full pytest tests run in a torch-free venv and with the image Python;
• vendor_cake.py --check and pre-commit run --all-files, with CUDA hidden.

At the layer head, two more checks ran: vendor_cake.py --check moe_grouped_gemm --upstream and pytest -m 'not multi_gpu' tests/cake/moe_grouped_gemm.

The image is radixark/miles:dev-202611032236 (Python 3.12.16, torch 2.14.1+cu130, apache-tvm-ffi 0.1.11).

Commit	Torch-free venv	Image Python	Other checks	Verdict
220e42b (vendor)	486 passed, 123 skipped, 1 failed	1082 passed, 260 skipped, 1 failed	vendor --check rc 0; pre-commit rc 0, no diff	only test_vendored.py::test_every_vendored_op_is_registered fails, as the per-commit rule allows
2edabb2 (op)	492 passed, 125 skipped, 1 failed	1090 passed, 260 skipped, 1 failed	vendor --check rc 0; pre-commit rc 0, no diff	only test_op_contract.py::test_the_op_has_cpu_tests_and_a_prebuild_test[moe_grouped_gemm] fails, as allowed
701c627 (tests)	513 passed, 128 skipped, 0 failed (rc 0)	1120 passed, 262 skipped, 0 failed (rc 0)	vendor --check rc 0; pre-commit rc 0, no diff; --check moe_grouped_gemm --upstream rc 0; tests/cake/moe_grouped_gemm (image, CPU) 29 passed, 2 skipped	all green

vendor_cake.py --check --upstream at the top of the stack (layer 13, with the same ops.moe_grouped_gemm block and package trees) returned rc 0.

GPU

The GPU tests ran on the top of the stack (layer 13). Its moe_grouped_gemm files and its ops.moe_grouped_gemm block are byte-identical to this layer's. These are cake/moe_grouped_gemm.py, _impl/moe_grouped_gemm/, both _vendor/cake_moe_grouped_* trees and tests/cake/moe_grouped_gemm/.

The container was radixark/miles:dev-202611032236 (amd64 on B200, arm64 on GB300), with Python 3.12.16, torch 2.14.1+cu130, nvcc 13.0.88 (CUDA 13.0), apache-tvm-ffi 0.1.11 and pytest 9.2.0. Each op had its own JIT cache.

Commands:

• Cold, then warm: PYTHONPATH=. python -m pytest -rfEs -p no:cacheprovider tests/cake/moe_grouped_gemm -m "not multi_gpu".
• Multi-GPU: -m multi_gpu on 2 GPUs.
• Prebuild first: as in layer 8.

Site	GPU	Driver	Cold (empty cache)	Warm	Prebuild first, empty cache	multi_gpu (2 GPUs)
B200	NVIDIA B200, sm_100a, 148 SMs	580.82.07	66 passed, 0 skipped, 135.5 s (24 libraries)	66 passed, 0 skipped, 115.5 s	4 passed, 35.1 s (24 libraries)	2 passed
GB300	NVIDIA GB300, sm_103a, 152 SMs	580.159.03	66 passed, 0 skipped, 77.3 s (prebuild test first; 28 libraries)	66 passed, 0 skipped, 31.9 s	in the cold run: 4 passed	2 passed
R200	sm_107a	—	pending	pending	pending	pending

Gates on both B200 and GB300, with 0 failures and 0 skips:

• det bitwise. test_det_matches_the_reference_and_repeats_bit_for_bit (6 layer shapes, within the rel-RMS budget), test_det_repeats_bit_for_bit_in_a_new_process and test_short_groups_between_large_ones_repeat_bit_for_bit[w13/w2] passed.
• Default and side stream in turn. test_calls_alternating_between_the_default_and_a_side_stream_give_the_same_bits[det/fast] and test_release_drops_the_stream_workspaces_and_later_calls_give_the_same_bits passed.
• fast. test_fast_gives_the_det_bits_where_its_reductions_are_ordered (6), test_the_fast_weight_gradient_adds_its_split_tail_unordered_within_the_budget and test_torch_deterministic_algorithms_run_det_and_say_so_once passed.
• 32-bit bound. test_the_32_bit_row_limit_fails_at_the_forward_before_any_launch passed.
• Another device. test_a_call_for_a_device_other_than_the_current_one_gives_its_bits[det/fast] passed on 2 GPUs.
• Prebuild.
  • Starting from an empty cache, one prebuild served all 18 routes with both variants, under no_build() and without any inline tvm-ffi build.
  • A variant configured after a call ran without building.
  • Every planner-bound program was prebuilt.
  • A second prebuild was idempotent.
• Pool. The op does not use the workspace pool, so no pool gate applies.
• sm_107a (R200). Pending: no R200 run of this tree has been recorded yet.

Performance

The performance section does not fit in this description. It is posted in full in the three comments below: Performance (continued, 1 of 3), (2 of 3) and (3 of 3). Every number there is copied from FlashInfer pull request bodies; nothing was measured for this PR. The comments cover:

• det. perf(cake_moe_grouped_det_gemm): skip forward recompute under checkpoint early stop flashinfer-ai/flashinfer#8247 delivered exactly the vendored det package. It measures B200, GB300 and a 208-SM sm_107a part against two baselines: the previous det wrapper, and torch._grouped_mm with autograd. The rows are checkpointed single GEMMs and checkpointed expert blocks, with early stop on and off. Its body states: "Every early-stop-on row is faster than both baselines."
• fast on sm_100a and sm_103a. perf(cake_moe_grouped_gemm): run the split k512 wgrad tail whole on the least-loaded clusters on sm_100a and sm_103a flashinfer-ai/flashinfer#8269 delivered exactly the vendored fast package. It reports dW and checkpointed rows on B200 and GB300, against det Cake, the previous fast package and the fastest of torch._grouped_mm, a per-expert cuBLAS loop and cuDNN.
• fast on sm_107a. perf(cake_moe_grouped_gemm): run the split k512 wgrad tail whole on the least-loaded clusters on sm_107a flashinfer-ai/flashinfer#8258 measured the revision before perf(cake_moe_grouped_gemm): run the split k512 wgrad tail whole on the least-loaded clusters on sm_100a and sm_103a flashinfer-ai/flashinfer#8269, which has the same sm_107a programs. Its det column predates perf(cake_moe_grouped_det_gemm): skip forward recompute under checkpoint early stop flashinfer-ai/flashinfer#8247, so its ckpt ES ratios against det do not describe the vendored det package.
• Missing.
  • Forward-only and input-gradient-only timings.
  • A det / non-Cake SOTA ratio for single operations.
  • The fast package's expert block.
  • Any timing of the fast package on sm_107a at the vendored revision.
  • GB200, B300, 212-SM or 216-SM sm_107a measurements.
  • A training-step timing in miles.

Notes

• sm_107a: cake_moe_grouped_det_gemm ships 20 sm_107a libraries and cake_moe_grouped_gemm ships 16. Both need CUDA 13.4 to build them, and an older toolkit raises RuntimeError. The queue probe is built on sm_100a / sm_103a only. No GPU run on sm_107a has been recorded for this tree yet.
• Supported inputs: BF16, with N and K multiples of 256. The backend requires torchtitan.
• 32-bit element offsets bound both E * N * K and the rows per call (max_grouped_rows). Rows per call vary with routing and are checked at the forward.
• Under fast, the weight gradient with the 512-wide k tile differs in its last bits from run to run.
• The per-stream workspaces live outside the pool and grow to the largest call until release().

🤖 Generated with Claude Code

[Zhichenzzz]

Review of layer 10 (moe_grouped_gemm) — suggest approve with nits

Confirmed on GPU

• The int32 addressing limit only shows up mid-step. At K=7168, N=4096, sum_m=300000 passes check_grouped_dims and the forward, then the backward raises ValueError; 299,593 works. Please check this at config time.
• The test gate is loose. The kernel's rel-Fro error is about 1.7e-3, yet y×1.02 and y×1.05 both pass assert_close(atol=rtol=1e-2). A relative Frobenius gate would catch them.

Cleared on GPU

• No sign of the upstream short-group wgrad race. I ran 1,200 fwd+bwd passes, det and fast, with 64 or 32 groups whose sizes come from {0, 1, 7, 31, 63, 64, 65, …}, interleaved with large groups:
  • 0 failures;
  • det bit-exact across iterations;
  • fast y/dx bit-equal to det.
• Empty input. sum_m=0 gives an exact-zero dw.
• Rows past offs[-1] (torchtitan padding) leave the valid y/dx and every dw correct, for both zero and random tails.

---

Verified on two GB300 nodes (sm_103a), CUDA 13.0, torch 2.14.1+cu130, at the stack head 2e422c50ca (this layer's files are unchanged there). The full suite passes (35 files, 809 tests). Each item above was checked with a dedicated GPU script; numbers are measured. sm_100a (GB200) paths were not exercised.