refactor(cake_moe_grouped_gemm): det-only names, explicit deterministic=False, and deterministic weight-gradient accumulation by yyihuang · Pull Request #8924 · flashinfer-ai/flashinfer

yyihuang · 2026-11-06T05:51:46Z

Summary

This PR changes the Cake MoE grouped GEMM in FlashInfer host code only. Kernels and programs are unchanged.

Det-only names. The deterministic grouped GEMM can now be reached only through names that contain det. All old det-free names and compatibility aliases are removed with no shim. Old imports raise ModuleNotFoundError and old attributes raise AttributeError.
Explicit deterministic=False. The non-deterministic package and the flashinfer.experimental.cake.grouped_mm facade run only when deterministic=False is passed explicitly. PyTorch's non-determinism alert fires only right before an order-varying launch.
Deterministic weight-gradient accumulation. grouped_det_gemm_wgrad(..., out=main_grad, accumulate=True) and cake_grouped_det_mm(..., fuse_wgrad_accumulation=True) accumulate into main_grad. The result is bitwise equal to AccumulateGrad of the deterministic dW and bitwise identical from run to run.

In this description, the det package is flashinfer.experimental.cake_moe_grouped_det_gemm and the non-det package is flashinfer.experimental.cake.cake_moe_grouped_gemm.

Breaking change: det-only names

No alias, shim, module __getattr__ or compatibility entry is left anywhere.

Old (removed)	Error now	Det name to use
package `flashinfer.experimental.cake_moe_grouped_gemm` and submodules `cake_autograd`, `cake_backend`, `cake_build`, `cake_jit` (directory deleted)	`ModuleNotFoundError`	`flashinfer.experimental.cake_moe_grouped_det_gemm` (same submodules)
module `flashinfer.grouped_mm.cake` (renamed with `git mv`)	`ModuleNotFoundError`	`flashinfer.grouped_mm.cake_det`
`flashinfer.grouped_mm.{grouped_gemm_fwd, grouped_gemm_dgrad, grouped_gemm_wgrad, cake_grouped_mm}`	`AttributeError`	`flashinfer.grouped_mm.{grouped_det_gemm_fwd, grouped_det_gemm_dgrad, grouped_det_gemm_wgrad, cake_grouped_det_mm}`
det package root `grouped_gemm_{fwd,dgrad,wgrad}`, `prepare_grouped_gemm_{fwd,dgrad,wgrad}`, `GroupedGemmLaunch`, `CakeGroupedMm`, `cake_grouped_mm`	`AttributeError`	`grouped_det_gemm_{fwd,dgrad,wgrad}`, `prepare_grouped_det_gemm_{fwd,dgrad,wgrad}`, `GroupedDetGemmLaunch`, `CakeGroupedDetMm`, `cake_grouped_det_mm`
det `cake_backend.{GroupedGemmLaunch, grouped_mm_bf16_cake, prepare_grouped_gemm_*}`	`AttributeError`	`cake_backend.{GroupedDetGemmLaunch, grouped_det_mm_bf16_cake, prepare_grouped_det_gemm_*}`
det `cake_autograd.{CakeGroupedMm, cake_grouped_mm}`	`AttributeError`	`cake_autograd.{CakeGroupedDetMm, cake_grouped_det_mm}`
`_RENAMED` maps and alias `__getattr__` in det `__init__`, `cake_autograd`, `cake_backend`, `grouped_mm/cake(_det).py`, `grouped_mm/__init__.py` (and its `_Any` import)	deleted	-
`deterministic` parameter of every det entry (`cake_grouped_det_mm` in the det package and in `grouped_mm.cake_det`)	`TypeError`	omit it
`CakeGroupedDetMm.apply(x, w, offs, <mode>)` (stale 4th positional mode)	`TypeError` "takes no fourth argument"	`apply(x, w, offs)`; fuse: `apply(x, w, offs, None, True)`
test-helper key `PACKAGES["cake_moe_grouped_gemm"]` (MoE entry of the shared `tests/test_helpers/cake_package.py` only)	`KeyError`	`"cake_moe_grouped_det_gemm"` (and `"cake.cake_moe_grouped_gemm"` for the non-det package)
`tests/experimental/test_cake_moe_grouped_det_gemm_alias.py`	deleted	absence tests in `test_cake_moe_grouped_det_gemm_api.py` / `test_cake_moe_grouped_gemm_api.py`
det benchmark `benchmarks/bench_cake_moe_grouped_det_gemm.py`: arms `cake`, `cake_inforward`; functions `cake_arm`, `cake_inforward_arm`, `_cake_vs_inforward`; JSON key `cake_vs_inforward`	`parser.error` / `AttributeError`	arms `cake_det`, `cake_det_inforward`; `cake_det_arm`, `cake_det_inforward_arm`, `_cake_det_vs_inforward`; key `cake_det_vs_inforward`
non-det `cake_backend._check_accumulate_call` (used only by the old facade det route)	deleted	-

