[yyihuang]

Stack 9/13 · base: main

Summary

This PR adds miles_ext_kernels.cake.dense_gemm on the shared cake runtime. It provides two BF16 GEMMs, each with forward and backward:

• batched_projection(x, weight, *, weight_layout) is a per-head projection x [T, H, D_in] → [T, H, D_out]. It serves MLA q-absorb ("in_out") and the value and output-group projections ("out_in").
• router_gating(x, weight) computes FP32 MoE router logits x @ weight.T. It matches Megatron's RouterGatingLinearFunction with --moe-router-dtype fp32 and no bias.

The PR vendors two packages from FlashInfer c5db8d3:

• dense_projection_det_gemm (det);
• dense_projection_gemm (fast), which runs its deterministic rows on the det package.

In miles, --batched-projection-gemm-backend cake and --moe-router-gemm-backend cake select the op (radixark/miles#4684). Either one requires --train-backend megatron. The command line offers det only, so every miles process runs det. Direct callers can reach fast through configure(cake_dense_gemm_variant="fast"). The previous version of this layer shipped det only.

Other changes from the previous version:

• cake.prebuild("dense_gemm", ...) takes model shapes and no token count.
• The results are documented as not batch-invariant.
• The runtime owns the variant and the arguments.
• A new extra, dense_gemm (cuda-bindings>=12.8), is listed in the README Install extras row.

Stack

This is layer 9 of 13, on top of layer 8 (gqa). 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 9. The previous version of this layer was reviewed in #9 ([Cake 9/11] dense_gemm: det batched projections and FP32 router GEMM).

This layer's commits:

• ef381a0 vendor: dense_gemm packages
• 20ccf43 dense_gemm: batched projection and FP32 router GEMMs on the cake runtime
• 9eead0b tests: dense_gemm

Changes

From the layer 8 head to the layer 9 head: 238 files, +162,232 / −0. Per commit: ef381a0 221 files +159,637; 20ccf43 6 files +640; 9eead0b 11 files +1,955.

Group	Files	Lines
Vendored cake/_vendor/dense_projection_det_gemm/	190	+138,713
Vendored cake/_vendor/dense_projection_gemm/	30	+20,680
cake/_vendor/VENDORED.json (ops.dense_gemm block)	1	+244
Op module and implementation: cake/dense_gemm.py, cake/_impl/dense_gemm/{__init__,ops}.py	3	+633
Tests: tests/cake/dense_gemm/	11	+1,955
Registry, README (Install extras, Ops and Prebuild rows), pyproject.toml (extra dense_gemm = ["cuda-bindings>=12.8"])	3	+7

Packages

Package	Upstream	Role	Libraries (sm_100a / sm_103a / sm_107a)	Minimum CUDA
dense_projection_det_gemm	flashinfer-ai/flashinfer@c5db8d3, flashinfer/experimental/dense_projection_det_gemm (190 files, Cake revision 72d0af54b2ddd96cd9f931508e70999338e18cce)	det, and the deterministic sibling that the fast package binds	268 = 87 / 90 / 91	12.8 / 12.9 / 13.4
dense_projection_gemm	same commit, flashinfer/experimental/dense_projection_gemm (30 files, same Cake revision)	fast: non-deterministic split-K (K1) and stream-K (K2) weight-gradient programs; binds the det package through use_deterministic_backend for every other row	30 = 9 / 10 / 11	12.8 / 12.9 / 13.4

Both packages come from one commit, because the fast package runs its deterministic rows on its sibling. There are no allow entries and no local files. There is no KCP variant, because the op has no context-parallel form.

python tools/vendor_cake.py fi --checkout ../flashinfer --commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 \
    --op dense_gemm --package flashinfer/experimental/dense_projection_det_gemm \
    --package flashinfer/experimental/dense_projection_gemm

FlashInfer revision

The message of commit 20ccf43 explains the choice of c5db8d3 from the FlashInfer PR history:

• From fix(cake_dense_projection_gemm): route K1 entries that fail the accuracy rule to the deterministic program flashinfer-ai/flashinfer#8671 and fix(cake_dense_projection_gemm): route router windows that fail the accuracy rule to the deterministic program flashinfer-ai/flashinfer#8689 on, the non-deterministic package routes weight-gradient rows and router windows to the deterministic program. It does this to meet a per-seed comparison with the deterministic program's error. That rule no longer applies; the deterministic package's absolute budgets remain the bar.
• perf(cake_dense_projection_gemm): restore the K1 weight-gradient routes that meet the deterministic package's absolute FP64 budgets flashinfer-ai/flashinfer#8743 restored the weight-gradient rows that meet those budgets and moved two 148-SM sm_103a rows to the deterministic program. perf(cake_dense_projection_gemm): restore the SIMT and head/tail router windows on 152-SM and 212-SM parts flashinfer-ai/flashinfer#8730 restored the router windows.
• perf(cake_dense_projection_gemm): restore the K1 weight-gradient routes on 216-SM sm_107a parts flashinfer-ai/flashinfer#8782 restores the weight-gradient rows of 216-SM sm_107a parts. The commit message calls it open. It was merged on 2026-11-05 at 14:16Z (merge commit f8bedda6038a3eff3b6953664c385c6b9eb0f1a3). This layer still vendors c5db8d3.
• c5db8d3 carries only the earlier rows of perf(cake_dense_projection_gemm): deterministic K1 weight gradient at measured slow or less accurate token counts on sm_100a / sm_103a flashinfer-ai/flashinfer#8326. On sm_103a parts with 152 or more SMs, the single-head 5120 → 32 / 48 / 64 / 128 BF16 weight gradients run the deterministic program from 32768 tokens. miles calls none of these shapes.
• The FP32-output K1 weight-gradient rows of fix(cake_dense_projection_gemm): run the deterministic program for four FP32 sm_107a K1 rows at large token counts on 216-SM parts flashinfer-ai/flashinfer#8509 (216-SM sm_107a) are not reachable through these entry points, because projection weight gradients have BF16 output only.

API

OP is name="dense_gemm", variants=("det", "fast"), cli_variants=("det",), selected_by=(("batched_projection_gemm_backend", "cake"), ("moe_router_gemm_backend", "cake")), with packages=("dense_projection_det_gemm", "dense_projection_gemm") and no extra flags. The module is torch-free. It exports OPERATIONS (("forward", "input_grad", "weight_grad")), WEIGHT_LAYOUTS (("in_out", "out_in")), batched_projection, check_projection_dims, fast_projection_programs, fast_router_programs and router_gating.

• batched_projection:
  • Weight layouts: "in_out" takes weight [H, D_in, D_out] (einsum("thd,hdm->thm")); "out_in" takes weight [H, D_out, D_in] (einsum("thm,hdm->thd")).
  • It is differentiable in x and weight.
  • Each result is an FP32 sum of exact BF16 products, rounded once. These are cuBLAS numerics, not cuBLAS's bits.
  • Views are read in place when they are 16-byte aligned, with unit inner stride and 16-byte outer strides. Other views are copied.
  • D_in and D_out must be positive multiples of 8.
• router_gating:
  • It takes BF16 x [..., H] and weight [E, H]. E and H must be positive multiples of 8.
  • The forward is a K1 program with FP32 output.
  • The backward computes dx = g @ weight.float() and dweight = g.T @ x.float() in FP32 through the exact BF16×3 split (K2). The results are rounded to the dtypes of x / weight.
• check_projection_dims(d_in, d_out) raises ValueError unless both are positive multiples of 8.
• fast_projection_programs(...) / fast_router_programs(...) return the program that fast runs for each operation: a template name, or None for the deterministic program. They need torch, but no device.
• OP.validate: a selected backend requires --train-backend megatron; otherwise it raises ValueError.
• An unsupported device raises NotImplementedError. 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; prebuild parameters projections, router; tests in tests/cake/dense_gemm/.

Variant routing

Situation	Result
Process never configured	det
miles CLI with either selector, flag unset, deterministic or not	det: --cake-dense-gemm-variant is registered but hidden (one CLI choice, det), and validate_args writes det back
Direct caller, configure(cake_dense_gemm_variant="fast")	fast; refused under deterministic_mode=True
At a call	every forward reads config.variant("dense_gemm"); the backward runs the variant of its forward
fast under torch.use_deterministic_algorithms(True)	runs det, logged once (warn_runs_det)

• What fast changes. It changes only weight gradients, and only on the rows and token counts its tables select for the device's architecture and SM count:

  • the K1 weight gradients of single-batch rows with d_in 4096–8192 and d_out 32–1024;
  • the router weight gradient of 256 experts at hidden 4096 / 5120 / 6144, inside measured token windows.

  Forwards, input gradients and every other row run the deterministic programs and give det's bits.

• miles' own calls under fast. In the recorded plan table (token counts 64 to 65536), only one of miles' calls runs a program of its own: the router weight gradient of 256 experts at hidden 4096 / 5120 / 6144. It does so inside token windows, on 152-SM sm_100a / sm_103a parts and on 208–216-SM sm_107a parts. On 148-SM parts and 160-SM sm_103a parts, every miles call runs the deterministic program. fast becomes a CLI choice once it is shown to be reachable and correct for miles' calls.

• Known accuracy misses of fast. The _impl/dense_gemm/ops.py docstring lists these, taken from perf(cake_dense_projection_gemm): restore the K1 weight-gradient routes that meet the deterministic package's absolute FP64 budgets flashinfer-ai/flashinfer#8743 measurements.

  • On 148-SM sm_103a parts, the two-split K1 weight gradient of the single-head 5120 → 48 and 5120 → 64 BF16 projections missed the det package's BF16 budget (atol = rtol = 1e-2, no element outside) for one of 10 seeds.
  • The misses were at 40961 and 40962 tokens (5120 → 48) and at 47105 tokens (5120 → 64).
  • The deterministic program met the budget on every seed.
  • miles calls neither shape and runs det.

• Batch invariance. Neither variant is batch-invariant.

  • The plan depends on the token count. A token's outputs, logits and gradients can therefore change in their last bits with the token count of the call, within the accuracy budget.
  • As a result, a router top-k can depend on packing when two logits are that close.
  • det repeats bit for bit for identical inputs of the same shape on the same device.
  • The README Ops row, the op module, ops.py and both entry points' docstrings all say so.

Prebuild

cake.prebuild("dense_gemm", *, device=None, max_workers=None, projections=(), router=()) takes (num_heads, d_in, d_out, weight_layout) tuples in projections and (num_experts, hidden_size) tuples in router. It takes no token count.

• It builds every program that the two packages register for the device's architecture, for the GEMM kinds given.
• It checks the shapes first (ValueError), before the device. With both lists empty, it returns [].
• Afterwards, no call of those GEMM kinds compiles on that device, at any token count, in either variant.

The README Prebuild row reads: dense_gemm | projections: (num_heads, d_in, d_out, weight_layout) tuples; router: (num_experts, hidden_size) tuples. The plan of a GEMM depends on the token count, so it builds every program the two packages register for the architecture for the GEMM kinds given (the projections, the router), not one plan per shape.

Memory

• The op does not use the runtime workspace pool, so the pool bound (default 0) does not affect it.
• The vendored entry points plan once per (device, stream, view geometry) and keep a bounded cache of planned launches.
• The cached launches of a stream share a growable workspace for that device and stream. Calls on two streams never share one.
• release() (clear_convenience_caches of both packages) drops the plans and workspaces. The README Memory section lists these as out-of-pool per-stream workspaces that grow to the largest call.
• The op is not CUDA-graph safe. A later growth replaces the workspace that a graph would replay, and an eviction frees it.
• The packages use the dense_gemm extra (cuda-bindings>=12.8) to call the CUDA driver. They use it to read the L2 size when torch does not report it, and to zero-fill the FP32 router gradients under fast.

Review items

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

#	Reviewer item	Resolution	Tests
9-1	Mid-step JIT. After a prebuild at T = 33 / 1001 / 2049, 50 random T × 4 ops compiled 2 new libraries at T = 239 (10.1 s and 4.6 s). The reviewer asked for a T-independent prebuild.	prebuild now takes model shapes only. It builds every registered program of the given GEMM kinds for the architecture.	CPU: test_prebuild_plans.py::test_every_planned_program_at_any_token_count_is_prebuilt[8 parts]; test_package.py::test_prebuild_takes_shapes_only, ::test_prebuild_checks_the_shapes_before_the_device, ::test_every_registered_program_has_a_kind_and_the_kinds_split_the_registry, ::test_the_router_kind_holds_the_k2_programs_the_split_and_the_fp32_forwards, ::test_an_unknown_template_has_no_kind; test_fast_plans.py::test_every_recorded_program_is_built_by_the_package. GPU: the 4 tests of test_prebuild_gpu.py, with T = 239 included; they passed on B200 and GB300 starting from an empty cache
9-2	Not batch-invariant: 2 / 16 router comparisons differed (logits max diff 3.4e-5), and 8 / 32 projection comparisons differed (output-group max diff 9.8e-4). The reviewer asked to document it or fix the split policy.	Documented in the README row, the op module, ops.py and the entry-point docstrings, including the consequence for top-k vs packing. The split policy is unchanged.	CPU: test_package.py::test_the_docs_state_that_the_results_are_not_batch_invariant. GPU: test_determinism_gpu.py::test_det_is_not_batch_invariant_but_every_token_count_is_within_budget[5 cases]. At T ∈ {1, 33, 257, 2049, 16384}, each T must be repeatable and within the rel-RMS budget against FP64. Router top-k must be equal on rows without a near tie. Cross-T bit differences are recorded, not asserted; the B200 counts are under Validation
9-3	Router backward memory was 1.41 GiB vs 1.32 GiB (Megatron-style) at T = 32768, H = 7168, judged "comparable, not a regression".	No change.	—
cleared	Partial gradients were bit-consistent and FP64-correct, prebuild covered the listed T, and arena growth was 90 MiB.	Kept. ROUTES adds input-gradient-only and weight-gradient-only routes.	GPU: test_prebuild_gpu.py routes ...-input-gradient-only and ...-weight-gradient-only

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

• The unconfigured default is now det, under one runtime rule.
• The mid-call nvcc compiles named for this layer are resolved by the shape-only prebuild.
• Argument ownership moved to the runtime. The op declares OP.validate, and the hidden variant flag is written back.
• A one-GPU run now deselects the two-device module with -m "not multi_gpu" instead of skipping it at collection. Collecting the module touches no device.

Tests

All tests are in tests/cake/dense_gemm/.

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

• CPU, with torch but no device:

  • test_package.py (7).
  • test_fast_plans.py (9). It covers the fast plans, including the known misses.
  • test_prebuild_plans.py (1). It checks that every planned program is in the prebuilt set, for both variants, over 8 parts: sm_100a with 148 / 152 SMs, sm_103a with 148 / 152 / 160, and sm_107a with 208 / 212 / 216. The token counts are:
    • 1 to 300;
    • every tile / split / window edge ±1 (32 to 65536);
    • 239, 1001, 2049, 10240, 16172, 16231, 20480 and 49152;
    • 24 seeded random counts up to 70000.

• GPU:

  • test_batched_projection_gpu.py (6) and test_router_gpu.py (5).
  • test_fast_gpu.py (4), including test_fast_runs_det_under_torch_deterministic_algorithms and test_the_known_accuracy_misses_of_fast_are_measured_and_recorded. The latter runs 10 seeds at the 3 documented points. It runs only on 148-SM sm_103a parts and expects a failure where an element is outside the budget.
  • test_determinism_gpu.py (3): det bit for bit in one process and in a fresh one; det on the default stream and a side stream in turn; and the batch-invariance test above.
  • test_multi_gpu.py (2, multi_gpu): det gives the same bits on both devices, and fast keeps det's forward and input-gradient bits on both. The module's skip conditions are marks, so collecting it touches no device.

• GPU, test_prebuild_gpu.py (4): one prebuild, then 20 ROUTES. Each route runs in det and fast under no_build() at 15 token counts: 1, 33, 64, 128, 239, 1001, 2049, 4097, 10240, 16384, 20480, and 4 seeded random counts in 2 to 20480. The routes are:

  • 6 projection shapes: (16,192,512,in_out), (16,512,256,out_in), (64,256,512,in_out), (8,4096,1024,out_in), (1,4096,1024,out_in) and (1,4096,32,out_in);
  • an input slice of a wider head slot together with a weight split of a parameter (both read in place);
  • a weight split, and a weight that is not 16-byte aligned (copied);
  • input-gradient-only and weight-gradient-only routes, and a side stream;
  • 4 routers (256 × 4096, 256 × 5120, 256 × 6144, 384 × 7168), a router with a [T, H] input, and a router on a side stream;
  • a router and a projection under torch.use_deterministic_algorithms(True).

  The other tests check four things: that one prebuild of the model shapes serves every route and token count; that a fresh process running both variants imports no flashinfer module; that the routes reach this device's fast programs; and that a second prebuild compiles nothing and returns the same libraries.

• Budgets:

  • BF16 results: rel-RMS 3e-3, plus the det package's elementwise BF16 bar (atol = rtol = 1e-2).
  • FP32 logits: rel-RMS 5e-5, the package's bar.
  • Router FP32 weight gradient under fast: rel-Frobenius 2e-6.
  • Router top-k tie margin: 1e-4.
  • Each test records its measured errors (record_property).

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, three more checks ran:

• vendor_cake.py --check dense_gemm --upstream;
• pytest -m 'not multi_gpu' tests/cake/dense_gemm;
• pytest -m 'not multi_gpu' tests/cake/dense_gemm/test_multi_gpu.py, which must deselect the whole module.

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
ef381a0 (vendor)	465 passed, 98 skipped, 1 failed	1024 passed, 249 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
20ccf43 (op)	471 passed, 100 skipped, 1 failed	1032 passed, 249 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[dense_gemm] fails, as allowed
9eead0b (tests)	487 passed, 123 skipped, 0 failed (rc 0)	1083 passed, 260 skipped, 0 failed (rc 0)	vendor --check rc 0; pre-commit rc 0, no diff; --check dense_gemm --upstream rc 0; tests/cake/dense_gemm (image, CPU) 50 passed, 5 skipped, 6 deselected; test_multi_gpu.py with -m 'not multi_gpu': 6 deselected, nothing run (pytest rc 5, as expected)	all green

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

GPU

The GPU tests ran on the top of the stack (layer 13). Its dense_gemm files and its ops.dense_gemm block are byte-identical to this layer's. These are cake/dense_gemm.py, _impl/dense_gemm/, both _vendor/dense_projection_* trees and tests/cake/dense_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, cuda-bindings 13.4.4 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/dense_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	168 passed, 3 skipped, 290.8 s (97 libraries built; 8 ops compiled in parallel)	168 passed, 3 skipped, 160.6 s	4 passed, 115.4 s (96 libraries)	6 passed
GB300	NVIDIA GB300, sm_103a, 152 SMs	580.159.03	168 passed, 3 skipped, 238.7 s (prebuild test first; 101 libraries)	168 passed, 3 skipped, 89.5 s	in the cold run: 4 passed	6 passed
R200	sm_107a	—	pending	pending	pending	pending

• Skips. The 3 skips on both sites are test_fast_gpu.py::test_the_known_accuracy_misses_of_fast_are_measured_and_recorded at the 3 documented points. These measurements run only on 148-SM sm_103a parts. A 148-SM sm_100a B200 and a 152-SM GB300 are not such parts, so the measurements have not been run in this round.
• det bitwise. test_determinism_gpu.py::test_det_repeats_bit_for_bit_in_one_process_and_in_a_fresh_one, ::test_det_gives_the_same_bits_on_the_default_stream_and_a_side_stream_in_turn and test_router_gpu.py::test_repeated_router_launches_are_bitwise_identical passed on B200 and GB300, cold and warm.
• Two devices. test_det_gives_the_same_bits_on_both_devices (3 cases) and test_fast_keeps_det_forward_and_input_gradient_bits_on_both_devices (3 cases) passed on both sites.
• Prebuild. Starting from an empty cache, one prebuild of the model shapes served all 20 routes at 15 token counts, in det and fast, under no_build(). The routes reached this device's fast programs. A second prebuild compiled nothing. A fresh process ran both variants without importing flashinfer.
• Deterministic algorithms. fast under torch.use_deterministic_algorithms(True) ran det (test_fast_runs_det_under_torch_deterministic_algorithms).
• Measured errors against FP64 on B200. These come from the junit properties; cold and warm runs were identical. BF16 outputs and gradients had rel-RMS of at most 1.71e-3 (budget 3e-3). Router FP32 logits had rel-RMS of at most 7.92e-6 (budget 5e-5).
• Batch invariance on B200. These counts were recorded, not asserted. Each count is the number of elements whose bits differ from the same tokens of the 16384-token call.
  • Projections: the three cases (16,192,512,in_out), (8,4096,1024,out_in) and (1,4096,32,out_in) had 0 differing output and input-gradient elements at T = 1, 33, 257 and 2049.
  • Router logits differed at every T, all within the 5e-5 budget. At T = 1 / 33 / 257 / 2049, 256 experts × 5120 had 252 / 8,371 / 65,182 / 519,702 differing elements, and 384 × 7168 had 382 / 12,597 / 98,058 / 781,731.
  • Router input gradients: 0 differing elements.
• 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 first comment below, Performance (continued). Every number there is copied from FlashInfer pull request bodies; nothing was measured for this PR. It covers:

• Packages and history. The two vendored package trees, and the FlashInfer pull requests whose merges match them, perf(cake_dense_projection_gemm): cache the router entry points' calls by their own tensors flashinfer-ai/flashinfer#8539 (the exact fast package) and perf(cake_dense_projection_det_gemm): run the half-tail TMA-store program on the sm_103a 5120 x 17408 weight gradients at T 16231 / 16384 flashinfer-ai/flashinfer#8522 (the exact det package), and the earlier pull requests perf(cake_dense_projection_gemm): sm_103a router weight-gradient head/tail windows at hidden 4096 and 6144 flashinfer-ai/flashinfer#8518, perf(cake_dense_projection_gemm): 212-SM sm_107a router weight-gradient windows run the head/tail program with the A parts in tensor memory flashinfer-ai/flashinfer#8517, perf(cake_dense_projection_gemm): sm_103a router weight-gradient head/tail program with the A parts in tensor memory flashinfer-ai/flashinfer#8444, perf(cake_dense_projection_det_gemm, cake_dense_projection_gemm): programmatic dependent launch on every program, safe by default flashinfer-ai/flashinfer#8389, perf(cake_dense_projection_det_gemm): swapped small-token forward programs with per-call PDL flashinfer-ai/flashinfer#8372 and fix(cake_dense_projection_det_gemm): take the cluster rendezvous before the 2-CTA tensor-memory allocation, and add sm_107a row rules flashinfer-ai/flashinfer#8303, which measured earlier trees of one or both packages.
• det batched projection.
  • GB300 forward against cuBLAS and the previous Cake program (perf(cake_dense_projection_det_gemm): swapped small-token forward programs with per-call PDL flashinfer-ai/flashinfer#8372, perf(cake_dense_projection_det_gemm, cake_dense_projection_gemm): programmatic dependent launch on every program, safe by default flashinfer-ai/flashinfer#8389).
  • sm_107a qabs / vproj rows against the previous version (fix(cake_dense_projection_det_gemm): take the cluster rendezvous before the 2-CTA tensor-memory allocation, and add sm_107a row rules flashinfer-ai/flashinfer#8303).
• Router backward.
  • perf(cake_dense_projection_gemm): cache the router entry points' calls by their own tensors flashinfer-ai/flashinfer#8539 on GB200 and GB300 at T 64 / 256 / 1024. It compares det Cake, previous Cake and a non-Cake SOTA, which per row is the fastest of the cuBLAS / cuBLASLt / torch.mm arms that meet the router's FP32 budget.
  • perf(cake_dense_projection_gemm): sm_103a router weight-gradient head/tail windows at hidden 4096 and 6144 flashinfer-ai/flashinfer#8518 and perf(cake_dense_projection_gemm): sm_103a router weight-gradient head/tail program with the A parts in tensor memory flashinfer-ai/flashinfer#8444 on GB300 at large token counts.
  • perf(cake_dense_projection_gemm): 212-SM sm_107a router weight-gradient windows run the head/tail program with the A parts in tensor memory flashinfer-ai/flashinfer#8517 on a 212-SM R200.
• Missing.
  • No FlashInfer measurement of the router forward this layer runs (BF16 operands, FP32 output).
  • No ratio of the det package against the non-Cake SOTA on the router rows.
  • No copied measurement of fast's K1 weight gradients.
  • No FlashInfer number for B200 (148-SM sm_100a) or B300 on these call paths.

Notes

• sm_107a: both packages ship sm_107a libraries (91 det, 11 fast) and need CUDA 13.4 to build them; an older toolkit raises RuntimeError. The CPU planner sweep covers sm_107a parts with 208, 212 and 216 SMs. No GPU run on sm_107a has been recorded for this tree yet.
• perf(cake_dense_projection_gemm): restore the K1 weight-gradient routes on 216-SM sm_107a parts flashinfer-ai/flashinfer#8782, which restores the weight-gradient rows of 216-SM sm_107a parts, was merged after these packages were vendored. Its merge commit is f8bedda6038a3eff3b6953664c385c6b9eb0f1a3, merged on 2026-11-05. The message of commit 20ccf43 still calls it open.
• Neither variant is batch-invariant, and the op is not CUDA-graph safe.
• D_in, D_out, E and H must be multiples of 8.
• fast is not a CLI choice yet. It has the known accuracy misses listed above, on shapes miles does not call. Those misses were not re-measured in this round, because no 148-SM sm_103a part was in the GPU runs.

🤖 Generated with Claude Code

[Zhichenzzz]

Review of layer 9 (dense_gemm) — suggest approve with nits

Confirmed on GPU

• Mid-step JIT. After prebuild at T = 33/1001/2049, a sweep of 50 random T × 4 ops compiled 2 new libraries at T=239: glm5_qabs took 10.1 s and dsv4_ogroup 4.6 s. This needs a T-independent prebuild.

• The ops are not batch-invariant. The same row in batches of T ∈ {1, 33, 257, 2049, 16384} changes bits in:

  • 2 of 16 router comparisons (logits max diff 3.4e-5);
  • 8 of 32 projection comparisons (dsv4 o-group out max diff 9.8e-4).

  Routing can therefore depend on packing, which matters for R3 and true-on-policy. Please document it, or fix the split policy.

• Router backward memory. Peak extra at T=32768, H=7168 is 1.41 GiB (bf16 dx plus an FP32 [T, H]), vs 1.32 GiB for a Megatron-style torch backward. Comparable, not a regression.

Cleared on GPU

• Partial grads. 12 cases are bit-consistent and FP64-correct.
• Prebuild coverage. The listed T values are covered for MLA, DSV4 and the router.
• Arena growth is 90 MiB over the sweep.

---

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.