Repository navigation
perf(cake_moe_grouped_det_gemm): skip forward recompute under checkpoint early stop - #8247
Conversation
…int early stop cake_grouped_det_mm now launches the forward GEMM after CakeGroupedDetMm.apply returns, after autograd saved x, w and offs, as cake_grouped_mm of the non-deterministic package does. Under a non-reentrant torch.utils.checkpoint with early stop on (the torch default), the recomputation then stops before the GEMM when these are the last saves of the innermost checkpointed region, as for torch._grouped_mm. A later save in the region makes it run, with the same bits. - The outputs and both gradients are bitwise unchanged; the signature of cake_grouped_det_mm is unchanged. A direct CakeGroupedDetMm.apply, and the wrapper under torch.jit.trace and torch.compile, keep the launch inside forward. - Tests: the order of the saves and the launch, the forward launch and CUDA kernel counts of checkpointed steps with early stop on and off (also against torch._grouped_mm), regions with later saves, an MoE expert block, reentrant and nested checkpoints, gradient accumulation, retain_graph, frozen operands, CUDA graph capture, memory, no host synchronisation, jit.trace and torch.compile; the facade and core entry points for the deterministic package. - bench_cake_moe_grouped_det_gemm.py gains checkpointed step rows with early stop off, so every arm recomputes its forward (equal work), and opt-in step_es / block / block_es rows, a cake_inforward control arm and a --census kernel count.
|
@flashinfer-bot run tests/experimental/test_cake_moe_grouped_det_gemm.py tests/experimental/test_cake_moe_grouped_det_gemm_api.py tests/experimental/test_cake_moe_grouped_gemm_api.py |
|
No actionable comments were generated in the recent review.
|
| Layer / File(s) | Summary |
|---|---|
Deferred forward launch flashinfer/experimental/cake_moe_grouped_det_gemm/cake_autograd.py, flashinfer/experimental/cake_moe_grouped_det_gemm/README.md, flashinfer/grouped_mm/cake.py, flashinfer/experimental/cake/grouped_mm.py, benchmarks/bench_cake_moe_grouped_gemm.py |
The wrapper defers the forward GEMM until after CakeGroupedDetMm.apply returns and autograd saves the inputs. Direct apply and tracing or compilation use the direct-launch path. The documentation describes the checkpoint behavior. |
Checkpoint and execution-mode tests tests/experimental/test_cake_moe_grouped_det_gemm.py, tests/experimental/test_cake_moe_grouped_det_gemm_api.py, tests/experimental/test_cake_moe_grouped_gemm_api.py |
Tests check saved tensors, launch counts, outputs, and gradients across checkpoint variants, MoE blocks, CUDA graphs, tracing, compilation, and validation cases. |
Checkpointed benchmark workloads benchmarks/bench_cake_moe_grouped_det_gemm.py, benchmarks/bench_cake_moe_grouped_gemm.py |
The benchmark adds checkpointed step and block workloads, a direct-forward Cake arm, percentile timings, kernel counts, and output and gradient comparisons. It rejects unsupported autograd workloads for cuBLAS and cuDNN arms. |
Estimated code review effort: 4 (Complex) | ~50 minutes
Change: Refactor
Sequence Diagram(s)
sequenceDiagram
participant Wrapper
participant CakeGroupedDetMm
participant Autograd
participant CUDA
Wrapper->>CakeGroupedDetMm: apply with deferred launch
CakeGroupedDetMm->>Autograd: save x, w, and offs
CakeGroupedDetMm-->>Wrapper: return output
Wrapper->>CUDA: launch forward GEMM
Suggested reviewers: aleozlx
Merge Risk: 
Minimal · up to 24782
-
No concrete issue remains that warrants delaying merge beyond normal validation.
Pre-merge checks: 
3 · 
2
-
Failed checks: 2 warnings Passed checks: 3 passedFull details: Description check
-
Explanation
The description gives a detailed and relevant explanation of the implementation, behavior, compatibility paths, tests, and benchmarks. However, it omits the repository template sections and required checklist information, including related issues, pre-commit status, test status and targets, and the experimental-track tracking issue, ownership, graduation plan, and experimental test targets.
Resolution
Add the required template sections. State related issues or explicitly indicate none. Complete the pre-commit and test checklist, including test results and suggested or required CI targets. Because this PR changes flashinfer/experimental/ code and tests, complete the Experimental Track section with a tracking issue, owner, rationale, graduation plan, compliance confirmations, and an experimental-tests block listing the relevant tests.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 61.22% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 49 functions across 8 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
Generate unit tests (beta)
-
- Create a new PR
-
- 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.
flashinfer-bot
left a comment
There was a problem hiding this comment.
CODEOWNERS gap attestation (not a code review).
A code owner of a changed path has approved this PR:
- @yzh119 -> flashinfer/grouped_mm/cake.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
- flashinfer/experimental/cake/grouped_mm.py
- flashinfer/experimental/cake_moe_grouped_det_gemm/cake_autograd.py
- flashinfer/experimental/cake_moe_grouped_det_gemm/README.md
- 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_api.py
Head: 2478225
…he least-loaded clusters on sm_100a and sm_103a (#8269) The Cake ragged BF16 grouped GEMM for MoE expert training has two variants. Both use FP32 accumulation. Each variant has forward (`Y = X W^T`), data-gradient (`dX = G W`) and weight-gradient (`dW = G^T X`, BF16 or FP32 output) programs. The deterministic variant, `flashinfer.experimental.cake_moe_grouped_det_gemm`, is bitwise reproducible from run to run. Callers select the non-deterministic variant, `flashinfer.experimental.cake.cake_moe_grouped_gemm`, with `deterministic=False`. The facade `flashinfer.experimental.cake.grouped_mm` also reaches it, but the facade default stays deterministic. This PR changes only the sm_100a (B200) and sm_103a (GB300) k512 wgrad half-release programs of the non-deterministic variant (BF16 and FP32 output), which take the planned tile walk. They gain the slack placement of the tail that #8258 added on sm_107a: one epilogue warp ranks the clusters by the load of their full-wave tiles. When every tail tile fits in the room the full waves leave on the least-loaded clusters, each tail tile runs whole on one cluster, with no tail sum and no ticket; otherwise the tail splits as before. The plain planned half-release program they used before is removed, so one program per output dtype now serves the half-release stage on sm_100a, sm_103a and sm_107a. The sm_107a programs, the other k512 programs, the forward and data-gradient programs, the Python API and FP32 accumulation are unchanged. As before for this variant, the summation order of a split tail tile can differ. The "vs previous non-deterministic Cake" column is this PR's own effect on GB300 (sm_103a) and B200 (sm_100a). On GB300, the geomeans vs the previous kernel are listed below, with the lowest and highest row in parentheses: - dW owner 1.000 (0.989-1.009), dW unsorted 1.001 (0.995-1.010) - ckpt owner 1.000 (0.995-1.005), ckpt unsorted 1.001 (0.999-1.004) - ckpt ES owner 1.001 (0.999-1.008), ckpt ES unsorted 1.000 (0.996-1.004) Six rows are above 1.0 vs previous over their whole [min, max] group range: dW down E4 M112576 skew asc 1.010 [1.006, 1.016], dW down E4 M136064 skew asc 1.008 [1.002, 1.015], dW down E4 M112576 skew 1.007 [1.001, 1.010], dW gate_up E32 M131072 bal 1.004 [1.002, 1.020], ckpt ES down E4 M136064 skew asc 1.004 [1.003, 1.006] and dW down E4 M136064 1-empty asc 1.002 [1.001, 1.003]. No row is below 1.0 over its whole group range. Regressions, reported as is: 24 of the 100 row medians are below 1.0 vs previous. The lowest are dW down E32 M131072 bal 0.989 [0.977, 1.001] and dW gate_up E4 M112576 bal 0.991 [0.978, 1.009]; all other rows are at 0.995 or above. On B200 (sm_100a), the geomeans vs the previous kernel are listed below, with the lowest and highest row in parentheses: - dW owner 1.003 (0.989-1.018), dW unsorted 1.000 (0.981-1.020) - ckpt owner 1.003 (0.995-1.018), ckpt unsorted 1.003 (0.997-1.016) - ckpt ES owner 1.005 (0.992-1.017), ckpt ES unsorted 1.002 (0.991-1.017) One row is above 1.0 vs previous over its whole [min, max] group range: ckpt ES down E4 M112576 skew 1.017 [1.006, 1.025]. One row is below 1.0 over its whole group range: ckpt ES gate_up E4 M136064 skew asc 0.991 [0.983, 0.997]. For the same shape, the dW row (the only part of that step this PR changes) is 1.000 [0.978, 1.022], the ckpt row is 1.002 [0.997, 1.022], and on GB300 the same ckpt ES row is 1.002 [0.988, 1.007]. Regressions, reported as is: 27 of the 100 row medians are below 1.0 vs previous. The lowest are dW gate_up E4 M112576 skew asc 0.981 [0.980, 1.000] and dW down E4 M112576 bal 0.989 [0.978, 1.021]. Against the deterministic variant on GB300, the geomeans (lowest-highest row) are dW owner 1.001 (0.992-1.014), dW unsorted 1.040 (0.983-1.115), ckpt owner 1.000 (0.996-1.005), ckpt unsorted 1.008 (0.994-1.027), ckpt ES owner 1.001 (0.991-1.008) and ckpt ES unsorted 1.013 (0.998-1.038). The lowest row is dW gate_up E4 M136064 skew asc 0.983 [0.980, 0.996] and the highest is dW down E4 M136064 1-empty asc 1.115. Against the non-Cake SOTA on GB300, the geomeans (owner / unsorted) are dW 1.045 / 1.053, ckpt 1.187 / 1.188 and ckpt ES 1.185 / 1.187; no row is below 1.0. Against the deterministic variant on B200, the geomeans (lowest-highest row) are dW owner 1.004 (0.990-1.035), dW unsorted 1.058 (1.008-1.164), ckpt owner 1.004 (0.992-1.016), ckpt unsorted 1.014 (1.004-1.039), ckpt ES owner 1.004 (0.989-1.012) and ckpt ES unsorted 1.018 (1.003-1.040). No B200 row is below 1.0 over its whole group range. The lowest row is ckpt ES gate_up E32 M131072 bal 0.989 [0.980, 1.010] and the highest is dW down E4 M136064 1-empty asc 1.164 [1.121, 1.184]. Against the non-Cake SOTA on B200, the geomeans (owner / unsorted) are dW 1.058 / 1.062, ckpt 1.237 / 1.254 and ckpt ES 1.228 / 1.255; no row is below 1.0. The ckpt ES rows are at parity with the deterministic variant because the deterministic variant on main now also skips the forward recompute under checkpoint early stop (#8247). Before #8247 the non-deterministic variant was ahead there because of #8114. **Baselines** - deterministic Cake: `cake_moe_grouped_det_gemm` on main (with #8247). - previous non-deterministic Cake: `cake_moe_grouped_gemm` on main with #8258, before this PR. - non-Cake SOTA: per row, the fastest supported of `torch._grouped_mm` (PyTorch 2.14.1+cu130), a per-expert cuBLAS loop (`torch.mm` per group) and the cuDNN grouped GEMM. It is the strongest applicable one because it is the fastest of the three on that row. Key in the SOTA cell: t = `torch._grouped_mm`, c = per-expert cuBLAS loop, d = cuDNN. - Coverage limits: cuDNN has no K-offset (weight gradient) form. The cuBLAS loop has no autograd form, so checkpointed rows compare with `torch._grouped_mm`. `torch._grouped_mm` requires the output dtype of `mat_a`, so FP32 dW rows compare with the cuBLAS loop. **Operations and shapes** - dW: weight gradient, BF16 output unless the shape says fp32. - ckpt: one grouped GEMM in a non-reentrant checkpoint, forward + backward (dX and dW), early stop off (2 x fwd + bwd); ckpt ES: early stop on. - BF16; gate_up N=4096 K=6144, down N=6144 K=2048. E4 M112576 bal = E=4 experts, M=112576 total rows, balanced; skew / 1-empty = skewed / one empty group (descending); skew asc / 1-empty asc = the same sizes ascending; random = lognormal sizes (CV 0.35, seed 0, drawn order); fp32 = FP32 dW output. - Row class: unsorted = skew asc, 1-empty asc and random rows; owner = all other rows. - Only dW and checkpointed rows are listed, because they are the rows that run the changed programs. Speedup = baseline latency / new latency, the median of the per-group paired ratios over 6 groups per row; the per-row table gives [min, max] over the groups. On GB300, their spread over the 100 rows, (max - min) / median, is median 1.37 %, p90 3.18 % vs deterministic Cake; median 1.43 %, p90 3.51 % vs previous non-deterministic Cake; and median 1.75 %, p90 3.55 % vs non-Cake SOTA. On B200, it is median 3.28 %, p90 6.00 % vs deterministic Cake; median 3.42 %, p90 6.13 % vs previous non-deterministic Cake; and median 4.39 %, p90 7.25 % vs non-Cake SOTA. Geometric mean of the per-row speedups. | GPU | operation | rows | count | vs deterministic Cake | vs non-Cake SOTA | vs previous non-deterministic Cake | |---|---|---|---|---|---|---| | GB300 (sm_103a) | dW | owner | 22 | 1.001 (0.992-1.014) | 1.045 | 1.000 | | GB300 (sm_103a) | dW | unsorted | 14 | 1.040 (0.983-1.115) | 1.053 | 1.001 | | GB300 (sm_103a) | ckpt | owner | 18 | 1.000 (0.996-1.005) | 1.187 | 1.000 | | GB300 (sm_103a) | ckpt | unsorted | 14 | 1.008 (0.994-1.027) | 1.188 | 1.001 | | GB300 (sm_103a) | ckpt ES | owner | 18 | 1.001 (0.991-1.008) | 1.185 | 1.001 | | GB300 (sm_103a) | ckpt ES | unsorted | 14 | 1.013 (0.998-1.038) | 1.187 | 1.000 | | B200 (sm_100a) | dW | owner | 22 | 1.004 (0.990-1.035) | 1.058 | 1.003 | | B200 (sm_100a) | dW | unsorted | 14 | 1.058 (1.008-1.164) | 1.062 | 1.000 | | B200 (sm_100a) | ckpt | owner | 18 | 1.004 (0.992-1.016) | 1.237 | 1.003 | | B200 (sm_100a) | ckpt | unsorted | 14 | 1.014 (1.004-1.039) | 1.254 | 1.003 | | B200 (sm_100a) | ckpt ES | owner | 18 | 1.004 (0.989-1.012) | 1.228 | 1.005 | | B200 (sm_100a) | ckpt ES | unsorted | 14 | 1.018 (1.003-1.040) | 1.255 | 1.002 | ### GB300 (sm_103a) NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4, driver 580.159.03, torch 2.14.1+cu130 (CUDA 13.0); CUPTI GPU time of one call, cold L2, eager; median of paired groups in rotated order | GPU / architecture | workload / shape / mode | operation | deterministic Cake latency (ms) | non-Cake SOTA latency (ms) | previous non-deterministic Cake latency (ms) | new non-deterministic Cake latency (ms) | speedup vs deterministic Cake | speedup vs non-Cake SOTA | speedup vs previous non-deterministic Cake | |---|---|---|---|---|---|---|---|---|---| | sm_103a | gate_up E4 M112576 bal | dW | 3.1664 | 3.3368 c | 3.1481 | 3.1727 | 0.998 [0.991, 1.008] | 1.051 [1.033, 1.063] | 0.991 [0.978, 1.009] | | sm_103a | gate_up E4 M112576 skew | dW | 3.2938 | 3.3325 c | 3.2967 | 3.2940 | 1.000 [0.995, 1.002] | 1.012 [1.006, 1.015] | 1.000 [0.989, 1.006] | | sm_103a | gate_up E4 M112576 1-empty | dW | 3.2723 | 3.2834 c | 3.2569 | 3.2754 | 1.000 [0.998, 1.002] | 1.004 [1.001, 1.011] | 0.997 [0.987, 1.001] | | sm_103a | gate_up E4 M112576 bal fp32 | dW | 3.2252 | 3.3082 c | 3.2206 | 3.2008 | 1.004 [0.998, 1.012] | 1.032 [1.024, 1.054] | 1.003 [0.987, 1.013] | | sm_103a | gate_up E4 M136064 bal | dW | 3.8339 | 4.0270 c | 3.8319 | 3.8366 | 0.999 [0.995, 1.006] | 1.051 [1.044, 1.052] | 0.998 [0.994, 1.003] | | sm_103a | gate_up E4 M136064 skew | dW | 3.9911 | 4.0418 c | 3.9970 | 3.9970 | 0.998 [0.996, 1.001] | 1.012 [1.007, 1.015] | 1.000 [0.995, 1.007] | | sm_103a | gate_up E4 M136064 1-empty | dW | 3.9701 | 4.0361 c | 3.9287 | 3.9598 | 1.004 [0.997, 1.032] | 1.019 [0.997, 1.048] | 0.996 [0.975, 1.030] | | sm_103a | gate_up E4 M136064 bal fp32 | dW | 3.9346 | 4.0701 c | 3.9244 | 3.9356 | 1.003 [0.992, 1.006] | 1.035 [1.032, 1.051] | 1.000 [0.992, 1.001] | | sm_103a | gate_up E8 M131072 bal | dW | 3.7203 | 3.9216 c | 3.7045 | 3.7105 | 1.003 [0.993, 1.006] | 1.057 [1.048, 1.060] | 0.998 [0.990, 1.005] | | sm_103a | gate_up E16 M131072 bal | dW | 3.6980 | 3.8986 t | 3.7082 | 3.7095 | 0.997 [0.996, 0.998] | 1.051 [1.051, 1.054] | 1.000 [0.995, 1.001] | | sm_103a | gate_up E32 M131072 bal | dW | 3.6352 | 3.7648 t | 3.6778 | 3.6531 | 0.995 [0.986, 1.000] | 1.032 [1.023, 1.040] | 1.004 [1.002, 1.020] | | sm_103a | gate_up E4 M112576 skew asc | dW | 3.2793 | 3.3505 c | 3.3029 | 3.2955 | 0.994 [0.988, 1.003] | 1.016 [1.007, 1.023] | 1.001 [0.996, 1.005] | | sm_103a | gate_up E4 M112576 1-empty asc | dW | 3.2443 | 3.3886 c | 3.2775 | 3.2818 | 0.990 [0.984, 1.027] | 1.039 [1.018, 1.068] | 1.000 [0.991, 1.031] | | sm_103a | gate_up E4 M136064 skew asc | dW | 3.9035 | 4.0407 c | 3.9701 | 3.9712 | 0.983 [0.980, 0.996] | 1.017 [1.012, 1.030] | 1.000 [0.998, 1.008] | | sm_103a | gate_up E4 M136064 1-empty asc | dW | 4.0163 | 4.1055 c | 4.0292 | 4.0460 | 0.994 [0.986, 1.007] | 1.014 [1.009, 1.030] | 0.995 [0.990, 1.007] | | sm_103a | gate_up E8 M131072 random | dW | 3.8180 | 3.8463 c | 3.6960 | 3.6891 | 1.034 [1.033, 1.039] | 1.042 [1.031, 1.071] | 1.002 [0.994, 1.006] | | sm_103a | gate_up E16 M131072 random | dW | 3.9325 | 4.0392 c | 3.7553 | 3.7601 | 1.045 [1.034, 1.048] | 1.074 [1.071, 1.076] | 0.999 [0.997, 1.000] | | sm_103a | gate_up E32 M131072 random | dW | 3.8478 | 3.9913 t | 3.7883 | 3.7896 | 1.016 [0.997, 1.025] | 1.055 [1.040, 1.069] | 1.001 [0.999, 1.009] | | sm_103a | down E4 M112576 bal | dW | 1.5830 | 1.6822 c | 1.5925 | 1.5775 | 1.004 [0.989, 1.009] | 1.059 [1.044, 1.083] | 1.008 [0.998, 1.034] | | sm_103a | down E4 M112576 skew | dW | 1.6370 | 1.6920 c | 1.6294 | 1.6184 | 1.010 [1.006, 1.016] | 1.044 [1.039, 1.049] | 1.007 [1.001, 1.010] | | sm_103a | down E4 M112576 1-empty | dW | 1.5785 | 1.6550 c | 1.5864 | 1.5867 | 0.994 [0.982, 1.011] | 1.043 [1.035, 1.063] | 1.000 [0.989, 1.008] | | sm_103a | down E4 M112576 bal fp32 | dW | 1.5961 | 1.6638 c | 1.5889 | 1.5916 | 1.003 [0.995, 1.011] | 1.046 [1.041, 1.052] | 1.001 [0.996, 1.004] | | sm_103a | down E4 M136064 bal | dW | 1.9327 | 2.0381 c | 1.9177 | 1.9244 | 1.004 [0.997, 1.011] | 1.063 [1.051, 1.066] | 0.999 [0.990, 1.001] | | sm_103a | down E4 M136064 skew | dW | 2.0012 | 2.0400 c | 1.9911 | 1.9730 | 1.014 [1.012, 1.015] | 1.034 [1.029, 1.036] | 1.009 [0.983, 1.011] | | sm_103a | down E4 M136064 1-empty | dW | 1.9253 | 2.0177 c | 1.9222 | 1.9195 | 1.003 [1.001, 1.005] | 1.051 [1.046, 1.057] | 1.001 [0.997, 1.018] | | sm_103a | down E4 M136064 bal fp32 | dW | 1.9469 | 2.0505 c | 1.9505 | 1.9476 | 0.992 [0.986, 1.001] | 1.053 [1.040, 1.055] | 1.002 [0.989, 1.006] | | sm_103a | down E8 M131072 bal | dW | 1.8187 | 1.9639 c | 1.8183 | 1.8148 | 1.002 [0.999, 1.005] | 1.082 [1.072, 1.086] | 1.002 [0.991, 1.008] | | sm_103a | down E16 M131072 bal | dW | 1.8721 | 2.0858 c | 1.8690 | 1.8699 | 1.001 [1.000, 1.016] | 1.116 [1.113, 1.120] | 0.999 [0.980, 1.008] | | sm_103a | down E32 M131072 bal | dW | 1.8647 | 1.9556 t | 1.8512 | 1.8786 | 0.992 [0.990, 1.011] | 1.045 [1.037, 1.059] | 0.989 [0.977, 1.001] | | sm_103a | down E4 M112576 skew asc | dW | 1.6412 | 1.6908 c | 1.6032 | 1.5861 | 1.034 [1.027, 1.045] | 1.063 [1.056, 1.072] | 1.010 [1.006, 1.016] | | sm_103a | down E4 M112576 1-empty asc | dW | 1.7711 | 1.7099 c | 1.6047 | 1.6029 | 1.107 [1.091, 1.146] | 1.069 [1.053, 1.100] | 1.006 [0.974, 1.041] | | sm_103a | down E4 M136064 skew asc | dW | 2.0348 | 2.0424 c | 1.9930 | 1.9711 | 1.033 [1.023, 1.035] | 1.036 [1.027, 1.043] | 1.008 [1.002, 1.015] | | sm_103a | down E4 M136064 1-empty asc | dW | 2.1324 | 2.0732 c | 1.9167 | 1.9147 | 1.115 [1.102, 1.118] | 1.084 [1.075, 1.096] | 1.002 [1.001, 1.003] | | sm_103a | down E8 M131072 random | dW | 2.0384 | 1.9910 c | 1.8849 | 1.8844 | 1.082 [1.074, 1.083] | 1.057 [1.055, 1.058] | 1.000 [0.999, 1.003] | | sm_103a | down E16 M131072 random | dW | 2.1025 | 2.0670 c | 1.8997 | 1.9089 | 1.101 [1.076, 1.122] | 1.083 [1.064, 1.094] | 0.997 [0.978, 1.013] | | sm_103a | down E32 M131072 random | dW | 1.9379 | 2.0551 t | 1.8589 | 1.8630 | 1.040 [1.029, 1.058] | 1.106 [1.093, 1.112] | 1.000 [0.985, 1.025] | | sm_103a | gate_up E4 M112576 bal | ckpt | 12.6579 | 15.5120 t | 12.6585 | 12.6513 | 1.001 [0.997, 1.007] | 1.229 [1.219, 1.236] | 1.001 [0.998, 1.002] | | sm_103a | gate_up E4 M112576 skew | ckpt | 12.8493 | 15.7940 t | 12.9281 | 12.8874 | 0.998 [0.994, 1.000] | 1.228 [1.210, 1.232] | 1.005 [0.999, 1.013] | | sm_103a | gate_up E4 M112576 1-empty | ckpt | 12.6127 | 15.5400 t | 12.5693 | 12.6125 | 0.999 [0.993, 1.003] | 1.232 [1.221, 1.238] | 0.996 [0.987, 1.008] | | sm_103a | gate_up E4 M136064 bal | ckpt | 15.2573 | 18.8586 t | 15.2798 | 15.2624 | 1.000 [0.988, 1.003] | 1.237 [1.222, 1.246] | 1.000 [0.990, 1.010] | | sm_103a | gate_up E4 M136064 skew | ckpt | 15.5979 | 19.2668 t | 15.6169 | 15.6181 | 0.996 [0.994, 1.003] | 1.234 [1.218, 1.257] | 1.000 [0.996, 1.003] | | sm_103a | gate_up E4 M136064 1-empty | ckpt | 15.4011 | 18.8718 t | 15.3614 | 15.3635 | 1.002 [0.993, 1.006] | 1.229 [1.225, 1.232] | 1.000 [0.987, 1.003] | | sm_103a | gate_up E8 M131072 bal | ckpt | 14.7011 | 17.5611 t | 14.7309 | 14.7249 | 0.998 [0.996, 1.001] | 1.191 [1.187, 1.200] | 1.002 [0.996, 1.004] | | sm_103a | gate_up E16 M131072 bal | ckpt | 15.0497 | 16.6615 t | 15.0024 | 14.9718 | 1.003 [0.994, 1.011] | 1.113 [1.108, 1.118] | 1.001 [0.993, 1.011] | | sm_103a | gate_up E32 M131072 bal | ckpt | 14.7770 | 15.2884 t | 14.8059 | 14.8132 | 0.997 [0.980, 0.999] | 1.030 [1.018, 1.036] | 1.000 [0.988, 1.003] | | sm_103a | gate_up E4 M112576 skew asc | ckpt | 12.6952 | 15.7037 t | 12.7622 | 12.7579 | 0.995 [0.994, 1.001] | 1.232 [1.226, 1.237] | 1.001 [0.996, 1.011] | | sm_103a | gate_up E4 M112576 1-empty asc | ckpt | 12.9375 | 16.1023 t | 12.9064 | 12.8959 | 1.000 [0.999, 1.012] | 1.246 [1.227, 1.261] | 1.001 [0.988, 1.006] | | sm_103a | gate_up E4 M136064 skew asc | ckpt | 15.3430 | 18.9009 t | 15.3727 | 15.3180 | 1.003 [0.996, 1.004] | 1.232 [1.224, 1.246] | 1.003 [0.994, 1.008] | | sm_103a | gate_up E4 M136064 1-empty asc | ckpt | 15.4231 | 19.3197 t | 15.5300 | 15.5385 | 0.994 [0.990, 0.997] | 1.244 [1.239, 1.251] | 1.000 [0.993, 1.006] | | sm_103a | gate_up E8 M131072 random | ckpt | 15.1947 | 18.3366 t | 15.0414 | 15.0282 | 1.011 [1.007, 1.016] | 1.220 [1.202, 1.227] | 1.002 [0.997, 1.005] | | sm_103a | gate_up E16 M131072 random | ckpt | 15.0851 | 16.8810 t | 14.9736 | 14.9818 | 1.008 [1.006, 1.014] | 1.127 [1.116, 1.138] | 1.000 [0.998, 1.005] | | sm_103a | gate_up E32 M131072 random | ckpt | 15.4443 | 16.1078 t | 15.4255 | 15.4361 | 1.001 [0.999, 1.017] | 1.044 [1.039, 1.051] | 0.999 [0.992, 1.005] | | sm_103a | down E4 M112576 bal | ckpt | 6.2910 | 7.5948 t | 6.2791 | 6.2707 | 1.005 [0.996, 1.014] | 1.211 [1.195, 1.229] | 1.005 [0.983, 1.023] | | sm_103a | down E4 M112576 skew | ckpt | 6.4125 | 7.9380 t | 6.4145 | 6.3908 | 1.004 [0.989, 1.013] | 1.245 [1.226, 1.267] | 1.004 [0.992, 1.014] | | sm_103a | down E4 M112576 1-empty | ckpt | 6.2826 | 7.7901 t | 6.2906 | 6.2811 | 1.002 [0.996, 1.012] | 1.243 [1.235, 1.248] | 1.001 [0.996, 1.004] | | sm_103a | down E4 M136064 bal | ckpt | 7.6785 | 9.3028 t | 7.6260 | 7.6777 | 0.999 [0.993, 1.003] | 1.214 [1.209, 1.225] | 0.995 [0.991, 1.003] | | sm_103a | down E4 M136064 skew | ckpt | 7.8223 | 9.7387 t | 7.7944 | 7.7981 | 1.003 [0.997, 1.009] | 1.248 [1.241, 1.257] | 1.000 [0.995, 1.011] | | sm_103a | down E4 M136064 1-empty | ckpt | 7.6135 | 9.5114 t | 7.5969 | 7.5951 | 1.003 [0.996, 1.018] | 1.253 [1.244, 1.268] | 1.001 [0.994, 1.010] | | sm_103a | down E8 M131072 bal | ckpt | 7.4072 | 8.4058 t | 7.4080 | 7.4073 | 1.000 [0.998, 1.000] | 1.135 [1.125, 1.146] | 1.000 [0.994, 1.001] | | sm_103a | down E16 M131072 bal | ckpt | 7.4652 | 8.2022 t | 7.4901 | 7.4943 | 0.996 [0.980, 1.004] | 1.092 [1.077, 1.105] | 0.999 [0.982, 1.008] | | sm_103a | down E32 M131072 bal | ckpt | 7.4224 | 7.7417 t | 7.4220 | 7.4310 | 0.999 [0.990, 1.007] | 1.041 [1.034, 1.050] | 0.998 [0.994, 1.006] | | sm_103a | down E4 M112576 skew asc | ckpt | 6.4784 | 7.9554 t | 6.4527 | 6.4242 | 1.008 [0.999, 1.010] | 1.236 [1.220, 1.253] | 1.004 [0.999, 1.005] | | sm_103a | down E4 M112576 1-empty asc | ckpt | 6.4520 | 7.8480 t | 6.2979 | 6.2790 | 1.026 [1.019, 1.031] | 1.249 [1.235, 1.261] | 1.002 [0.997, 1.006] | | sm_103a | down E4 M136064 skew asc | ckpt | 7.7877 | 9.5875 t | 7.7691 | 7.7600 | 1.005 [1.000, 1.017] | 1.236 [1.215, 1.258] | 1.002 [0.995, 1.016] | | sm_103a | down E4 M136064 1-empty asc | ckpt | 7.9959 | 9.8420 t | 7.7974 | 7.7823 | 1.027 [1.025, 1.041] | 1.268 [1.252, 1.285] | 1.001 [1.000, 1.012] | | sm_103a | down E8 M131072 random | ckpt | 7.5057 | 8.5003 t | 7.3919 | 7.3713 | 1.016 [1.013, 1.027] | 1.146 [1.142, 1.161] | 1.000 [0.972, 1.005] | | sm_103a | down E16 M131072 random | ckpt | 7.6812 | 8.3222 t | 7.5756 | 7.5734 | 1.014 [1.012, 1.016] | 1.099 [1.097, 1.103] | 1.000 [0.999, 1.002] | | sm_103a | down E32 M131072 random | ckpt | 7.8150 | 8.3989 t | 7.7617 | 7.7680 | 1.004 [1.002, 1.011] | 1.080 [1.065, 1.090] | 1.000 [0.988, 1.006] | | sm_103a | gate_up E4 M112576 bal | ckpt ES | 9.3663 | 11.4013 t | 9.3264 | 9.3447 | 1.006 [0.991, 1.008] | 1.218 [1.204, 1.234] | 1.000 [0.992, 1.004] | | sm_103a | gate_up E4 M112576 skew | ckpt ES | 9.5468 | 11.5057 t | 9.5965 | 9.5800 | 0.997 [0.990, 1.003] | 1.202 [1.188, 1.215] | 1.001 [0.989, 1.010] | | sm_103a | gate_up E4 M112576 1-empty | ckpt ES | 9.6244 | 11.7650 t | 9.6485 | 9.6418 | 0.999 [0.995, 1.007] | 1.218 [1.210, 1.227] | 1.001 [0.990, 1.009] | | sm_103a | gate_up E4 M136064 bal | ckpt ES | 11.3815 | 13.8066 t | 11.3579 | 11.3533 | 1.000 [0.999, 1.005] | 1.215 [1.209, 1.222] | 1.001 [0.994, 1.006] | | sm_103a | gate_up E4 M136064 skew | ckpt ES | 11.7204 | 14.3336 t | 11.6415 | 11.6077 | 0.999 [0.995, 1.032] | 1.215 [1.195, 1.258] | 1.000 [0.989, 1.013] | | sm_103a | gate_up E4 M136064 1-empty | ckpt ES | 11.6324 | 14.3628 t | 11.6246 | 11.5513 | 1.008 [0.988, 1.022] | 1.239 [1.223, 1.267] | 1.004 [0.996, 1.023] | | sm_103a | gate_up E8 M131072 bal | ckpt ES | 10.7533 | 12.8192 t | 10.9026 | 10.8549 | 1.002 [0.979, 1.005] | 1.179 [1.146, 1.207] | 1.005 [0.982, 1.023] | | sm_103a | gate_up E16 M131072 bal | ckpt ES | 11.0500 | 11.9749 t | 11.0517 | 11.0566 | 1.001 [0.998, 1.008] | 1.082 [1.081, 1.094] | 1.001 [0.998, 1.009] | | sm_103a | gate_up E32 M131072 bal | ckpt ES | 11.3156 | 12.0762 t | 11.3330 | 11.3186 | 1.000 [0.993, 1.006] | 1.069 [1.059, 1.070] | 1.001 [0.994, 1.006] | | sm_103a | gate_up E4 M112576 skew asc | ckpt ES | 9.4904 | 11.5859 t | 9.4975 | 9.4971 | 1.000 [0.976, 1.007] | 1.212 [1.208, 1.252] | 1.001 [0.990, 1.007] | | sm_103a | gate_up E4 M112576 1-empty asc | ckpt ES | 9.5237 | 11.6802 t | 9.5434 | 9.5470 | 0.998 [0.988, 1.006] | 1.220 [1.213, 1.232] | 0.998 [0.982, 1.014] | | sm_103a | gate_up E4 M136064 skew asc | ckpt ES | 11.7361 | 14.3327 t | 11.7139 | 11.6922 | 1.002 [0.986, 1.007] | 1.224 [1.215, 1.233] | 1.002 [0.988, 1.007] | | sm_103a | gate_up E4 M136064 1-empty asc | ckpt ES | 11.5073 | 14.1531 t | 11.5036 | 11.4805 | 0.998 [0.993, 1.004] | 1.231 [1.227, 1.238] | 1.002 [0.994, 1.009] | | sm_103a | gate_up E8 M131072 random | ckpt ES | 11.2200 | 13.2241 t | 11.0895 | 11.1049 | 1.009 [1.009, 1.013] | 1.190 [1.185, 1.201] | 0.998 [0.994, 1.000] | | sm_103a | gate_up E16 M131072 random | ckpt ES | 11.5895 | 12.8866 t | 11.3670 | 11.3570 | 1.019 [1.015, 1.027] | 1.134 [1.131, 1.139] | 1.001 [0.998, 1.003] | | sm_103a | gate_up E32 M131072 random | ckpt ES | 11.4684 | 11.8086 t | 11.4087 | 11.4100 | 1.004 [0.997, 1.008] | 1.034 [1.024, 1.047] | 0.999 [0.996, 1.002] | | sm_103a | down E4 M112576 bal | ckpt ES | 4.6835 | 5.6540 t | 4.6898 | 4.6771 | 1.001 [0.984, 1.030] | 1.209 [1.200, 1.248] | 1.002 [0.980, 1.032] | | sm_103a | down E4 M112576 skew | ckpt ES | 4.6944 | 5.8074 t | 4.7272 | 4.7181 | 0.991 [0.976, 1.003] | 1.230 [1.208, 1.238] | 1.002 [0.984, 1.018] | | sm_103a | down E4 M112576 1-empty | ckpt ES | 4.7842 | 5.9623 t | 4.7875 | 4.7729 | 1.002 [0.995, 1.011] | 1.251 [1.237, 1.258] | 1.002 [0.989, 1.008] | | sm_103a | down E4 M136064 bal | ckpt ES | 5.6247 | 6.8500 t | 5.6192 | 5.6219 | 1.000 [0.997, 1.002] | 1.219 [1.216, 1.245] | 0.999 [0.995, 1.004] | | sm_103a | down E4 M136064 skew | ckpt ES | 5.7158 | 7.0419 t | 5.7068 | 5.7121 | 1.002 [0.994, 1.005] | 1.235 [1.229, 1.255] | 1.000 [0.994, 1.010] | | sm_103a | down E4 M136064 1-empty | ckpt ES | 5.8110 | 7.2810 t | 5.8096 | 5.8082 | 1.001 [0.999, 1.004] | 1.255 [1.243, 1.263] | 1.001 [0.996, 1.013] | | sm_103a | down E8 M131072 bal | ckpt ES | 5.4216 | 6.3190 t | 5.4210 | 5.4276 | 0.999 [0.984, 1.003] | 1.163 [1.157, 1.168] | 1.000 [0.996, 1.002] | | sm_103a | down E16 M131072 bal | ckpt ES | 5.4970 | 6.0144 t | 5.5158 | 5.4857 | 1.001 [0.980, 1.020] | 1.096 [1.087, 1.112] | 1.008 [0.981, 1.018] | | sm_103a | down E32 M131072 bal | ckpt ES | 5.6244 | 6.0064 t | 5.6105 | 5.6179 | 1.000 [0.996, 1.007] | 1.069 [1.062, 1.083] | 0.999 [0.994, 1.005] | | sm_103a | down E4 M112576 skew asc | ckpt ES | 4.7019 | 5.8377 t | 4.6899 | 4.6960 | 1.011 [0.983, 1.021] | 1.243 [1.229, 1.265] | 1.004 [0.987, 1.018] | | sm_103a | down E4 M112576 1-empty asc | ckpt ES | 4.9442 | 5.9699 t | 4.7702 | 4.7619 | 1.038 [1.031, 1.044] | 1.253 [1.248, 1.273] | 1.003 [0.996, 1.009] | | sm_103a | down E4 M136064 skew asc | ckpt ES | 5.7645 | 7.1457 t | 5.7245 | 5.7020 | 1.012 [1.010, 1.013] | 1.253 [1.243, 1.259] | 1.004 [1.003, 1.006] | | sm_103a | down E4 M136064 1-empty asc | ckpt ES | 5.8676 | 7.1629 t | 5.6824 | 5.7041 | 1.030 [1.028, 1.033] | 1.256 [1.241, 1.262] | 0.996 [0.995, 1.000] | | sm_103a | down E8 M131072 random | ckpt ES | 5.7750 | 6.6888 t | 5.6084 | 5.6299 | 1.026 [1.022, 1.036] | 1.190 [1.182, 1.198] | 0.998 [0.993, 1.001] | | sm_103a | down E16 M131072 random | ckpt ES | 5.6758 | 6.2531 t | 5.5402 | 5.5429 | 1.023 [1.019, 1.037] | 1.126 [1.120, 1.144] | 0.998 [0.998, 1.019] | | sm_103a | down E32 M131072 random | ckpt ES | 5.7573 | 6.1208 t | 5.6975 | 5.6771 | 1.013 [1.004, 1.028] | 1.078 [1.067, 1.086] | 1.000 [0.999, 1.011] | ### B200 (sm_100a) NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.4, driver 580.82.07, torch 2.14.1+cu130 (CUDA 13.0); CUPTI GPU time of one call, cold L2, eager; median of paired groups in rotated order | GPU / architecture | workload / shape / mode | operation | deterministic Cake latency (ms) | non-Cake SOTA latency (ms) | previous non-deterministic Cake latency (ms) | new non-deterministic Cake latency (ms) | speedup vs deterministic Cake | speedup vs non-Cake SOTA | speedup vs previous non-deterministic Cake | |---|---|---|---|---|---|---|---|---|---| | sm_100a | gate_up E4 M112576 bal | dW | 3.9976 | 4.2457 c | 3.9870 | 3.9885 | 1.003 [0.961, 1.025] | 1.066 [1.041, 1.086] | 1.010 [0.978, 1.021] | | sm_100a | gate_up E4 M112576 skew | dW | 4.1702 | 4.2587 c | 4.1527 | 4.1624 | 1.002 [0.984, 1.003] | 1.022 [1.004, 1.027] | 0.992 [0.979, 1.001] | | sm_100a | gate_up E4 M112576 1-empty | dW | 4.0958 | 4.2650 c | 4.0841 | 4.0885 | 1.002 [0.984, 1.021] | 1.043 [1.043, 1.047] | 1.000 [0.981, 1.018] | | sm_100a | gate_up E4 M112576 bal fp32 | dW | 4.0618 | 4.2551 c | 4.0615 | 4.0599 | 1.001 [0.958, 1.022] | 1.048 [1.012, 1.064] | 1.000 [0.979, 1.022] | | sm_100a | gate_up E4 M136064 bal | dW | 4.8575 | 5.1252 c | 4.8972 | 4.8423 | 1.001 [0.984, 1.031] | 1.058 [1.044, 1.082] | 1.007 [0.997, 1.022] | | sm_100a | gate_up E4 M136064 skew | dW | 5.0294 | 5.1806 c | 5.0146 | 5.0195 | 1.003 [1.002, 1.039] | 1.033 [1.016, 1.052] | 1.000 [0.982, 1.001] | | sm_100a | gate_up E4 M136064 1-empty | dW | 4.9063 | 5.1470 c | 4.8889 | 4.9223 | 1.002 [0.984, 1.022] | 1.049 [1.036, 1.064] | 0.993 [0.986, 1.014] | | sm_100a | gate_up E4 M136064 bal fp32 | dW | 5.0028 | 5.1651 c | 5.0181 | 5.0179 | 0.997 [0.979, 1.023] | 1.029 [1.023, 1.070] | 1.000 [0.986, 1.024] | | sm_100a | gate_up E8 M131072 bal | dW | 4.6836 | 4.9684 c | 4.6894 | 4.6893 | 0.999 [0.981, 1.006] | 1.060 [1.044, 1.068] | 1.000 [0.974, 1.008] | | sm_100a | gate_up E16 M131072 bal | dW | 4.5996 | 5.0050 t | 4.6180 | 4.6719 | 0.990 [0.980, 1.002] | 1.080 [1.056, 1.103] | 1.002 [0.981, 1.007] | | sm_100a | gate_up E32 M131072 bal | dW | 4.5906 | 4.7800 t | 4.6094 | 4.6398 | 0.996 [0.986, 1.026] | 1.027 [1.020, 1.043] | 0.999 [0.988, 1.017] | | sm_100a | gate_up E4 M112576 skew asc | dW | 4.2118 | 4.2625 c | 4.1668 | 4.1967 | 1.008 [0.993, 1.011] | 1.016 [1.001, 1.022] | 0.981 [0.980, 1.000] | | sm_100a | gate_up E4 M112576 1-empty asc | dW | 4.2738 | 4.2947 c | 4.0839 | 4.0836 | 1.047 [1.045, 1.062] | 1.052 [1.050, 1.081] | 1.000 [0.984, 1.014] | | sm_100a | gate_up E4 M136064 skew asc | dW | 4.9440 | 5.1221 c | 4.9421 | 4.8964 | 1.010 [0.989, 1.011] | 1.045 [1.022, 1.062] | 1.000 [0.978, 1.022] | | sm_100a | gate_up E4 M136064 1-empty asc | dW | 5.2634 | 5.1861 c | 5.0200 | 5.0272 | 1.039 [1.027, 1.051] | 1.031 [1.013, 1.033] | 0.999 [0.981, 1.019] | | sm_100a | gate_up E8 M131072 random | dW | 4.8594 | 4.9952 c | 4.6333 | 4.6324 | 1.038 [0.991, 1.061] | 1.063 [1.045, 1.081] | 0.992 [0.970, 1.031] | | sm_100a | gate_up E16 M131072 random | dW | 4.8748 | 5.0676 c | 4.6498 | 4.6453 | 1.049 [1.014, 1.059] | 1.082 [1.067, 1.104] | 1.001 [0.979, 1.020] | | sm_100a | gate_up E32 M131072 random | dW | 4.9194 | 5.0602 t | 4.7956 | 4.7932 | 1.026 [1.007, 1.048] | 1.056 [1.019, 1.059] | 1.000 [1.000, 1.017] | | sm_100a | down E4 M112576 bal | dW | 2.0133 | 2.1235 c | 1.9720 | 2.0137 | 1.000 [0.999, 1.021] | 1.061 [1.048, 1.083] | 0.989 [0.978, 1.021] | | sm_100a | down E4 M112576 skew | dW | 2.0484 | 2.1367 c | 2.0358 | 2.0075 | 1.017 [1.007, 1.029] | 1.065 [1.062, 1.069] | 1.015 [0.999, 1.028] | | sm_100a | down E4 M112576 1-empty | dW | 1.9687 | 2.1103 c | 1.9709 | 1.9769 | 1.004 [0.953, 1.005] | 1.059 [1.030, 1.088] | 0.998 [0.963, 1.016] | | sm_100a | down E4 M112576 bal fp32 | dW | 2.0121 | 2.1209 c | 2.0011 | 2.0004 | 0.997 [0.995, 1.015] | 1.060 [1.040, 1.063] | 1.000 [0.985, 1.001] | | sm_100a | down E4 M136064 bal | dW | 2.4190 | 2.5454 c | 2.4017 | 2.4031 | 0.998 [0.981, 1.019] | 1.058 [1.034, 1.069] | 0.999 [0.981, 1.001] | | sm_100a | down E4 M136064 skew | dW | 2.4836 | 2.6037 c | 2.4814 | 2.4187 | 1.035 [0.986, 1.048] | 1.076 [1.020, 1.085] | 1.016 [0.986, 1.048] | | sm_100a | down E4 M136064 1-empty | dW | 2.4237 | 2.5360 c | 2.4593 | 2.3796 | 1.002 [0.957, 1.047] | 1.066 [1.063, 1.073] | 1.018 [0.998, 1.043] | | sm_100a | down E4 M136064 bal fp32 | dW | 2.4582 | 2.5711 c | 2.4668 | 2.4403 | 1.018 [0.981, 1.019] | 1.055 [1.043, 1.065] | 1.012 [1.000, 1.022] | | sm_100a | down E8 M131072 bal | dW | 2.3249 | 2.4791 c | 2.3405 | 2.3302 | 0.998 [0.964, 1.052] | 1.064 [1.028, 1.112] | 1.001 [0.967, 1.035] | | sm_100a | down E16 M131072 bal | dW | 2.3628 | 2.6612 c | 2.3580 | 2.3640 | 1.000 [0.981, 1.021] | 1.126 [1.073, 1.151] | 0.999 [0.961, 1.017] | | sm_100a | down E32 M131072 bal | dW | 2.3488 | 2.5136 t | 2.3451 | 2.3228 | 1.016 [0.981, 1.026] | 1.082 [1.045, 1.094] | 1.010 [0.998, 1.037] | | sm_100a | down E4 M112576 skew asc | dW | 2.1135 | 2.1706 c | 2.0273 | 2.0313 | 1.041 [1.031, 1.064] | 1.069 [1.039, 1.086] | 0.998 [0.982, 1.014] | | sm_100a | down E4 M112576 1-empty asc | dW | 2.2088 | 2.1334 c | 1.9797 | 1.9872 | 1.115 [1.096, 1.154] | 1.075 [1.063, 1.087] | 1.000 [0.982, 1.015] | | sm_100a | down E4 M136064 skew asc | dW | 2.6051 | 2.6086 c | 2.4670 | 2.4423 | 1.058 [1.056, 1.077] | 1.068 [1.054, 1.079] | 1.007 [0.994, 1.027] | | sm_100a | down E4 M136064 1-empty asc | dW | 2.7647 | 2.6315 c | 2.4211 | 2.3734 | 1.164 [1.121, 1.184] | 1.092 [1.069, 1.111] | 1.005 [0.964, 1.038] | | sm_100a | down E8 M131072 random | dW | 2.5173 | 2.4921 c | 2.3593 | 2.3577 | 1.067 [1.044, 1.110] | 1.058 [1.051, 1.076] | 0.999 [0.979, 1.023] | | sm_100a | down E16 M131072 random | dW | 2.6664 | 2.6173 c | 2.4221 | 2.3987 | 1.119 [1.062, 1.171] | 1.091 [1.041, 1.126] | 1.020 [0.961, 1.021] | | sm_100a | down E32 M131072 random | dW | 2.5157 | 2.5734 t | 2.4102 | 2.4059 | 1.046 [1.005, 1.079] | 1.074 [1.039, 1.111] | 1.001 [0.953, 1.035] | | sm_100a | gate_up E4 M112576 bal | ckpt | 15.8466 | 20.3868 t | 15.6836 | 15.6994 | 1.006 [0.992, 1.022] | 1.286 [1.252, 1.333] | 1.003 [0.978, 1.030] | | sm_100a | gate_up E4 M112576 skew | ckpt | 16.1417 | 20.7842 t | 16.1533 | 16.1237 | 1.002 [0.983, 1.010] | 1.278 [1.255, 1.371] | 0.997 [0.980, 1.016] | | sm_100a | gate_up E4 M112576 1-empty | ckpt | 15.7466 | 20.6752 t | 15.8363 | 15.7299 | 1.003 [0.990, 1.015] | 1.316 [1.289, 1.363] | 1.009 [0.992, 1.018] | | sm_100a | gate_up E4 M136064 bal | ckpt | 19.0178 | 24.2402 t | 18.9531 | 19.0144 | 1.002 [0.982, 1.010] | 1.275 [1.268, 1.309] | 0.999 [0.978, 1.024] | | sm_100a | gate_up E4 M136064 skew | ckpt | 19.4603 | 26.1626 t | 19.4689 | 19.4465 | 0.999 [0.986, 1.019] | 1.333 [1.291, 1.360] | 0.995 [0.986, 1.017] | | sm_100a | gate_up E4 M136064 1-empty | ckpt | 19.1663 | 25.0398 t | 19.2085 | 19.1684 | 0.999 [0.985, 1.009] | 1.308 [1.300, 1.318] | 1.001 [0.995, 1.007] | | sm_100a | gate_up E8 M131072 bal | ckpt | 18.4489 | 22.8121 t | 18.3309 | 18.3292 | 1.006 [0.978, 1.019] | 1.243 [1.219, 1.290] | 0.999 [0.987, 1.020] | | sm_100a | gate_up E16 M131072 bal | ckpt | 18.7873 | 21.3988 t | 18.7435 | 18.7504 | 1.003 [0.992, 1.009] | 1.139 [1.113, 1.190] | 0.999 [0.976, 1.015] | | sm_100a | gate_up E32 M131072 bal | ckpt | 18.7197 | 19.7054 t | 18.7163 | 18.7545 | 1.001 [0.997, 1.013] | 1.051 [1.036, 1.073] | 0.998 [0.985, 1.023] | | sm_100a | gate_up E4 M112576 skew asc | ckpt | 15.8669 | 21.0350 t | 16.0721 | 15.7258 | 1.009 [0.988, 1.027] | 1.332 [1.306, 1.366] | 1.015 [0.996, 1.028] | | sm_100a | gate_up E4 M112576 1-empty asc | ckpt | 16.2500 | 21.8791 t | 16.1957 | 16.1326 | 1.006 [1.002, 1.011] | 1.355 [1.269, 1.369] | 1.001 [0.993, 1.007] | | sm_100a | gate_up E4 M136064 skew asc | ckpt | 19.3207 | 25.9406 t | 19.3153 | 19.1767 | 1.004 [0.992, 1.019] | 1.333 [1.326, 1.383] | 1.002 [0.997, 1.022] | | sm_100a | gate_up E4 M136064 1-empty asc | ckpt | 19.5528 | 25.4783 t | 19.4943 | 19.4801 | 1.004 [0.998, 1.040] | 1.333 [1.278, 1.360] | 1.000 [0.997, 1.020] | | sm_100a | gate_up E8 M131072 random | ckpt | 18.9272 | 24.4975 t | 18.7762 | 18.7787 | 1.008 [1.006, 1.022] | 1.305 [1.249, 1.337] | 1.000 [1.000, 1.018] | | sm_100a | gate_up E16 M131072 random | ckpt | 18.9697 | 22.5539 t | 18.6717 | 18.6835 | 1.013 [1.006, 1.021] | 1.211 [1.176, 1.258] | 1.000 [0.982, 1.026] | | sm_100a | gate_up E32 M131072 random | ckpt | 19.2953 | 20.3427 t | 19.3690 | 19.2416 | 1.004 [0.989, 1.027] | 1.053 [1.047, 1.078] | 1.008 [0.986, 1.023] | | sm_100a | down E4 M112576 bal | ckpt | 7.8800 | 9.8458 t | 7.8234 | 7.8392 | 1.007 [0.999, 1.014] | 1.264 [1.227, 1.286] | 1.002 [0.979, 1.014] | | sm_100a | down E4 M112576 skew | ckpt | 7.9549 | 10.1525 t | 7.9520 | 7.8998 | 1.007 [1.003, 1.007] | 1.282 [1.220, 1.299] | 1.006 [0.954, 1.025] | | sm_100a | down E4 M112576 1-empty | ckpt | 7.7975 | 10.0041 t | 7.7579 | 7.7874 | 1.016 [0.999, 1.020] | 1.288 [1.240, 1.332] | 1.000 [0.982, 1.034] | | sm_100a | down E4 M136064 bal | ckpt | 9.4939 | 12.4195 t | 9.5333 | 9.5240 | 0.992 [0.989, 1.018] | 1.298 [1.283, 1.334] | 1.000 [0.988, 1.023] | | sm_100a | down E4 M136064 skew | ckpt | 9.7430 | 12.5376 t | 9.7316 | 9.6976 | 1.005 [0.987, 1.027] | 1.302 [1.279, 1.345] | 1.005 [0.987, 1.023] | | sm_100a | down E4 M136064 1-empty | ckpt | 9.3983 | 12.2718 t | 9.4316 | 9.4044 | 1.003 [0.987, 1.012] | 1.301 [1.296, 1.322] | 1.003 [0.970, 1.013] | | sm_100a | down E8 M131072 bal | ckpt | 9.2061 | 10.7688 t | 9.2159 | 9.0914 | 1.013 [0.988, 1.021] | 1.182 [1.162, 1.214] | 1.011 [0.985, 1.022] | | sm_100a | down E16 M131072 bal | ckpt | 9.3238 | 10.4223 t | 9.3770 | 9.2016 | 1.005 [0.981, 1.020] | 1.120 [1.092, 1.137] | 1.018 [0.974, 1.039] | | sm_100a | down E32 M131072 bal | ckpt | 9.2812 | 9.8208 t | 9.2888 | 9.2571 | 1.006 [0.984, 1.029] | 1.065 [1.036, 1.084] | 1.006 [0.975, 1.027] | | sm_100a | down E4 M112576 skew asc | ckpt | 8.0466 | 10.2900 t | 7.9529 | 7.9044 | 1.018 [1.004, 1.019] | 1.299 [1.281, 1.374] | 1.006 [0.992, 1.007] | | sm_100a | down E4 M112576 1-empty asc | ckpt | 8.0651 | 10.0984 t | 7.8000 | 7.8037 | 1.035 [1.024, 1.041] | 1.296 [1.287, 1.305] | 0.998 [0.992, 1.001] | | sm_100a | down E4 M136064 skew asc | ckpt | 9.7130 | 12.7132 t | 9.5754 | 9.5794 | 1.016 [0.999, 1.032] | 1.332 [1.289, 1.410] | 0.998 [0.990, 1.014] | | sm_100a | down E4 M136064 1-empty asc | ckpt | 9.9020 | 13.1569 t | 9.6923 | 9.5397 | 1.039 [1.028, 1.053] | 1.370 [1.315, 1.410] | 1.016 [0.990, 1.019] | | sm_100a | down E8 M131072 random | ckpt | 9.4192 | 11.1804 t | 9.1971 | 9.1186 | 1.017 [1.014, 1.040] | 1.221 [1.170, 1.259] | 1.001 [0.982, 1.030] | | sm_100a | down E16 M131072 random | ckpt | 9.5329 | 10.3783 t | 9.3843 | 9.3841 | 1.019 [1.009, 1.035] | 1.106 [1.093, 1.157] | 1.000 [0.992, 1.022] | | sm_100a | down E32 M131072 random | ckpt | 9.7310 | 10.4517 t | 9.6720 | 9.6640 | 1.007 [0.974, 1.016] | 1.071 [1.060, 1.082] | 0.997 [0.980, 1.027] | | sm_100a | gate_up E4 M112576 bal | ckpt ES | 11.8729 | 14.8949 t | 11.9077 | 11.7782 | 1.006 [0.988, 1.018] | 1.269 [1.182, 1.291] | 1.013 [0.999, 1.023] | | sm_100a | gate_up E4 M112576 skew | ckpt ES | 11.8961 | 15.0676 t | 11.9358 | 11.8012 | 1.010 [0.989, 1.024] | 1.268 [1.255, 1.313] | 1.013 [0.980, 1.022] | | sm_100a | gate_up E4 M112576 1-empty | ckpt ES | 12.0663 | 15.2560 t | 11.9654 | 11.9633 | 1.005 [0.971, 1.021] | 1.273 [1.257, 1.287] | 1.000 [0.985, 1.012] | | sm_100a | gate_up E4 M136064 bal | ckpt ES | 14.2658 | 18.3608 t | 14.2694 | 14.2425 | 1.001 [0.995, 1.020] | 1.287 [1.262, 1.311] | 1.005 [1.000, 1.013] | | sm_100a | gate_up E4 M136064 skew | ckpt ES | 14.5425 | 18.2241 t | 14.3239 | 14.4151 | 1.001 [0.984, 1.022] | 1.268 [1.239, 1.311] | 0.992 [0.980, 1.016] | | sm_100a | gate_up E4 M136064 1-empty | ckpt ES | 14.6913 | 18.7803 t | 14.5212 | 14.5658 | 1.011 [0.983, 1.020] | 1.289 [1.260, 1.379] | 1.001 [0.982, 1.015] | | sm_100a | gate_up E8 M131072 bal | ckpt ES | 13.6119 | 16.8884 t | 13.7758 | 13.6072 | 0.999 [0.986, 1.016] | 1.241 [1.226, 1.253] | 1.009 [0.987, 1.017] | | sm_100a | gate_up E16 M131072 bal | ckpt ES | 13.8674 | 15.5597 t | 13.9236 | 13.7783 | 1.001 [0.985, 1.025] | 1.126 [1.102, 1.148] | 1.008 [0.988, 1.020] | | sm_100a | gate_up E32 M131072 bal | ckpt ES | 14.1112 | 15.0496 t | 14.2367 | 14.2698 | 0.989 [0.980, 1.010] | 1.055 [1.039, 1.091] | 0.996 [0.990, 1.007] | | sm_100a | gate_up E4 M112576 skew asc | ckpt ES | 11.9772 | 15.6856 t | 11.9212 | 11.7996 | 1.009 [1.005, 1.035] | 1.334 [1.306, 1.350] | 1.017 [1.000, 1.018] | | sm_100a | gate_up E4 M112576 1-empty asc | ckpt ES | 12.0144 | 15.4493 t | 11.8924 | 11.9302 | 1.010 [0.993, 1.039] | 1.294 [1.267, 1.303] | 1.000 [0.982, 1.016] | | sm_100a | gate_up E4 M136064 skew asc | ckpt ES | 14.7572 | 19.8451 t | 14.6340 | 14.7482 | 1.003 [0.993, 1.027] | 1.342 [1.274, 1.364] | 0.991 [0.983, 0.997] | | sm_100a | gate_up E4 M136064 1-empty asc | ckpt ES | 14.6179 | 19.1221 t | 14.4399 | 14.4515 | 1.010 [1.003, 1.025] | 1.322 [1.297, 1.336] | 0.999 [0.984, 1.014] | | sm_100a | gate_up E8 M131072 random | ckpt ES | 14.0762 | 17.8375 t | 13.8015 | 13.8183 | 1.017 [1.004, 1.030] | 1.281 [1.268, 1.312] | 1.000 [0.977, 1.006] | | sm_100a | gate_up E16 M131072 random | ckpt ES | 14.4723 | 16.7824 t | 14.2621 | 14.2518 | 1.021 [0.987, 1.035] | 1.172 [1.164, 1.205] | 1.002 [0.967, 1.024] | | sm_100a | gate_up E32 M131072 random | ckpt ES | 14.5908 | 15.2804 t | 14.3928 | 14.2734 | 1.013 [0.991, 1.031] | 1.067 [1.055, 1.081] | 1.001 [0.994, 1.021] | | sm_100a | down E4 M112576 bal | ckpt ES | 5.7918 | 7.0925 t | 5.8270 | 5.7619 | 1.000 [0.992, 1.019] | 1.223 [1.162, 1.240] | 1.009 [0.992, 1.020] | | sm_100a | down E4 M112576 skew | ckpt ES | 5.8593 | 7.4696 t | 5.9009 | 5.8066 | 1.010 [1.008, 1.025] | 1.286 [1.271, 1.302] | 1.017 [1.006, 1.025] | | sm_100a | down E4 M112576 1-empty | ckpt ES | 5.9467 | 7.6880 t | 5.9327 | 5.9460 | 1.000 [0.997, 1.013] | 1.297 [1.258, 1.313] | 0.999 [0.965, 1.010] | | sm_100a | down E4 M136064 bal | ckpt ES | 7.0628 | 9.0143 t | 7.0844 | 7.0581 | 1.001 [0.986, 1.016] | 1.262 [1.219, 1.308] | 1.004 [0.966, 1.037] | | sm_100a | down E4 M136064 skew | ckpt ES | 7.1628 | 9.2080 t | 7.1768 | 7.1385 | 1.003 [0.989, 1.028] | 1.308 [1.244, 1.332] | 1.008 [0.989, 1.028] | | sm_100a | down E4 M136064 1-empty | ckpt ES | 7.2625 | 9.3095 t | 7.1643 | 7.1636 | 1.012 [1.001, 1.019] | 1.298 [1.260, 1.407] | 1.001 [0.994, 1.017] | | sm_100a | down E8 M131072 bal | ckpt ES | 6.8311 | 8.1482 t | 6.8018 | 6.8143 | 1.009 [0.983, 1.023] | 1.195 [1.185, 1.263] | 1.000 [0.982, 1.012] | | sm_100a | down E16 M131072 bal | ckpt ES | 6.8208 | 7.6266 t | 6.8268 | 6.7864 | 1.000 [0.960, 1.029] | 1.124 [1.075, 1.126] | 1.001 [0.963, 1.037] | | sm_100a | down E32 M131072 bal | ckpt ES | 7.0760 | 7.5530 t | 7.0087 | 6.9505 | 1.010 [0.981, 1.024] | 1.082 [1.046, 1.098] | 1.008 [0.981, 1.020] | | sm_100a | down E4 M112576 skew asc | ckpt ES | 5.9513 | 7.7195 t | 5.8617 | 5.8067 | 1.026 [1.008, 1.031] | 1.333 [1.288, 1.355] | 1.009 [0.992, 1.016] | | sm_100a | down E4 M112576 1-empty asc | ckpt ES | 6.1629 | 7.7365 t | 5.9436 | 5.9416 | 1.037 [1.019, 1.066] | 1.314 [1.301, 1.326] | 1.007 [0.999, 1.020] | | sm_100a | down E4 M136064 skew asc | ckpt ES | 7.1822 | 9.3037 t | 7.1717 | 7.0899 | 1.014 [0.992, 1.025] | 1.313 [1.280, 1.384] | 1.008 [0.998, 1.024] | | sm_100a | down E4 M136064 1-empty asc | ckpt ES | 7.3673 | 9.4665 t | 7.1147 | 7.1138 | 1.035 [1.015, 1.042] | 1.333 [1.296, 1.345] | 1.001 [0.964, 1.010] | | sm_100a | down E8 M131072 random | ckpt ES | 7.2057 | 8.7284 t | 6.9452 | 6.9467 | 1.040 [1.021, 1.070] | 1.267 [1.235, 1.302] | 1.001 [0.965, 1.035] | | sm_100a | down E16 M131072 random | ckpt ES | 7.1159 | 8.1205 t | 6.9600 | 7.0642 | 1.007 [1.006, 1.028] | 1.154 [1.142, 1.174] | 0.992 [0.983, 1.006] | | sm_100a | down E32 M131072 random | ckpt ES | 7.1903 | 7.7637 t | 7.1047 | 7.0878 | 1.014 [0.994, 1.051] | 1.089 [1.052, 1.114] | 1.003 [0.964, 1.022] | The family's other Cake variant (deterministic, `flashinfer.experimental.cake_moe_grouped_det_gemm`) does not change; both variants ship on sm_100a (B200), sm_103a (GB300) and sm_107a (Rubin); this PR changes only sm_100a and sm_103a programs. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * The k512 unordered-tail routes now support slack-tail placement on SM100a, SM103a, and SM107a. When available capacity allows, each tail tile runs whole on one of the least-loaded clusters; otherwise, tiles continue to split as before. * **Documentation** * Updated the architecture and behavior descriptions to reflect this placement support and clarify when tail tiles run whole or split. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com>
The Cake ragged BF16 grouped GEMM for MoE expert training has two variants. Both use FP32 accumulation. Each variant has forward (
Y = X W^T), data-gradient (dX = G W) and weight-gradient (dW = G^T X, BF16 or FP32 output) programs. The deterministic variant,flashinfer.experimental.cake_moe_grouped_det_gemm, is bitwise reproducible from run to run. Its autograd wrapper,cake_grouped_det_mm, is also exposed asflashinfer.grouped_mm.cake_grouped_det_mm. This wrapper is the default (deterministic=True) of the facadeflashinfer.experimental.cake.grouped_mm.cake_grouped_mm. Callers select the non-deterministic variant,flashinfer.experimental.cake.cake_moe_grouped_gemm, withdeterministic=False.This PR changes only
cake_grouped_det_mm, the autograd wrapper of the deterministic variant. The wrapper now launches the forward GEMM afterCakeGroupedDetMm.applyreturns. At that point, autograd has savedx,wandoffs, as it does for a built-in op.torch.utils.checkpoint(..., use_reentrant=False)has early stop on by default. With early stop on, the backward recomputation stops before this GEMM if these are the last saves of the innermost checkpointed region. The kernels do not change. Outputs and gradients have the same bits as before.The non-deterministic wrapper has launched its forward GEMM after
applysince #8114. With this PR, both variants skip the recomputed forward GEMM under early stop. On early-stop-on rows, the speedup over the previous deterministic wrapper comes from one skipped forward GEMM, not from a faster kernel. A directCakeGroupedDetMm.apply(x, w, offs, deterministic)call still launches the GEMM insideforward, and the backward recomputes it as before. The wrapper undertorch.jit.traceortorch.compiledoes the same.Baselines (table columns):
cake_grouped_det_mmon main, with the launch insideforward. It is measured in the same process asCakeGroupedDetMm.apply(x, w, offs, True)on this branch. That call uses the same programs and the same launch place as the wrapper on main.torch._grouped_mm(PyTorch 2.14.1+cu130) with autograd. It is the only non-Cake implementation with an autograd form for these ragged rows. The per-expert cuBLAS loop and the cuDNN grouped GEMM have none. As a built-in op, it also skips its recomputed forward under early stop. Thus both sides of this comparison have the same recomputation behavior.Operations (column "operation"):
checkpointed fwd + bwd, early stop onand..., early stop off, 18 shapes each. One grouped GEMM in a non-reentrant checkpoint, with forward and backward (data and weight gradients).checkpointed gate_up -> silu*mul -> down + bwd, early stop onand..., early stop off, 9 shapes each. One checkpointed MoE expert block: gate_up GEMM,silu(a) * bon its two halves, and down GEMM. The backward computes the gradients ofx,w1andw2. SiLU and the product save tensors after the gate_up GEMM. Thus early stop can skip only the down GEMM recomputation, and the gain is smaller.Shapes (BF16): gate_up N=4096 K=6144 and down N=6144 K=2048. E=4 with balanced, skewed and one-empty groups at M=112576 and M=136064. E=8, 16 and 32 with balanced groups at M=131072.
Speedup = baseline latency / latency of this PR. Each row gives the median of the per-group paired ratios, with [min, max] over the groups. B200, GB300 and Rubin have 54 rows each. Every early-stop-on row is faster than both baselines. The lowest early-stop-on row median is 1.043 against the previous deterministic wrapper and 1.019 against
torch._grouped_mm. With early stop off, every implementation recomputes its forward. Early-stop-off rows are therefore at parity with the previous deterministic wrapper. Every row below 1.0 is an early-stop-off row against the previous deterministic wrapper, at 0.981 or above.Geometric mean of the per-row speedups. The parentheses give the lowest and highest row median.
B200 (sm_100a)
NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.4, driver 580.82.07, torch 2.14.1+cu130 (CUDA 13.0); CUPTI GPU time of one call, cold L2, eager; median of paired groups in rotated / mirrored order
GB300 (sm_103a)
NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4, driver 580.159.03, torch 2.14.1+cu130 (CUDA 13.0); CUPTI GPU time of one call, cold L2, eager; median of paired groups in rotated / mirrored order
Rubin (sm_107a)
NVIDIA Graphics Device (sm_107a, 208 SMs), CUDA 13.4, driver 615.71.10, torch 2.14.1+cu130 (CUDA 13.0); CUPTI GPU time of one call, cold L2, eager; median of paired groups in rotated / mirrored order
The family's other Cake variant (non-deterministic,
flashinfer.experimental.cake.cake_moe_grouped_gemm) does not change in this PR and has skipped the recompute since #8114, and both variants ship on sm_100a (B200), sm_103a (GB300) and sm_107a (Rubin).Summary by CodeRabbit
New Features
Documentation