Unchanged names:

flashinfer.grouped_mm.grouped_mm_bf16(..., backend="cake") keeps the value "cake". It runs the forward pass only, and the det and non-det forward programs are bitwise identical.
The non-det package, the facade flashinfer.experimental.cake.grouped_mm and the non-det program names stay as they are.

Editable installs: a stale flashinfer/experimental/cake_moe_grouped_gemm/__pycache__ makes the removed path import as an empty namespace package. To fix this, run git clean -fdX flashinfer/experimental/cake_moe_grouped_gemm. Fresh clones, wheels and CI are not affected.

Mode rules

Entry	`deterministic`	`torch.use_deterministic_algorithms(True)`	`warn_only=True`
det package, `flashinfer.grouped_mm.cake_det`, `flashinfer.grouped_mm.{grouped_det_gemm_*, cake_grouped_det_mm}`	No parameter. Passing it (`True`/`False`/`None`) raises `TypeError`.	No alert, no warning	same
`CakeGroupedDetMm.apply(x, w, offs[, None[, fuse]])`	A 4th argument other than `None` (e.g. a stale `True`/`False`/`0`/`1`) raises `TypeError` before any launch. A non-bool `fuse` raises `ValueError`.	No change	same
`grouped_mm_bf16(backend="cake")`	No parameter. Passing it raises `TypeError`.	Det forward, no alert	same
det accumulate: `grouped_det_gemm_wgrad(..., out=out, accumulate=True)` / `cake_grouped_det_mm(..., fuse_wgrad_accumulation=True)`	n/a	No alert. Result has the AccumulateGrad bits of the det dW and is bitwise identical run to run.	same
non-det package entries (`grouped_gemm_`, `prepare_grouped_gemm_`, `cake_grouped_mm`) and the four facade entries	Keyword-only `deterministic: bool = True`. Only the bool `False` runs. Omitted/`True`/`None`/`0`/`1`/`"False"`/`np.False_` raise `ValueError`, which names `flashinfer.experimental.cake_moe_grouped_det_gemm.<det entry>` and says "pass deterministic=False for the non-deterministic path".	`RuntimeError` naming the entry and its det entry, raised right before an order-varying launch only	`UserWarning` with the same names, then the launch runs
`CakeGroupedMm.apply(x, w, offs[, mode[, None[, fuse]]])`	The 4th argument is the mode, with default `True`, so omitting it raises `ValueError` like every value other than `False`. Anything but `None` in the 5th (internal) slot raises `TypeError`, so `apply(x, w, offs, False, True)` no longer drops the fuse flag. The 6th argument is the fuse flag.	as above	as above

No routing to det. Nothing routes to the det package and there is no fallback. The facade's old det route is gone; call the det entry directly.
Order-varying launches. Only the weight gradient's split tail varies in order: the 512-wide k tile with unordered_tail and tail_splits > 1, in overwrite or accumulate mode. It is reached through:
- grouped_gemm_wgrad
- GroupedGemmLaunch.launch(), checked at every launch
- CakeGroupedMm.backward when it computes dW
What the alert names. The alert names the det replacement: grouped_det_gemm_wgrad, GroupedDetGemmLaunch.launch or CakeGroupedDetMm.
What never alerts. fwd, dgrad, k256 or unsplit wgrad, and preparation never alert.
Where the warning points. Under warn_only, the warning points at the caller's line for direct calls and at a PyTorch autograd frame for the backward.
No shared state between packages. No tape, saved tensor or workspace crosses between the two packages.
Other Cake families. This is the rule the other Cake families adopt.

Deterministic weight-gradient accumulation

API

grouped_det_gemm_wgrad. New signature grouped_det_gemm_wgrad(g, x, offs, out=None, out_dtype=None, num_groups=None, *, accumulate=False), available in the det package, flashinfer.grouped_mm.cake_det and flashinfer.grouped_mm.
cake_grouped_det_mm. New signature cake_grouped_det_mm(x, w, offs, *, fuse_wgrad_accumulation=False), in the same three places, plus CakeGroupedDetMm.forward(ctx, x, w, offs, _launch=None, fuse_wgrad_accumulation=False). With fuse_wgrad_accumulation=True the backward accumulates into w.main_grad, following the Megatron main_grad / grad_added_to_main_grad convention used by the non-det cake_grouped_mm.
Prepared launches. cake_backend.prepare_grouped_det_gemm_wgrad(..., *, accumulate=False), cake_backend.ACCUMULATE_MODES = ("off", "composed"), plan["accumulate"] and GroupedDetGemmLaunch.add_from.
Benchmark. The non-det benchmark benchmarks/bench_cake_moe_grouped_gemm.py gains the arm det_mgX and the ratios det_mgX/det_acc, det_mgX/det_mgB and mgX/det_mgX.

Design and numerics

Composed accumulation. The unchanged deterministic dW programs write to a launch-owned scratch of out's dtype, and then out.add_ runs on the same stream. Accumulation uses exactly the same plan and stages as the overwrite path.
Bitwise equality. For fp32 out the result is RN32(out + dW). For bf16 out, the bf16 dW is added in bf16. Both are bitwise equal to AccumulateGrad of the deterministic dW and bitwise identical from run to run.
Overlap checks. A w.main_grad with overlapping rows or experts (expanded, row stride < K) raises ValueError at the recorded forward. Memory shared with the operands is still refused at the backward.
CUDA graphs. Both autograd backwards (det and non-det) query the current stream only outside a capture, guarded by torch.cuda.is_current_stream_capturing(). As a result, capturing the fused backward in a CUDA graph no longer fails at a stream query.
Scratch lifetime. The scratch has no record_stream. If you launch a prepared launch on a stream other than the one it was prepared on, keep the launch alive until that launch completes.

Cost

The transient scratch takes 4*E*N*K bytes for an fp32 out. On the #8602 MoE shapes:

shape	E=4	E=8	E=16	E=32
gate_up (N=4096, K=6144)	384 MiB	768 MiB	1.5 GiB	3 GiB
down (N=6144, K=2048)	192 MiB	384 MiB	768 MiB	1.5 GiB

Under torch.use_deterministic_algorithms(True) with fill_uninitialized_memory on, PyTorch NaN-fills the scratch when it is allocated, which adds the cost of that fill. The extra write and read of the scratch cost time; see Performance.

A fused deterministic accumulate program, which would not need the scratch, is future work.

Performance

Measured per architecture (B200 sm_100a, GB300 sm_103a, R200 sm_107a); table follows.

Tests

CPU suites pass; GPU suites on B200 / GB300 / R200: results follow.

…y aliases The deterministic MoE grouped GEMM kept its names from before the deterministic-family rename as aliases. They are removed with no replacement: - the package flashinfer.experimental.cake_moe_grouped_gemm and its cake_autograd, cake_backend, cake_build and cake_jit submodules; - the _RENAMED maps and their __getattr__ branches in the det package (CakeGroupedMm, GroupedGemmLaunch, cake_grouped_mm, grouped_gemm_*, prepare_grouped_gemm_*), its cake_autograd (CakeGroupedMm, cake_grouped_mm) and its cake_backend (GroupedGemmLaunch, grouped_mm_bf16_cake, prepare_grouped_gemm_*); - flashinfer.grouped_mm.{grouped_gemm_fwd, grouped_gemm_dgrad, grouped_gemm_wgrad, cake_grouped_mm} and the core module's map; - the test-helper key cake_moe_grouped_gemm. The det-free spellings now name only the non-deterministic package flashinfer.experimental.cake.cake_moe_grouped_gemm and its facade.

…e to cake_det The thin core module holds only the deterministic entry points (grouped_det_gemm_fwd / _dgrad / _wgrad and cake_grouped_det_mm), so its name now carries det like the other deterministic thin modules. flashinfer.grouped_mm imports from .cake_det; the old module name is removed (no alias module). flashinfer.grouped_mm.grouped_mm_bf16 and the exported names of flashinfer.grouped_mm are unchanged.

… of the det entries cake_grouped_det_mm (det package and flashinfer.grouped_mm) took a deterministic argument and silently ran the deterministic programs for deterministic=False. The deterministic entries now take no mode argument: - cake_grouped_det_mm(x, w, offs): passing deterministic raises TypeError (unexpected keyword); - CakeGroupedDetMm.apply(x, w, offs): forward accepts only the internal deferred-launch marker in the fourth position, so a stale positional apply(x, w, offs, True) or apply(x, w, offs, False) raises TypeError instead of binding the mode to the launch slot. torch.use_deterministic_algorithms(True) changes nothing for these entries. The facade's deterministic route and the det benchmark's direct apply arm call the new signatures.

…deterministic only The selectable facade flashinfer.experimental.cake.grouped_mm ran the deterministic package by default, rerouted deterministic=False to it under torch.use_deterministic_algorithms(True) with a one-time warning, composed the deterministic accumulate itself and raised NotImplementedError for the deterministic fuse_wgrad_accumulation. It now hands every call to the non-deterministic package with the same arguments; the deterministic routing, the fallback and its warning, the facade compose and the NotImplementedError route are removed (the deterministic accumulate belongs to the deterministic package). Mode rules of the package and the facade: - deterministic is keyword-only on every entry, cake_grouped_mm included; every value but the bool False (the default True, None, 0, 1, "False") raises ValueError naming the deterministic entry (grouped_det_gemm_*, prepare_grouped_det_gemm_*, cake_grouped_det_mm) and saying to pass deterministic=False for the non-deterministic path; - under torch.use_deterministic_algorithms(True), PyTorch's non-determinism alert (RuntimeError; warn_only=True: UserWarning, then the launch runs) is raised immediately before an order-varying launch only: GroupedGemmLaunch.order_varying, the weight gradient whose 512-wide k tile splits its tail (overwrite or accumulate), at launch(), in grouped_gemm_wgrad and in CakeGroupedMm.backward. The preparation, the forward, the data gradient, the 256-wide k tile and an unsplit tail run in a fixed order and never alert. _check_accumulate_call, used only by the removed facade compose, is removed from the package.

…rministic mode contract - Delete test_cake_moe_grouped_det_gemm_alias.py (it asserted the removed aliases). - Deterministic API suite: absence tests for the removed package path, its submodules and flashinfer.grouped_mm.cake (ModuleNotFoundError), every removed det-free name on the det package, its cake_autograd and cake_backend, flashinfer.grouped_mm.cake_det and flashinfer.grouped_mm (AttributeError / ImportError) and the test-helper key (KeyError); no det entry has a deterministic parameter; cake_grouped_det_mm and the det functions refuse deterministic= (TypeError); a stale positional CakeGroupedDetMm.apply(x, w, offs, True|False) raises TypeError; torch deterministic mode (strict and warn_only) changes nothing for the det entries. - Non-deterministic API suite: the facade hands every call to the package (deterministic objects are tripwires), refuses every value but the bool False with ValueError naming the det entry, has keyword-only deterministic on every entry; the alert is raised immediately before an order-varying launch only (recorded launches on CPU: strict RuntimeError with nothing launched, warn_only UserWarning naming the caller then the launch; fixed-order launches never alert; the grouped_gemm_wgrad name through the package and the facade); the det-free names are only non-deterministic objects; on a GPU, a split and an ordered weight gradient and the autograd backward. - Accumulate suite: order_varying follows the plan (meta operands at every SM count), the accumulate alert contract, the facade accumulate and fuse rows hand off to the package and refuse the default mode; the facade's deterministic compose test is removed; the unsplit autograd case runs under torch deterministic mode without an alert, and a split-tail fuse alerts before touching main_grad. - Det backend suite: call sites without deterministic; the direct apply arm refuses a stale mode argument. - Reproducibility suite: the facade fwd / dgrad (deterministic=False) give the deterministic bits.

…e the det benchmark arms - docs/api/cake_moe_grouped_det_gemm.rst and the det package README: a "Mode and dispatch rules" section (deterministic only, no deterministic argument: TypeError, also for a fourth argument of CakeGroupedDetMm.apply; torch deterministic mode changes nothing; grouped_mm_bf16(backend="cake") is the deterministic forward) and the removal of the det-free names with no alias; the deprecated-alias sections and the alias test are gone. - docs/api/cake_moe_grouped_gemm.rst and the ND package README: the non-deterministic-only rules (keyword-only deterministic, ValueError naming the det entry for every value but False, the alert only before an order-varying launch, no deterministic route or fallback), the facade's accumulate rows, examples with deterministic=False and the det / ND name table without aliases. - benchmarks/bench_cake_moe_grouped_det_gemm.py: the det arms are cake_det and cake_det_inforward (the JSON comparison key is cake_det_vs_inforward); the ND benchmark's docstring no longer names a facade deterministic route.

prepare_grouped_det_gemm_wgrad and grouped_det_gemm_wgrad (also flashinfer.grouped_mm.grouped_det_gemm_wgrad) take a keyword-only accumulate=False. accumulate=True adds the deterministic weight gradient into the caller's fp32 or bf16 out: out[e] += G[e].T @ X[e]. The launch runs the unchanged overwrite plan into a contiguous scratch of out.dtype owned by the prepared launch, then out.add_ on the same stream outside autograd (inference mode for an inference out). out then holds AccumulateGrad's bits of the deterministic weight gradient in its dtype, bitwise reproducible run to run. plan["accumulate"] is "composed" or "off"; every other plan key is the overwrite's. The checks of out (given, out_dtype None or out.dtype, bf16 or fp32, no self overlap, no byte shared with g, x or offs) run before anything is allocated or bound. sum_m == 0 launches nothing and leaves out unchanged. Every launch and every CUDA graph replay adds again.

…n into main_grad cake_grouped_det_mm (also flashinfer.grouped_mm.cake_grouped_det_mm) takes a keyword-only fuse_wgrad_accumulation=False with the Megatron / Transformer Engine convention of the non-deterministic wrapper: - w.main_grad (fp32 or bf16, w's shape and device, unit stride along K) is checked at a forward call that autograd records and read again at the backward; - the backward adds the deterministic dW into it with the package's accumulate mode (AccumulateGrad's bits of that dW, run to run bitwise) and returns a fresh placeholder when w has grad_added_to_main_grad (zeros under zero_out_wgrad), else None; - every backward adds: retain_graph twice adds twice, and so does an autograd.grad whose inputs exclude w; - create_graph, torch.jit.trace and torch.compile raise NotImplementedError; a reentrant checkpoint raises ValueError. CakeGroupedDetMm.forward takes fuse_wgrad_accumulation as its fifth argument and refuses a non-bool value (ValueError). The fourth argument stays the internal deferred-launch slot: anything but None or the internal sentinel still raises TypeError. Torch's deterministic mode changes nothing. The docstrings state the scratch and NaN-fill cost.

…nt accumulation New tests/experimental/test_cake_moe_grouped_det_gemm_acc.py. CPU (no program loaded): keyword-only switches and signatures; the accumulate plans equal the overwrite plans on the owner rows at every architecture and part SM count, and every bound program writes a contiguous launch-owned scratch; forced-form shapes with and without a split tail exist on every architecture; the out checks refuse before anything is allocated or bound; sum_m == 0; the launch adds once after the programs, outside autograd, also into an inference tensor, with no alert under torch's deterministic mode; the fuse flag's refusals, trace and compile; and the Megatron contract of the autograd wrapper with CPU stand-ins for the preparations (placeholder, zero_out_wgrad, retain_graph, autograd.grad without w, frozen w, re-pointed main_grad, create_graph, direct apply, reentrant / non-reentrant / nested checkpoints). GPU (oracle: AccumulateGrad of the deterministic dW, out0.add_(dW) in out's dtype, bitwise): every weight-gradient form with and without a split tail, fp32 and bf16, run to run over 10 launches and against the fp64 reference; owner rows; empty groups; sum_m == 0; row padding, one group, m_indptr, inference tensor; two streams; CUDA graph replays; allocations; version counter; torch's deterministic mode; and the autograd flag with the generated programs, including a bf16 main_grad against autograd's own AccumulateGrad. The det API suite also runs the accumulate and fuse entries under torch's deterministic mode.

…ation and its cost The package README and the API page of the deterministic package describe accumulate=True and fuse_wgrad_accumulation=True: the composed launch (the overwrite plan into a launch-owned scratch, then out.add_), the numerical contract (AccumulateGrad's bits of the deterministic dW in out's dtype, bitwise reproducible run to run, empty groups, sum_m == 0), the checks, prepared launches and CUDA graphs, the Megatron convention of the autograd flag, and the cost: a scratch of 4 * E * N * K bytes for an fp32 out (192 MiB to 3 GiB on the MoE projection rows), 8 * E * N * K bytes of extra traffic, and torch's NaN fill of the scratch and the placeholder under its deterministic mode. The mode-rule tables name the fifth argument of CakeGroupedDetMm.apply, and the non-deterministic API page points to the deterministic accumulation.

…ccumulation The accumulation rows of benchmarks/bench_cake_moe_grouped_gemm.py gain the arm det_mgX: cake_grouped_det_mm(..., fuse_wgrad_accumulation=True) into the fp32 main_grad with the grad_added_to_main_grad placeholder in the checkpointed step and the backward, and prepare_grouped_det_gemm_wgrad(..., accumulate=True) in the weight-gradient row. acc_ratios adds det_mgX/det_acc, det_mgX/det_mgB and mgX/det_mgX.

…en the autograd checks Both autograd functions bound the CUDA context of the backward's thread with a stream query, which a CUDA graph capture of the backward forbids. The backward now makes the capture-status query first and queries the stream only outside a capture. CakeGroupedMm.apply refuses anything but None in its internal fifth slot (TypeError), so apply(x, w, offs, False, True) no longer drops the fuse flag, and its omitted mode is True and refused like every other value but False. The non-determinism alert names the deterministic entry. The deterministic main_grad check at the forward also refuses overlapping rows or experts. Tests: the capture of the fused backward (GPU, both packages), the capture-status logic, the internal slot, the alert's entry names, the facade under warn_only, the forward main_grad layout check, bf16 forced form cases, fresh-scratch repeats, every form under torch's deterministic mode, and a per-rounding bf16 budget.

…rules The accumulate launch's add bumps out._version at launch() and at capture, not at a graph replay; its scratch lives on the preparation stream, so a launch run on another stream must stay alive until that launch completes; the backward's peak rises by 2 * S bytes only for an fp32 main_grad. The forward checks main_grad's layout, the backward also its memory sharing. The warn_only warning names the caller's line for a direct call, a frame of PyTorch's autograd for a backward; the alert names the deterministic entry. CakeGroupedMm.apply's omitted mode and internal slot are stated. The deterministic API page title names its mode, and the README tells how to drop a stale cache of the removed package.

coderabbitai · 2026-11-06T05:51:53Z

Caution

Review failed

The pull request is closed.

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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.

Share

yyihuang · 2026-11-06T05:51:56Z

@flashinfer-bot run tests/experimental/test_cake_moe_grouped_det_gemm.py tests/experimental/test_cake_moe_grouped_det_gemm_acc.py tests/experimental/test_cake_moe_grouped_det_gemm_api.py tests/experimental/test_cake_moe_grouped_det_gemm_determinism.py tests/experimental/test_cake_moe_grouped_det_gemm_standalone.py tests/experimental/test_cake_moe_grouped_gemm.py tests/experimental/test_cake_moe_grouped_gemm_acc.py tests/experimental/test_cake_moe_grouped_gemm_api.py tests/experimental/test_cake_moe_grouped_gemm_reproducibility.py tests/experimental/test_cake_moe_grouped_gemm_standalone.py tests/grouped_mm

yzh119

LGTM. Nit: the scratch in accumulate mode is a full 4EN*K allocation per backward, so a fused accumulate program would be worth a follow-up.

flashinfer-bot

CODEOWNERS gap attestation (not a code review).

A code owner of a changed path has approved this PR:

@yzh119 -> flashinfer/grouped_mm/cake_det.py
@yzh119 -> flashinfer/grouped_mm/init.py
@yzh119 -> tests/test_helpers/cake_package.py

The following files match only the catch-all rule; this approval covers that technical gap under the one-approval policy. GitHub still requires each owned path's own owner.

benchmarks/bench_cake_moe_grouped_det_gemm.py
benchmarks/bench_cake_moe_grouped_gemm.py
docs/api/cake_moe_grouped_det_gemm.rst
docs/api/cake_moe_grouped_gemm.rst
flashinfer/experimental/cake/cake_moe_grouped_gemm/cake_autograd.py
flashinfer/experimental/cake/cake_moe_grouped_gemm/cake_backend.py
flashinfer/experimental/cake/cake_moe_grouped_gemm/init.py
flashinfer/experimental/cake/cake_moe_grouped_gemm/README.md
flashinfer/experimental/cake/grouped_mm.py
flashinfer/experimental/cake_moe_grouped_det_gemm/cake_autograd.py
flashinfer/experimental/cake_moe_grouped_det_gemm/cake_backend.py
flashinfer/experimental/cake_moe_grouped_det_gemm/init.py
flashinfer/experimental/cake_moe_grouped_det_gemm/README.md
flashinfer/experimental/cake_moe_grouped_gemm/cake_autograd.py
flashinfer/experimental/cake_moe_grouped_gemm/cake_backend.py
flashinfer/experimental/cake_moe_grouped_gemm/cake_build.py
flashinfer/experimental/cake_moe_grouped_gemm/cake_jit.py
flashinfer/experimental/cake_moe_grouped_gemm/init.py
tests/experimental/test_cake_moe_grouped_det_gemm_acc.py
tests/experimental/test_cake_moe_grouped_det_gemm_alias.py
tests/experimental/test_cake_moe_grouped_det_gemm_api.py
tests/experimental/test_cake_moe_grouped_det_gemm.py
tests/experimental/test_cake_moe_grouped_gemm_acc.py
tests/experimental/test_cake_moe_grouped_gemm_api.py
tests/experimental/test_cake_moe_grouped_gemm_reproducibility.py

Head: 97e1372

…policy Sync with main 59d4bcd (flashinfer-ai#8906 conv1d nondet rename, flashinfer-ai#8920, flashinfer-ai#8924, flashinfer-ai#8943, flashinfer-ai#8948 and others). Two conflicts, both in shared tables: - .pre-commit-config.yaml exclude: keep this branch's mHC csrc paths (experimental/cake/cake_mhc_{det,head_det,nondet}_train/csrc/) and take main's experimental/cake/cake_causal_conv1d_nondet_train/csrc/. - tests/test_helpers/cake_package.py PACKAGES: keep "cake.cake_mhc_nondet_train" and take main's "cake_causal_conv1d_nondet_train"; drop both former rows.

averyhNV added 13 commits November 5, 2026 19:49

This was referenced Nov 6, 2026

refactor(cake_mhc): name the mHC training modules by their mode: cake_mhc_det_train, cake_mhc_nondet_train, cake_mhc_head_det_train (#8928) #8907

Draft

[CAKE] GEMM / BMM Kernel Progress Tracker #6596

Open

flashinfer-bot added op: moe run-ci labels Nov 6, 2026

yyihuang marked this pull request as ready for review November 6, 2026 09:49

yyihuang requested review from Anerudhan, Jingfan-Sun, YangXu1990uiuc, aleozlx, bkryu, dierksen, jiahanc, righthandabacus, yongwww and yzh119 as code owners November 6, 2026 09:49

yyihuang requested review from IwakuraRein, StudyingShao, flashinfer-bot, jimmyzho, nv-yunzheq, samuellees and yanqinz2 as code owners November 6, 2026 09:49

yzh119 approved these changes Nov 6, 2026

View reviewed changes

flashinfer-bot approved these changes Nov 6, 2026

View reviewed changes

yyihuang merged commit 3260247 into flashinfer-ai:main Nov 6, 2026
43 of 45 checks passed

yyihuang mentioned this pull request Nov 6, 2026

refactor(cake_moe_grouped_gemm): name the MoE grouped GEMM modules by their mode (#8928) #8963

Draft

Navigation Menu

Repository navigation

refactor(cake_moe_grouped_gemm): det-only names, explicit deterministic=False, and deterministic weight-gradient accumulation - #8924

refactor(cake_moe_grouped_gemm): det-only names, explicit deterministic=False, and deterministic weight-gradient accumulation#8924
yyihuang merged 13 commits into
flashinfer-ai:mainfrom
yyihuang:moe-det-names-and-accumulate

yyihuang commented Nov 6, 2026

Uh oh!

coderabbitai Bot commented Nov 6, 2026 •

edited

Loading

Review failed

Uh oh!

yyihuang commented Nov 6, 2026

Uh oh!

yzh119 left a comment

Uh oh!

flashinfer-bot left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

yyihuang commented Nov 6, 2026

Summary

Breaking change: det-only names

Mode rules

Deterministic weight-gradient accumulation

API

Design and numerics

Cost

Performance

Tests

Uh oh!

coderabbitai Bot commented Nov 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review failed

Uh oh!

yyihuang commented Nov 6, 2026

Uh oh!

yzh119 left a comment

Choose a reason for hiding this comment

Uh oh!

flashinfer-bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

coderabbitai Bot commented Nov 6, 2026 •

edited

Loading