You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
perf(cake_rmsnorm): launch the non-deterministic column tasks at once and run the deterministic pair in two GB300 backward windows - #8577
flashinfer.cake_rmsnorm_train is RMSNorm for training. The forward returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and dw (FP32). The two modes, deterministic=True and deterministic=False, use the same forward and differ only in the backward. The family ships for sm_100a (B200), sm_103a (GB300) and sm_107a (R200). The tables cover these forms in both modes:
[T, H] rows with BF16 weight (H512 to H6144)
FP32 weight (w32): H5120 and H6144, plus H2048 and H4096 in the B200 small-T rows
fused residual add (res, H6144) with BF16 or FP32 weight
per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a shared w[256]
This follow-up to #8208 and #8253 changes only the deterministic=False backward. It makes two changes:
Column-task launch (all three GPUs): the column-task programs (H2048, H4096, H5120, H6144, the res H6144 route and their FP32-weight twins) now launch one CTA per row chunk plus one per column task, capped at the slot count, as the deterministic column-task programs already do; the grid changes only at small T, where a few CTAs used to run every column task in turn.
Two GB300 windows: at H6144 from T449 to T1152 (both weight dtypes) and at H5120 (BF16 weight) from T321 to T896, the backward now runs the deterministic pair instead of its single kernel, because the pair is faster there (the pair is the main kernel + fixed-order dw reduce, programs the package already ships).
Speedup = baseline latency / new latency. A value > 1 means the new kernel is faster. Geomeans over the 78-shape list, vs deterministic / vs SOTA:
fwd (same forward in both modes): B200 1.001x / 1.015x, GB300 1.001x / 1.017x, R200 0.999x / 1.012x
In a separate earlier run on B200 and GB300, T = 1 backward took 87-212 us with the previous non-deterministic programs and 3.8-4.5 us with this PR's programs.
torch.compile still beats the new kernel at T = 1 backward on all three GPUs. The exception is GB300 res and res w32. Example, B200 H2048 T1 bwd: det 3.81 us, SOTA 2.14 us (tcd), new 3.87 us. On R200, SOTA is faster on all 8 small-T bwd rows (0.518x to 0.936x vs SOTA).
The rows below are slower than deterministic. This is pre-existing: this PR does not change their GPU programs or launch. Latencies are new vs det:
R200 at large T. The lowest is H5120 T65536 bwd at 300.86 vs 161.73 us (0.54x). Others: H6144 T32768 bwd 162.40 vs 101.82 us (0.63x) and H6144 T65536 bwd 311.11 vs 192.77 us (0.62x). The full list is every R200 row below 0.95x vs det.
GB300: H1536 T65536 bwd 128.45 vs 101.47 us (0.79x); [T,4,256] T1024 bwd 9.70 vs 7.94 us (0.82x); H5120 T65536 and H5120 T65536 w32 bwd 0.85x to 0.87x; H6144 T64 res bwd 8.03 vs 7.36 us (0.92x).
B200: [T,24,256] T65536 bwd 427.77 vs 384.89 us (0.90x, !).
GB300 window rows: both modes launch the same GPU programs here. The new mode measures 0.962x to 0.996x of det on bwd (geomean 0.971). This gap comes from the timing order (see below), not from the GPU programs.
Baselines and measurement. Every row has two baselines:
SOTA: per row, the fastest non-Cake candidate that has the same outputs (BF16 y / dx / h_new, FP32 rstd / dw) and meets the accuracy budget. This is the strongest applicable comparison: every listed candidate is timed on every row, and the fastest in-budget one is taken. APIs that return dw in the weight dtype (Liger, QuACK, TE) do not have matched outputs and are never chosen. Apex FusedRMSNorm has no wheel.
new = the non-deterministic Cake kernel of this PR (deterministic=False). The accuracy budget takes the worst of 5 seeds: y and dx atol = rtol = 1e-2; rstd rel. error <= 1e-4; dw per element atol 1e-4 + rtol 1e-4 vs FP64. New and det meet it on every row on all three GPUs.
Latency = median of 804 CUPTI spans per implementation per row. One span is one CUDA-graph replay after an L2 flush (cold L2). It includes the implementation's own setup, reduction, cast and copy kernels. For each GPU type, 4 sessions ran in parallel on 4 GPUs of one node (2 in row order, 2 in reverse). Each session ran 3 interleaved passes x 67 replays. eps = 1e-6. Latencies print to 0.01 us and per-row ratios to 0.01x.
In each session the implementations run in a fixed order: det first, then the non-Cake candidates, then new. A later check on B200 and GB300 found that this order alone moves an implementation's latency by up to about 5 %: rows where both modes launch the same GPU programs read 0.95x to 1.05x (up to 1.08x on ~ rows), not 1.00x. A ratio within about 5 % of 1.00x is therefore not a measured difference. All implementations of a row run in one process that also loads every non-Cake candidate, including torch.compile. A later check on B200 and GB300 found that a torch.compile function in the process makes the GPU programs of every implementation run faster: on small rows the Cake kernels and the other libraries (quack, te, torch, liger) run up to about 30 % shorter than in a process without torch.compile, and torch.compile itself runs within about 3 % of the same speed in both. So the table compares all implementations in the same state. On small rows the Cake kernels gain somewhat more than the other libraries: the ratio vs those libraries reads up to about 1.10x (GB300) and 1.14x (B200) higher than in a process without torch.compile, so it can overstate the new kernel's lead by up to that factor. On large rows this difference is at most 1.02x. Ratios vs det are not affected (both are Cake kernels in the same process).
Key:
Shapes: H = hidden size, T = tokens. x / y / dx are BF16, rstd / dw are FP32, and w is BF16 unless marked w32. res adds the BF16 output h_new. [T,24,256] has strides (12288, 512, 1), and [T,4,256] is contiguous. Both share w[256].
SOTA:
tc = torch.compile of an explicit FP32-statistics reference (torch 2.14.1+cu130, inductor max-autotune-no-cudagraphs, static shape, emulate_precision_casts=True)
B200, 148 SMs, sm_100a, compute capability 10.0, driver 580.82.07, CUDA 13.4 (nvcc cuda_13.4.r13.4/compiler.38855100_0), torch 2.14.1+cu130. SM clock median 1965 MHz (99 % of samples at max), memory clock 3996 MHz, 4 GPUs of one node, no foreign GPU processes. Rows: 78 shapes + 36 small-T shapes, x 3 ops = 342 rows. The small-T shapes are T1, T4, T16 and T64 of H2048, H2048 w32, H4096, H4096 w32, H5120, H6144, H6144 w32, H6144 res and H6144 res w32. New and det meet the budget on every row. No matched candidate meets it on the 2 ! rows. 39 rows were re-timed in a second run on another B200 node with the same software, because tc did not run on them in the first run: H1536 T65536, plus T256/T1024/T4096 of H5120 w32, H6144 w32, H6144 res and H6144 res w32. All columns of those rows come from the second run.
~: 117 rows
!: 2 rows, [T,24,256] T65536 bwd and 2xfwd+bwd. The fastest matched candidate is shown (tc, dw err/tol 1.122).
GB300, 152 SMs, sm_103a, compute capability 10.3, driver 580.159.03, CUDA 13.4 (same nvcc as B200), torch 2.14.1+cu130. SM clock 2070 MHz (100 % of samples at max), memory clock 3996 MHz, 4 GPUs of one node. Rows: 78 shapes + 20 small-T shapes + 9 window shapes, x 3 ops = 321 rows. The small-T shapes are T1, T4, T16 and T64 of H5120, H6144, H6144 w32, H6144 res and H6144 res w32. The window shapes are H5120 T321/T512/T896, plus T449/T768/T1152 of H6144 and H6144 w32. New, det and SOTA meet the budget on every row.
R200, 212 SMs, sm_107a, compute capability 10.7, driver 620.62, CUDA 13.4, torch 2.14.1+cu130. SM clock 2400-2424 MHz (median 2412), memory clock 4752 MHz, 4 GPUs of one node. This is a 212-SM R200 node. 216-SM R200 boards are timed separately and never pooled with it, and that measurement is pending. Rows: 78 shapes + 8 small-T shapes, x 3 ops = 258 rows. The small-T shapes are T1, T4, T16 and T64 of H6144 res and H6144 res w32. New and det meet the budget on every row.
…column-task backward at once
The column-task programs of the deterministic=False backward (H = 2048, 4096,
5120 and 6144, the residual H = 6144 route and the FP32-weight twins) claim
H / col column tasks ahead of the row chunks, but launch_grid gave them
min(slots, n_chunks) CTAs, so at a small token count a few CTAs ran every
column task in turn. Their records now spell col_tasks and launch_grid
launches min(slots, n_chunks + col_tasks) CTAs, the rule of the deterministic
package's column-task programs. dw and dx keep their bits: each column task
sums its columns over every row in a fixed order whichever CTA runs it.
…ward windows
On sm_103a the deterministic=False backward ran its single kernel at
H = 6144 from T 449 to 1152 (both weight dtypes) and at H = 5120 (BF16
weight) from T 321 to 896, between windows of the deterministic pair,
which is faster there. Those windows now run that pair (the main kernel
and its fixed-order dw reduce, programs this package already ships), so
H = 6144 runs it from T 4 to 4096 (FP32 weight: from T 449; the single
kernel keeps T 321 to 448) and H = 5120 from T 23 to 5120.
No actionable comments were generated in the recent review.
Walkthrough
The backward RMSNorm launch grid now accounts for column-task CTAs alongside claimed chunks. Selected H5120 and H6144 backward module windows and shared-memory wait tables are consolidated or retuned. Tests update grid-layout checks and add small token counts.
Selected module grid rules include col_tasks. The claimed-chunk grid calculation adds column tasks to chunk claims and caps the result at slots. Tests check the updated grid width and add token counts 1 and 3.
Selected H5120 and H6144 backward routes and window boundaries change. The shared-memory wait tables, edge outcomes, and window documentation reflect the updated ranges. The H5120 route comment is revised.
The launch-grid and backward-window changes have no established blocking defect. Merge readiness remains subject to normal checks, including the requested GPU tests.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title clearly identifies both main changes: concurrent non-deterministic column-task launches and GB300 backward-window updates.
Docstring Coverage
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Description check
The description is detailed and relevant. It explains the scope, implementation changes, affected GPU architectures, performance results, accuracy results, measurement method, limitations, and related…
Full details: Docstring Coverage
Explanation
Docstring coverage is 100.00% which is sufficient. 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 2 functions across 5 files.
Full details: Description check
Explanation
The description is detailed and relevant. It explains the scope, implementation changes, affected GPU architectures, performance results, accuracy results, measurement method, limitations, and related issue. It does not reproduce the template headings or explicitly confirm pre-commit and test checklist items, but the required change context is substantially covered.
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.
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 4 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
cake_jit.py:14457 — bwd_h5120 now lists cake_rmsnorm_train_c5831b29d17867842e5f__22_5120, but the diff deletes both c5831b29 records (__22_320, __896_5120) and never adds __22_5120. That will be a KeyError at lookup. Fix: rename __22_320 to __22_5120 with at_most: 5120, as you did for the 2fc4932f and 5f5cdc51 records. If the record already exists somewhere else, ignore this.
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.
…00 backward (#8655)
Related issue: #8066
`flashinfer.cake_rmsnorm_train` is RMSNorm for training. The forward
returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and
dw (FP32). The two modes, `deterministic=True` and
`deterministic=False`, use the same forward and differ only in the
backward. Both modes support these forms on sm_100a (B200), sm_103a
(GB300) and sm_107a (R200):
- [T, H] rows with BF16 weight, or FP32 weight (`w32`)
- fused residual add (`res`, H6144)
- per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a
shared w[256]
This follow-up to #8577 and #8626 changes only the `deterministic=False`
backward at large T, on sm_103a and sm_107a. The bounds of the windows
below are unchanged. In each one, the backward now runs a new
single-kernel program (kblock pipe, blocked FP32 dw sums) in place of
the previous non-deterministic program:
- GB300 (sm_103a): H4096 with BF16 weight, T > 50176; H4096 `w32`, T >
9216; H5120 with either weight dtype, T > 61952.
- R200 (sm_107a): H5120 with either weight dtype, T > 49152.
This PR does not change the sm_100a (B200) routes or programs, so it has
no B200 table. The forward and the deterministic package are also
unchanged.
Speedup = baseline latency / new latency. A value > 1 means the new
kernel is faster. Geomeans over the shapes of the changed windows, vs
deterministic / vs SOTA:
- GB300 (12 shapes): fwd 0.999x / 0.994x, bwd 0.996x / 1.124x, 2xfwd+bwd
1.000x / 1.069x
- R200 (8 shapes): fwd 1.001x / 0.930x, bwd 0.951x / 1.371x, 2xfwd+bwd
0.978x / 1.207x
Against the program each window replaces, for the GB300 H4096 and R200
H5120 windows. These numbers come from a separate route screen on
another node of the same type: 4 sessions, with s2 and s4 in reverse
order. Compare its ratios with the tables below, not its latencies.
Values are previous -> new bwd latency (previous/new):
- GB300 H4096, BF16 weight: 257.4245 -> 254.2535 us at T65536 (1.0125x),
376.5300 -> 373.1205 us at T98304 (1.0091x), 498.1615 -> 494.3530 us at
T131072 (1.0077x)
- GB300 H4096 `w32`: 76.6405 -> 72.9120 us at T16384 (1.0511x), 136.0490
-> 131.6020 us at T32768 (1.0338x), 258.7055 -> 253.8730 us at T65536
(1.0190x)
- R200 H5120, BF16 weight: 301.4755 -> 170.1305 us at T65536 (1.7720x),
447.5260 -> 242.9000 us at T98304 (1.8424x), 594.8560 -> 323.3485 us at
T131072 (1.8397x)
- R200 H5120 `w32`: 297.7805 -> 175.9865 us at T65536 (1.6921x),
441.5740 -> 251.1555 us at T98304 (1.7582x), 586.7600 -> 334.8205 us at
T131072 (1.7525x)
- The new program was faster in every session. Per-session ratios were
1.00738 to 1.05222 on GB300 and 1.68171 to 1.85670 on R200.
- The R200 screen did not include T55296. At T53248 and T57344, the
per-session ratios were 1.75861 to 1.80047 with BF16 weight and 1.68194
to 1.72597 with `w32`.
Notable rows:
- Rows slower than det. In these windows det runs faster programs, but
at some T in these windows those programs do not meet the per-element dw
accuracy bar that the non-deterministic mode requires of its own
programs: the worst dw element must stay within 0.85 of the per-element
dw budget defined below. The non-deterministic mode therefore cannot
route to them. This PR narrows the gap to det but does not close it.
Latencies are new vs det:
- GB300 H4096, BF16 weight, bwd: 254.69 vs 249.97 us at T65536, 374.91
vs 368.90 us at T98304 and 494.75 vs 487.06 us at T131072 (0.98x each).
H4096 `w32` T65536 bwd: 254.27 vs 250.19 us (0.98x). On these four rows
new is slower in all 4 sessions. The other GB300 bwd rows read 0.99x to
1.01x. Of these, only H5120 `w32` T65536 bwd (314.69 vs 312.58 us,
0.99x) is also slower in all 4 sessions. All GB300 2xfwd+bwd rows read
0.99x to 1.02x.
- R200: every bwd row is slower, 0.96x to 0.97x with BF16 weight and
0.93x to 0.95x with `w32`. The 2xfwd+bwd rows read 0.97x to 0.99x. The
largest gap is H5120 T65536 `w32` bwd: 175.71 vs 163.52 us (0.93x).
- Rows below SOTA:
- GB300 H4096 T65536 bwd (BF16 weight): 254.69 vs 246.38 us (QuACK
autotuned, 0.97x).
- GB300 H4096 T131072 bwd (BF16 weight): 494.75 vs 490.08 us
(torch.compile, 0.99x).
- The 2xfwd+bwd rows of these two shapes read 0.99x.
- Every other bwd and 2xfwd+bwd row is at or above 1.00x vs SOTA. GB300
reaches up to 1.22x on bwd and 1.12x on 2xfwd+bwd. R200 reads 1.31x to
1.45x on bwd and 1.12x to 1.32x on 2xfwd+bwd.
- The forward, which this PR does not change, is below SOTA on 9 of 12
GB300 shapes (lowest 0.985x) and on 8 of 8 R200 shapes (0.903x to
1.000x).
- R200 candidates outside the budget. QuACK autotuned is faster than new
on the four `w32` bwd rows and on all eight 2xfwd+bwd rows. Its dw is
outside the budget there (dw err/tol 1.108 to 2.476 over 5 seeds), so it
is not SOTA. All 16 R200 bwd and 2xfwd+bwd rows are marked `o`. On the
four BF16 bwd rows, the out-of-budget candidates beat the SOTA shown but
not new. The candidates are listed below the R200 table.
Baselines and measurement. Every row has two baselines:
- det: the deterministic Cake kernel (`deterministic=True`). On every
row here it launches the same GPU programs as current FlashInfer main.
- SOTA: per row, the fastest non-Cake candidate that has the same
outputs (BF16 y / dx, FP32 rstd / dw) and meets the accuracy budget.
Every applicable candidate is timed on every row. Configurations that
return dw in the weight dtype (the public PyTorch, Liger, QuACK and TE
APIs on BF16-weight rows) do not have matched outputs and are never
chosen.
new = the non-deterministic Cake kernel (`deterministic=False`) with
this PR's programs. The accuracy budget takes the worst of 5 seeds: dx
atol = rtol = 1e-2, and dw per element atol 1e-4 + rtol 1e-4 vs FP64.
Latency = median of the CUPTI spans pooled over 4 sessions (s1 to s4) on
the 4 GPUs of one node. One span is one CUDA-graph replay after an L2
flush (cold L2), and it includes the implementation's own setup,
reduction, cast and copy kernels. eps = 1e-6. Latencies print to 0.01 us
and per-row ratios to 0.01x. The summary geomeans use the unrounded
medians. On GB300, s1 and s3 run the rows in order and s2 and s4 run
them in reverse, each in 3 interleaved rounds. The implementation order
rotates so that every implementation takes every position equally often,
so no implementation gets a better or worse slot from a fixed order. On
the fwd rows, where both modes launch the same programs, det/new reads
0.993 to 1.002 on GB300 and 1.000 to 1.002 on R200.
torch.compile in the process. All implementations of a row run in one
process that also loads every non-Cake candidate, including
torch.compile. The check reported in #8577 (B200 and GB300) found that a
torch.compile function in the process makes every non-compile
implementation run faster. On small rows, the Cake kernels gain somewhat
more than the other non-compile libraries, so the ratio vs those
libraries reads up to about 1.10x (GB300) and 1.14x (B200) higher than
in a process without torch.compile. On large rows this difference is at
most 1.02x, and every row here is large (T >= 16384). Ratios vs det are
not affected, because both are Cake kernels in the same process.
Key:
- Shapes: x [T, H] is BF16, and w is BF16 or FP32 (`w32`). `plain` = no
residual. y / dx are BF16, and rstd / dw are FP32.
- op: `2xfwd+bwd (checkpointed)` = two forwards and one backward.
- `=` after the op: det and new launch the same GPU programs (every fwd
row). The printed ratios are measured values; none is set to 1.00x.
- SOTA:
- torch.compile = torch.compile (torch 2.14.1+cu130, inductor
max-autotune-no-cudagraphs, emulate_precision_casts=True) of an explicit
FP32-statistics reference. In the R200 table it is labelled
torch.compile (inductor).
- torch.compile (default mode) = the same in torch.compile default mode.
In the R200 table it is labelled torch.compile default mode (inductor).
- QuACK = QuACK 0.6.5 (nvidia-cutlass-dsl 4.8.1), heuristic config.
- QuACK autotuned = QuACK 0.6.5 exhaustive per-shape autotune
(rmsnorm_fwd_tuned / rmsnorm_bwd_tuned, cached before timing).
- Liger-Kernel 0.8.5 (Triton 3.8.0, casting_mode gemma) and Transformer
Engine 2.20.2 (run as the FP32 problem, with the casts inside its span)
were never fastest.
- PyTorch's fused rms_norm cannot run this form (BF16 x with FP32 w).
Apex FusedRMSNorm has no wheel. FlashInfer's inference rmsnorm is
forward only.
- `o` (R200): a non-Cake candidate faster than the SOTA shown fails the
dw budget on this row. These candidates are listed below the R200 table.
**GB300 (sm_103a)**
GB300, 152 SMs, sm_103a, compute capability 10.3, driver 580.159.03,
CUDA 13.4 (nvcc cuda_13.4.r13.4/compiler.38855100_0), torch
2.14.1+cu130. SM clock median 2070 MHz (100 % of samples at max), memory
clock 3996 MHz, 4 GPUs of one node, no foreign GPU processes.
Rows: 12 shapes x 3 ops = 36 rows, all inside the four changed GB300
windows:
- H4096 with BF16 weight at T65536, T98304 and T131072
- H4096 `w32` at T16384, T32768 and T65536
- H5120 with either weight dtype at T65536, T98304 and T131072
Summary:
| op | det/new geomean [min, max], rows < 1 | SOTA/new geomean [min,
max], rows < 1 |
|---|---|---|
| fwd | 0.999x [0.993, 1.002], 9 of 12 < 1 | 0.994x [0.985, 1.006], 9 of
12 < 1 |
| bwd | 0.996x [0.981, 1.009], 6 of 12 < 1 | 1.124x [0.967, 1.224], 2 of
12 < 1 |
| 2xfwd+bwd (checkpointed) | 1.000x [0.992, 1.017], 6 of 12 < 1 | 1.069x
[0.991, 1.122], 2 of 12 < 1 |
| GPU / architecture | workload / shape / dtype / mode | operation |
deterministic Cake latency | non-Cake SOTA latency | new
non-deterministic Cake latency | speedup vs deterministic Cake | speedup
vs non-Cake SOTA |
|---|---|---|---|---|---|---|---|
| GB300 sm_103a | x [65536, 4096] BF16, w BF16, plain | fwd = | 150.85
us | 151.68 us torch.compile | 151.39 us | 1.00x | 1.00x |
| GB300 sm_103a | x [65536, 4096] BF16, w BF16, plain | bwd | 249.97 us
| 246.38 us QuACK autotuned | 254.69 us | 0.98x | 0.97x |
| GB300 sm_103a | x [65536, 4096] BF16, w BF16, plain | 2xfwd+bwd
(checkpointed) | 552.69 us | 552.34 us QuACK autotuned | 557.06 us |
0.99x | 0.99x |
| GB300 sm_103a | x [98304, 4096] BF16, w BF16, plain | fwd = | 226.69
us | 224.77 us torch.compile | 226.78 us | 1.00x | 0.99x |
| GB300 sm_103a | x [98304, 4096] BF16, w BF16, plain | bwd | 368.90 us
| 443.65 us torch.compile (default mode) | 374.91 us | 0.98x | 1.18x |
| GB300 sm_103a | x [98304, 4096] BF16, w BF16, plain | 2xfwd+bwd
(checkpointed) | 819.49 us | 890.79 us torch.compile (default mode) |
824.63 us | 0.99x | 1.08x |
| GB300 sm_103a | x [131072, 4096] BF16, w BF16, plain | fwd = | 300.96
us | 298.62 us torch.compile (default mode) | 301.07 us | 1.00x | 0.99x
|
| GB300 sm_103a | x [131072, 4096] BF16, w BF16, plain | bwd | 487.06 us
| 490.08 us torch.compile | 494.75 us | 0.98x | 0.99x |
| GB300 sm_103a | x [131072, 4096] BF16, w BF16, plain | 2xfwd+bwd
(checkpointed) | 1086.05 us | 1082.91 us torch.compile | 1093.06 us |
0.99x | 0.99x |
| GB300 sm_103a | x [65536, 5120] BF16, w BF16, plain | fwd = | 190.94
us | 188.51 us QuACK autotuned | 191.01 us | 1.00x | 0.99x |
| GB300 sm_103a | x [65536, 5120] BF16, w BF16, plain | bwd | 313.82 us
| 362.18 us torch.compile | 313.18 us | 1.00x | 1.16x |
| GB300 sm_103a | x [65536, 5120] BF16, w BF16, plain | 2xfwd+bwd
(checkpointed) | 692.35 us | 751.75 us torch.compile | 692.11 us | 1.00x
| 1.09x |
| GB300 sm_103a | x [98304, 5120] BF16, w BF16, plain | fwd = | 284.32
us | 280.13 us QuACK autotuned | 284.29 us | 1.00x | 0.99x |
| GB300 sm_103a | x [98304, 5120] BF16, w BF16, plain | bwd | 463.94 us
| 551.35 us torch.compile (default mode) | 460.13 us | 1.01x | 1.20x |
| GB300 sm_103a | x [98304, 5120] BF16, w BF16, plain | 2xfwd+bwd
(checkpointed) | 1028.42 us | 1133.76 us torch.compile | 1024.55 us |
1.00x | 1.11x |
| GB300 sm_103a | x [131072, 5120] BF16, w BF16, plain | fwd = | 377.09
us | 374.05 us QuACK autotuned | 377.28 us | 1.00x | 0.99x |
| GB300 sm_103a | x [131072, 5120] BF16, w BF16, plain | bwd | 614.53 us
| 703.27 us torch.compile (default mode) | 608.87 us | 1.01x | 1.16x |
| GB300 sm_103a | x [131072, 5120] BF16, w BF16, plain | 2xfwd+bwd
(checkpointed) | 1364.36 us | 1519.84 us torch.compile (default mode) |
1359.33 us | 1.00x | 1.12x |
| GB300 sm_103a | x [16384, 4096] BF16, w FP32, plain | fwd = | 40.51 us
| 40.93 us QuACK autotuned | 40.80 us | 0.99x | 1.00x |
| GB300 sm_103a | x [16384, 4096] BF16, w FP32, plain | bwd | 73.70 us |
87.68 us QuACK | 73.31 us | 1.01x | 1.20x |
| GB300 sm_103a | x [16384, 4096] BF16, w FP32, plain | 2xfwd+bwd
(checkpointed) | 154.59 us | 170.50 us QuACK | 152.02 us | 1.02x | 1.12x
|
| GB300 sm_103a | x [32768, 4096] BF16, w FP32, plain | fwd = | 78.66 us
| 79.17 us torch.compile (default mode) | 78.66 us | 1.00x | 1.01x |
| GB300 sm_103a | x [32768, 4096] BF16, w FP32, plain | bwd | 133.34 us
| 146.30 us torch.compile | 133.44 us | 1.00x | 1.10x |
| GB300 sm_103a | x [32768, 4096] BF16, w FP32, plain | 2xfwd+bwd
(checkpointed) | 287.36 us | 300.45 us torch.compile | 286.58 us | 1.00x
| 1.05x |
| GB300 sm_103a | x [65536, 4096] BF16, w FP32, plain | fwd = | 153.02
us | 152.93 us torch.compile | 153.09 us | 1.00x | 1.00x |
| GB300 sm_103a | x [65536, 4096] BF16, w FP32, plain | bwd | 250.19 us
| 257.35 us torch.compile | 254.27 us | 0.98x | 1.01x |
| GB300 sm_103a | x [65536, 4096] BF16, w FP32, plain | 2xfwd+bwd
(checkpointed) | 554.75 us | 558.51 us torch.compile | 557.63 us | 0.99x
| 1.00x |
| GB300 sm_103a | x [65536, 5120] BF16, w FP32, plain | fwd = | 191.30
us | 190.66 us QuACK autotuned | 191.62 us | 1.00x | 0.99x |
| GB300 sm_103a | x [65536, 5120] BF16, w FP32, plain | bwd | 312.58 us
| 370.43 us torch.compile | 314.69 us | 0.99x | 1.18x |
| GB300 sm_103a | x [65536, 5120] BF16, w FP32, plain | 2xfwd+bwd
(checkpointed) | 694.42 us | 763.01 us torch.compile | 697.31 us | 1.00x
| 1.09x |
| GB300 sm_103a | x [98304, 5120] BF16, w FP32, plain | fwd = | 285.73
us | 282.85 us QuACK autotuned | 285.57 us | 1.00x | 0.99x |
| GB300 sm_103a | x [98304, 5120] BF16, w FP32, plain | bwd | 462.50 us
| 565.81 us torch.compile (default mode) | 462.37 us | 1.00x | 1.22x |
| GB300 sm_103a | x [98304, 5120] BF16, w FP32, plain | 2xfwd+bwd
(checkpointed) | 1032.56 us | 1149.51 us torch.compile (default mode) |
1032.88 us | 1.00x | 1.11x |
| GB300 sm_103a | x [131072, 5120] BF16, w FP32, plain | fwd = | 381.14
us | 375.14 us QuACK autotuned | 380.35 us | 1.00x | 0.99x |
| GB300 sm_103a | x [131072, 5120] BF16, w FP32, plain | bwd | 613.17 us
| 720.05 us torch.compile (default mode) | 611.65 us | 1.00x | 1.18x |
| GB300 sm_103a | x [131072, 5120] BF16, w FP32, plain | 2xfwd+bwd
(checkpointed) | 1374.18 us | 1494.37 us torch.compile (default mode) |
1371.11 us | 1.00x | 1.09x |
**R200 (sm_107a)**
R200, 212 SMs, sm_107a, compute capability 10.7, CUDA 13.4 (nvcc
cuda_13.4.r13.4/compiler.38855100_0), torch 2.14.1+cu130. Memory clock
4752 MHz at the start and end of the run, 4 GPUs of one node, no foreign
GPU processes.
Rows: 8 shapes x 3 ops = 24 rows, all inside the changed R200 window:
H5120 with either weight dtype at T55296, T65536, T98304 and T131072.
Summary:
| op | det/new geomean [min, max], rows < 1 | SOTA/new geomean [min,
max], rows < 1 |
|---|---|---|
| fwd | 1.001x [1.000, 1.002], 1 of 8 < 1 | 0.930x [0.903, 1.000], 8 of
8 < 1 |
| bwd | 0.951x [0.931, 0.974], 8 of 8 < 1 | 1.371x [1.312, 1.445], 0 of
8 < 1 |
| 2xfwd+bwd (checkpointed) | 0.978x [0.969, 0.989], 8 of 8 < 1 | 1.207x
[1.122, 1.318], 0 of 8 < 1 |
| GPU / architecture | workload / shape / dtype / mode | operation |
deterministic Cake latency | non-Cake SOTA latency | new
non-deterministic Cake latency | speedup vs deterministic Cake | speedup
vs non-Cake SOTA |
|---|---|---|---|---|---|---|---|
| NVIDIA Graphics Device sm_107a (212 SMs) | x [55296, 5120] BF16, w
BF16, plain | fwd = | 95.55 us | 86.70 us QuACK (autotuned) | 95.55 us |
1.00x | 0.91x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [55296, 5120] BF16, w
BF16, plain | bwd | 138.85 us | 202.88 us torch.compile (inductor) `o` |
144.03 us | 0.96x | 1.41x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [55296, 5120] BF16, w
BF16, plain | 2xfwd+bwd (checkpointed) | 328.99 us | 399.46 us
torch.compile (inductor) `o` | 334.37 us | 0.98x | 1.19x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [65536, 5120] BF16, w
BF16, plain | fwd = | 112.35 us | 101.92 us QuACK (autotuned) | 112.35
us | 1.00x | 0.91x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [65536, 5120] BF16, w
BF16, plain | bwd | 161.83 us | 222.27 us torch.compile (inductor) `o` |
169.38 us | 0.96x | 1.31x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [65536, 5120] BF16, w
BF16, plain | 2xfwd+bwd (checkpointed) | 385.06 us | 450.40 us
torch.compile (inductor) `o` | 392.83 us | 0.98x | 1.15x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [98304, 5120] BF16, w
BF16, plain | fwd = | 166.69 us | 150.50 us QuACK (autotuned) | 166.66
us | 1.00x | 0.90x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [98304, 5120] BF16, w
BF16, plain | bwd | 236.74 us | 346.11 us torch.compile (inductor) `o` |
243.04 us | 0.97x | 1.42x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [98304, 5120] BF16, w
BF16, plain | 2xfwd+bwd (checkpointed) | 568.77 us | 683.07 us
torch.compile (inductor) `o` | 575.23 us | 0.99x | 1.19x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [131072, 5120] BF16, w
BF16, plain | fwd = | 220.90 us | 199.36 us QuACK (autotuned) | 220.51
us | 1.00x | 0.90x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [131072, 5120] BF16, w
BF16, plain | bwd | 312.55 us | 430.59 us torch.compile default mode
(inductor) `o` | 325.11 us | 0.96x | 1.32x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [131072, 5120] BF16, w
BF16, plain | 2xfwd+bwd (checkpointed) | 752.67 us | 987.35 us
torch.compile default mode (inductor) `o` | 765.27 us | 0.98x | 1.29x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [55296, 5120] BF16, w
FP32, plain | fwd = | 94.53 us | 94.43 us QuACK (autotuned) | 94.43 us |
1.00x | 1.00x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [55296, 5120] BF16, w
FP32, plain | bwd | 140.26 us | 207.39 us torch.compile (inductor) `o` |
149.81 us | 0.94x | 1.38x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [55296, 5120] BF16, w
FP32, plain | 2xfwd+bwd (checkpointed) | 325.63 us | 418.72 us
torch.compile (inductor) `o` | 335.58 us | 0.97x | 1.25x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [65536, 5120] BF16, w
FP32, plain | fwd = | 112.03 us | 105.73 us torch.compile (inductor) |
111.78 us | 1.00x | 0.95x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [65536, 5120] BF16, w
FP32, plain | bwd | 163.52 us | 233.38 us torch.compile (inductor) `o` |
175.71 us | 0.93x | 1.33x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [65536, 5120] BF16, w
FP32, plain | 2xfwd+bwd (checkpointed) | 383.20 us | 443.75 us
torch.compile (inductor) `o` | 395.49 us | 0.97x | 1.12x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [98304, 5120] BF16, w
FP32, plain | fwd = | 165.15 us | 155.20 us torch.compile (inductor) |
165.12 us | 1.00x | 0.94x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [98304, 5120] BF16, w
FP32, plain | bwd | 239.30 us | 363.30 us torch.compile (inductor) `o` |
251.42 us | 0.95x | 1.45x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [98304, 5120] BF16, w
FP32, plain | 2xfwd+bwd (checkpointed) | 564.90 us | 671.97 us
torch.compile (inductor) `o` | 577.92 us | 0.98x | 1.16x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [131072, 5120] BF16, w
FP32, plain | fwd = | 219.14 us | 204.67 us torch.compile (inductor) |
219.20 us | 1.00x | 0.93x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [131072, 5120] BF16, w
FP32, plain | bwd | 314.91 us | 452.80 us torch.compile default mode
(inductor) `o` | 335.62 us | 0.94x | 1.35x |
| NVIDIA Graphics Device sm_107a (212 SMs) | x [131072, 5120] BF16, w
FP32, plain | 2xfwd+bwd (checkpointed) | 748.61 us | 1014.21 us
torch.compile default mode (inductor) `o` | 769.73 us | 0.97x | 1.32x |
`o`: faster non-Cake candidates outside the accuracy budget on each
marked row:
- H5120 T55296 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 146.82 us (dx err/tol 0.311, dw err/tol 1.234, y err/tol
0.312 over 5 seeds), QuACK 198.24 us (dx err/tol 0.311, dw err/tol
1.234, y err/tol 0.312 over 5 seeds)
- H5120 T55296 2xfwd+bwd (checkpointed): faster but outside the accuracy
budget: QuACK (autotuned) 318.85 us (dx err/tol 0.311, dw err/tol 1.234,
y err/tol 0.312 over 5 seeds), QuACK 392.32 us (dx err/tol 0.311, dw
err/tol 1.234, y err/tol 0.312 over 5 seeds)
- H5120 T65536 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 170.56 us (dx err/tol 0.312, dw err/tol 1.355, y err/tol
0.311 over 5 seeds)
- H5120 T65536 2xfwd+bwd (checkpointed): faster but outside the accuracy
budget: QuACK (autotuned) 372.45 us (dx err/tol 0.312, dw err/tol 1.355,
y err/tol 0.311 over 5 seeds)
- H5120 T98304 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 248.10 us (dx err/tol 0.312, dw err/tol 1.829, y err/tol
0.312 over 5 seeds), QuACK 336.85 us (dx err/tol 0.312, dw err/tol
1.829, y err/tol 0.312 over 5 seeds)
- H5120 T98304 2xfwd+bwd (checkpointed): faster but outside the accuracy
budget: QuACK (autotuned) 547.30 us (dx err/tol 0.312, dw err/tol 1.829,
y err/tol 0.312 over 5 seeds), QuACK 674.69 us (dx err/tol 0.312, dw
err/tol 1.829, y err/tol 0.312 over 5 seeds)
- H5120 T131072 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 325.33 us (dx err/tol 0.312, dw err/tol 2.206, y err/tol
0.312 over 5 seeds), torch.compile (inductor) 420.58 us (dx err/tol
0.312, dw err/tol 1.080, y err/tol 0.312 over 5 seeds)
- H5120 T131072 2xfwd+bwd (checkpointed): faster but outside the
accuracy budget: QuACK (autotuned) 722.08 us (dx err/tol 0.312, dw
err/tol 2.206, y err/tol 0.312 over 5 seeds), torch.compile (inductor)
866.85 us (dx err/tol 0.312, dw err/tol 1.080, y err/tol 0.312 over 5
seeds), QuACK 891.52 us (dx err/tol 0.312, dw err/tol 2.206, y err/tol
0.312 over 5 seeds)
- H5120 T55296 w32 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 143.33 us (dx err/tol 0.312, dw err/tol 1.108, y err/tol
0.312 over 5 seeds), QuACK 194.56 us (dx err/tol 0.312, dw err/tol
1.108, y err/tol 0.312 over 5 seeds)
- H5120 T55296 w32 2xfwd+bwd (checkpointed): faster but outside the
accuracy budget: QuACK (autotuned) 330.46 us (dx err/tol 0.312, dw
err/tol 1.108, y err/tol 0.312 over 5 seeds), QuACK 416.10 us (dx
err/tol 0.312, dw err/tol 1.108, y err/tol 0.312 over 5 seeds)
- H5120 T65536 w32 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 166.62 us (dx err/tol 0.311, dw err/tol 1.327, y err/tol
0.312 over 5 seeds), QuACK 227.94 us (dx err/tol 0.311, dw err/tol
1.327, y err/tol 0.312 over 5 seeds)
- H5120 T65536 w32 2xfwd+bwd (checkpointed): faster but outside the
accuracy budget: QuACK (autotuned) 386.37 us (dx err/tol 0.311, dw
err/tol 1.327, y err/tol 0.312 over 5 seeds)
- H5120 T98304 w32 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 242.72 us (dx err/tol 0.312, dw err/tol 1.925, y err/tol
0.312 over 5 seeds), QuACK 333.22 us (dx err/tol 0.312, dw err/tol
1.925, y err/tol 0.312 over 5 seeds)
- H5120 T98304 w32 2xfwd+bwd (checkpointed): faster but outside the
accuracy budget: QuACK (autotuned) 568.74 us (dx err/tol 0.312, dw
err/tol 1.925, y err/tol 0.312 over 5 seeds)
- H5120 T131072 w32 bwd: faster but outside the accuracy budget: QuACK
(autotuned) 317.95 us (dx err/tol 0.312, dw err/tol 2.476, y err/tol
0.312 over 5 seeds), QuACK 438.40 us (dx err/tol 0.312, dw err/tol
2.476, y err/tol 0.312 over 5 seeds), torch.compile (inductor) 441.06 us
(dx err/tol 0.312, dw err/tol 1.114, y err/tol 0.312 over 5 seeds)
- H5120 T131072 w32 2xfwd+bwd (checkpointed): faster but outside the
accuracy budget: QuACK (autotuned) 748.42 us (dx err/tol 0.312, dw
err/tol 2.476, y err/tol 0.312 over 5 seeds), torch.compile (inductor)
848.96 us (dx err/tol 0.312, dw err/tol 1.114, y err/tol 0.312 over 5
seeds), QuACK 952.56 us (dx err/tol 0.312, dw err/tol 2.476, y err/tol
0.312 over 5 seeds)
The family has no other variants. It ships these forms in both modes on
sm_100a, sm_103a and sm_107a, and this PR changes and measures only
sm_103a and sm_107a.
Co-authored-by: Yingyi Huang <averyh@nvidia.com>
… token counts (#8704)
Related issue: #8066
`flashinfer.cake_rmsnorm_train` is RMSNorm for training. The forward
returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and
dw (FP32). The two modes, `deterministic=True` and
`deterministic=False`, use the same forward and differ only in the
backward. Both modes support the forms listed in #8655 on sm_100a
(B200), sm_103a (GB300) and sm_107a (R200):
- [T, H] rows with BF16 weight, or FP32 weight (`w32`)
- fused residual add (`res`, H6144)
- per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a
shared w[256]
This follow-up to #8577, #8626 and #8655 speeds up the
`deterministic=False` backward at small token counts on all three
architectures. It covers 12 backward routes per architecture: H1536,
H2048, H4096, H5120, H6144 and residual H6144, each with BF16 and FP32
weight. On each route, the backward for T in (0, cut] now runs a new
single-kernel program (fixed-claim TwoSum column task). The windows are:
|form|weight|B200 sm_100a|GB300 sm_103a|R200 sm_107a|
|---|---|---|---|---|
|H1536|BF16|(0, 64]|(0, 64]|(0, 64]|
|H2048|BF16|(0, 64]|(0, 64]|(0, 64]|
|H4096|BF16|(0, 64]|(0, 32]|(0, 64]|
|H5120|BF16|(0, 64]|(0, 64]|(0, 64]|
|H6144|BF16|(0, 64]|(0, 32]|(0, 64]|
|H6144 residual|BF16|(0, 64]|(0, 48]|(0, 64]|
|H1536|FP32|(0, 64]|(0, 64]|(0, 64]|
|H2048|FP32|(0, 64]|(0, 64]|(0, 64]|
|H4096|FP32|(0, 64]|(0, 32]|(0, 64]|
|H5120|FP32|(0, 64]|(0, 48]|(0, 64]|
|H6144|FP32|(0, 64]|(0, 48]|(0, 64]|
|H6144 residual|FP32|(0, 64]|(0, 48]|(0, 64]|
This PR does not change any other route or T, the forward or the
deterministic package.
Speedup = baseline latency / new latency. A value > 1 means the new
kernel is faster.
In the B200 and GB300 records, the new backward is faster than det on
every in-window bwd row: 1.18x to 1.65x on B200 and 1.05x to 1.59x on
GB300. 2xfwd+bwd is also faster than det on every in-window row except
two GB300 rows: H6144 w FP32 T48 (new 10.18 us vs det 9.79 us, 0.96x)
and H5120 w FP32 T48 (new 9.09 us vs det 9.06 us; the ratio prints as
1.00x). Against the strongest in-budget non-Cake SOTA, bwd reads 0.82x
to 2.09x on B200 and 0.70x to 1.84x on GB300. Most rows below 1 are at
T1, where torch.compile's backward is faster. A separate backward-only
timing also compares new with the `deterministic=False` program it
replaces. In that timing, new is faster than both det and the replaced
program on every in-window row of all three GPUs. Only the R200 rows of
that timing are shown below.
Rows inside the windows, B200 and GB300 (CUDA toolkit 13.4, CUDA runtime
13.0; values from the record tables below), min-max over the rows
(number of rows < 1 in parentheses):
|GPU|op|rows|det/new|SOTA/new|
|---|---|---|---|---|
|B200 sm_100a|fwd|48|0.97-1.04x (9)|0.82-1.13x (35)|
|B200 sm_100a|bwd|48|1.18-1.65x (0)|0.82-2.09x (6)|
|B200 sm_100a|2xfwd+bwd|48|1.10-1.28x (0)|0.96-1.50x (4)|
|GB300 sm_103a|fwd|36|0.96-1.04x (4)|0.75-1.07x (29)|
|GB300 sm_103a|bwd|36|1.05-1.59x (0)|0.70-1.84x (9)|
|GB300 sm_103a|2xfwd+bwd|36|0.96-1.31x (2)|0.94-1.91x (3)|
The GB300 2xfwd+bwd vs det count includes H5120 w FP32 T48, whose ratio
prints as 1.00x (new 9.09 us vs det 9.06 us). The B200 fwd vs det count
(9) uses the printed ratio. The record headline below counts 11 of 54
B200 fwd rows < 1 vs det before rounding. That is one more than the rows
that print below 1.00x (9 inside the windows plus 1 just above them),
and the source does not say which row it is.
The records' headline geomeans (B200 and GB300, CUDA toolkit 13.4, CUDA
runtime 13.0) cover all their rows. That includes the six T = cut + 1
rows per GPU and op, which this PR does not change (see below), and the
rows marked `~` (session spread above 5 %). No geomean exists for the
in-window rows alone.
|GPU|op|new vs det: geomean [min, max], rows < 1|new vs SOTA: geomean
[min, max], rows < 1|
|---|---|---|---|
|B200 sm_100a|fwd|1.000x [0.97, 1.04], 11 of 54 rows < 1|0.954x [0.82,
1.13], 40 of 54 rows < 1|
|B200 sm_100a|bwd|1.367x [0.70, 1.65], 2 of 54 rows < 1|1.315x [0.82,
2.09], 7 of 54 rows < 1|
|B200 sm_100a|2xfwd+bwd|1.169x [0.84, 1.28], 2 of 54 rows < 1|1.158x
[0.96, 1.50], 6 of 54 rows < 1|
|GB300 sm_103a|fwd|1.003x [0.96, 1.04], 5 of 42 rows < 1|0.945x [0.75,
1.07], 34 of 42 rows < 1|
|GB300 sm_103a|bwd|1.294x [0.92, 1.59], 2 of 42 rows < 1|1.154x [0.70,
1.84], 14 of 42 rows < 1|
|GB300 sm_103a|2xfwd+bwd|1.149x [0.96, 1.31], 5 of 42 rows < 1|1.247x
[0.94, 1.91], 4 of 42 rows < 1|
The R200 (sm_107a) two-baseline record is pending. It needs det, SOTA
and new with the full candidate set, under the same protocol as B200 and
GB300. The R200 table below is a backward-only timing of the same
programs, labelled as such.
Notable rows. Each latency pair below reads new vs baseline.
- **Rows slower than det, inside the windows.**
- GB300 H6144 plain T48 w FP32, 2xfwd+bwd: new 10.18 us vs det 9.79 us
(0.96x). On the same shape, bwd reads 1.05x (new 5.92 us vs det 6.21 us)
and fwd 1.00x. This row is not marked `~`.
- GB300 H5120 plain T48 w FP32, 2xfwd+bwd: new 9.09 us vs det 9.06 us
(prints as 1.00x). This row is not marked `~`.
- fwd: the forward is unchanged, and both modes launch the same forward
programs. These rows read 0.97x to 1.04x vs det on B200 (9 of 48 below 1
by printed ratio; see the note above) and 0.96x to 1.04x on GB300 (4 of
36 below 1).
- No in-window bwd row is slower than det on either GPU.
- **bwd rows slower than SOTA.**
- At T1, torch.compile beats new on H1536, H2048 and H4096 with either
weight dtype, on both GPUs. On B200 these rows read 0.82x to 0.92x, for
example H2048 plain T1 w BF16 at new 2.46 us vs TCd 2.03 us. On GB300
they read 0.70x to 0.97x, for example H2048 plain T1 w FP32 at new 2.85
us vs TCd 1.98 us. All twelve of these rows are marked `~`.
- GB300, FP32 weight, at the top of the window: H2048 T64 new 4.90 us vs
TCd 4.83 us (0.99x), H5120 T48 new 5.44 us vs TC 5.15 us (0.95x) and
H6144 T48 new 5.92 us vs TCd 5.70 us (0.96x `~`).
- **2xfwd+bwd rows slower than SOTA.** All are at T1 and all are marked
`~`.
- B200: H2048 w BF16 0.96x, H4096 w BF16 0.97x, H6144 w BF16 0.98x,
H4096 w FP32 0.97x.
- GB300: H1536 w BF16 0.98x, H2048 w BF16 0.98x, H2048 w FP32 0.94x.
- **fwd rows slower than SOTA.** This PR does not change the forward. It
is below SOTA on 35 of 48 B200 rows (down to 0.82x) and on 29 of 36
GB300 rows (down to 0.75x).
- **R200, backward-only timing.**
- new is faster than det and than main's `deterministic=False` program
on every row.
- torch.compile (the faster of its two modes) beats new on every T1 row
(0.479x to 0.759x), on every H4096 row and on most rows at T >= 32. It
also beats new on two H2048 rows: T4 w BF16 (new 3.17 us vs 3.09 us,
0.975x) and T2 w FP32 (new 3.18 us vs 3.04 us, 0.956x).
Rows just above the windows. Each record also timed every BF16-weight
form at T = cut + 1. That is T65, except T33 for GB300 H4096 and H6144
and T49 for GB300 residual H6144. These rows run main's
`deterministic=False` programs, which this PR does not change, so they
are not in the tables below. On these rows, new is slower than det in
these bwd and 2xfwd+bwd cases (new vs det):
- B200 H1536 T65: bwd 5.70 vs 3.97 us (0.70x), 2xfwd+bwd 9.15 vs 7.71 us
(0.84x)
- B200 H5120 T65: bwd 4.08 vs 4.03 us (0.99x)
- B200 H2048 T65: 2xfwd+bwd 7.14 vs 7.01 us (0.98x)
- GB300 H1536 T65: bwd 5.54 vs 5.09 us (0.92x), 2xfwd+bwd 8.16 vs 7.90
us (0.97x)
- GB300 H2048 T65: bwd 6.11 vs 5.79 us (0.95x), 2xfwd+bwd 8.80 vs 8.70
us (0.99x)
- GB300 residual H6144 T49: 2xfwd+bwd 11.78 vs 11.46 us (0.97x)
The fwd rows run the same programs in both modes. Two of them also read
below 1: B200 H1536 T65 (1.89 vs 1.86 us, 0.98x) and GB300 H6144 T33
(2.62 vs 2.59 us, 0.99x `~`).
Against SOTA, the bwd rows above the windows read 0.87x to 1.21x on B200
and 0.75x to 1.11x on GB300.
Baselines and measurement (B200 and GB300 tables). Every row has two
baselines:
- det: the deterministic Cake kernel (`deterministic=True`) of the same
tree. The FlashInfer deterministic package is unchanged from main.
- SOTA: the fastest non-Cake candidate that has these outputs and meets
the accuracy budget. The budget is dx atol = rtol = 1e-2, and dw per
element atol 1e-4 + rtol 1e-4 vs FP64, taking the worst of 5 seeds.
Candidates timed:
- QuACK 0.6.5, heuristic and autotuned
- torch.compile, max-autotune-no-cudagraphs and default mode
- Liger-Kernel 0.8.5 (Triton 3.8.0)
- Transformer Engine 2.20.2
new is the non-deterministic Cake kernel (`deterministic=False`) with
this PR's programs. Accuracy (new / det / SOTA) is pass / pass / pass on
every row.
Timing setup:
- Sessions: 4 sessions (s1/s3 forward, s2/s4 reverse) run in parallel on
the node's 4 GPUs.
- Latency: the median of the CUPTI span of one CUDA-graph replay with
cold L2, over 2 sessions x 3 interleaved rounds x 201 replays, pooled.
The p10-p90 ranges are omitted here.
- Span contents: each span includes the implementation's own setup,
reduction, cast and copy kernels.
- Workload: eps 1e-6. y / dx / h_new are BF16, and rstd and dw are FP32,
for every implementation.
- Process: all implementations of a row run in one process that also
loads every non-Cake candidate, including torch.compile.
- Order: the implementation order rotates, so every implementation takes
every position equally often.
Key:
- Workload `H<H> plain|res T<T> w BF16|FP32`: x [T, H] is BF16, and w is
BF16 or FP32. `plain` = no residual. `res` = fused residual add (h_new =
BF16(h + u), combined grad), H6144.
- op: `2xfwd+bwd` = two forwards and one backward (checkpointed). On
every fwd row, det and new launch the same programs. The printed fwd
ratios are measured values.
- SOTA names:
- TC = torch.compile (torch 2.14.1+cu130, inductor mode
max-autotune-no-cudagraphs, emulate_precision_casts=True) of the
explicit FP32-statistics reference
- TCd = the same with torch.compile default mode
- QA = QuACK 0.6.5 rmsnorm_fwd_tuned / rmsnorm_bwd_tuned (exhaustive
per-shape autotune, cached before timing)
- Liger = Liger-Kernel 0.8.5 (Triton 3.8.0), casting_mode gemma (FP32
math)
- Transformer Engine was never the fastest.
- `~` marks a ratio whose baseline (det or SOTA) or new latency has a
session spread above 5 % (|s1 - s2| / min of the per-session medians).
Such rows are noisy at this size.
**B200 (sm_100a): new non-deterministic Cake vs deterministic Cake vs
non-Cake SOTA, rows inside this PR's windows (T <= cut)**
GPU B200 sm_100a, 148 SMs, compute capability 10.0; driver 580.82.07;
CUDA toolkit Build cuda_13.4.r13.4/compiler.38855100_0; torch
2.14.1+cu130 (CUDA runtime 13.0). NVML SM clock during timing 1965-1965
MHz (median 1965, 100% of samples at max) in each of s1-s4; memory
3996-3996 MHz; foreign GPU processes at start / end 0 / 0. Residual
bias: Cake gains somewhat more from the context than the other
non-compile libraries on small rows: SOTA/new vs the non-compile
candidates is up to about 1.10x higher (GB300) / 1.14x higher (B200) in
the one-process measurement than in a process without torch.compile;
large rows <= 1.02x. So small-row ratios may overstate the new kernel's
lead by up to about that much. Ratios vs det are unaffected. All rows
here are small rows (T <= 65), and the kernels run 2-8 us.
Rows: 12 forms at T1, T4, T16 and T64, x 3 ops = 144 rows.
|workload (H / form / T / w)|op|det us|SOTA us (name)|new
us|det/new|SOTA/new|
|---|---|---|---|---|---|---|
|H1536 plain T1 w BF16|fwd|1.86|1.79 Liger|1.82|1.02x ~|0.98x ~|
|H1536 plain T1 w BF16|bwd|3.87|2.02 TC|2.40|1.61x|0.84x ~|
|H1536 plain T1 w BF16|2xfwd+bwd|7.42|6.14 TCd|6.11|1.21x|1.01x ~|
|H1536 plain T4 w BF16|fwd|1.82|1.76 Liger|1.86|0.98x ~|0.95x ~|
|H1536 plain T4 w BF16|bwd|3.87|3.36 TC|2.50|1.55x|1.35x ~|
|H1536 plain T4 w BF16|2xfwd+bwd|6.82|6.26 TC|5.41|1.26x ~|1.16x ~|
|H1536 plain T16 w BF16|fwd|1.86|1.82 Liger|1.86|1.00x ~|0.98x ~|
|H1536 plain T16 w BF16|bwd|3.87|3.71 TC|2.53|1.53x ~|1.47x ~|
|H1536 plain T16 w BF16|2xfwd+bwd|6.91|7.46 TC|5.44|1.27x ~|1.37x ~|
|H1536 plain T64 w BF16|fwd|1.89|1.82 Liger|1.89|1.00x|0.97x|
|H1536 plain T64 w BF16|bwd|3.97|3.97 TC|2.78|1.42x|1.43x|
|H1536 plain T64 w BF16|2xfwd+bwd|6.91|7.42 TC|5.70|1.21x|1.30x|
|H2048 plain T1 w BF16|fwd|1.70|1.66 Liger|1.63|1.04x ~|1.02x ~|
|H2048 plain T1 w BF16|bwd|3.81|2.03 TCd|2.46|1.55x ~|0.82x ~|
|H2048 plain T1 w BF16|2xfwd+bwd|6.75|5.09 TC|5.28|1.28x|0.96x ~|
|H2048 plain T4 w BF16|fwd|1.82|1.76 Liger|1.82|1.00x ~|0.96x ~|
|H2048 plain T4 w BF16|bwd|3.84|3.30 TC|2.50|1.54x ~|1.32x ~|
|H2048 plain T4 w BF16|2xfwd+bwd|6.82|6.11 TC|5.44|1.25x ~|1.12x ~|
|H2048 plain T16 w BF16|fwd|1.89|1.82 Liger|1.89|1.00x ~|0.97x|
|H2048 plain T16 w BF16|bwd|3.87|4.42 TC|2.56|1.51x|1.73x ~|
|H2048 plain T16 w BF16|2xfwd+bwd|7.58|8.22 TC|6.34|1.20x ~|1.30x ~|
|H2048 plain T64 w BF16|fwd|1.92|1.86 Liger|1.92|1.00x ~|0.97x ~|
|H2048 plain T64 w BF16|bwd|4.00|3.98 TC|2.94|1.36x|1.35x ~|
|H2048 plain T64 w BF16|2xfwd+bwd|7.04|8.00 TC|5.82|1.21x|1.37x|
|H4096 plain T1 w BF16|fwd|2.02|2.02 Liger|2.02|1.00x ~|1.00x ~|
|H4096 plain T1 w BF16|bwd|3.87|2.46 TC|3.01|1.29x ~|0.82x ~|
|H4096 plain T1 w BF16|2xfwd+bwd|7.09|6.05 TC|6.24|1.14x ~|0.97x ~|
|H4096 plain T4 w BF16|fwd|1.98|2.02 QA|1.98|1.00x ~|1.02x ~|
|H4096 plain T4 w BF16|bwd|3.90|3.74 TCd|3.01|1.30x|1.24x ~|
|H4096 plain T4 w BF16|2xfwd+bwd|7.17|6.97 TCd|6.24|1.15x|1.12x ~|
|H4096 plain T16 w BF16|fwd|1.98|2.02 Liger|2.05|0.97x ~|0.98x ~|
|H4096 plain T16 w BF16|bwd|3.94|4.13 TC|3.07|1.28x|1.34x|
|H4096 plain T16 w BF16|2xfwd+bwd|7.14|7.49 TCd|6.27|1.14x|1.19x|
|H4096 plain T64 w BF16|fwd|2.05|2.05 Liger|2.05|1.00x|1.00x|
|H4096 plain T64 w BF16|bwd|5.18|4.51 TCd|4.06|1.28x ~|1.11x|
|H4096 plain T64 w BF16|2xfwd+bwd|8.45|7.71 TCd|7.17|1.18x|1.08x|
|H5120 plain T1 w BF16|fwd|2.24|2.18 Liger|2.24|1.00x ~|0.97x ~|
|H5120 plain T1 w BF16|bwd|3.90|3.46 TC|2.85|1.37x|1.21x ~|
|H5120 plain T1 w BF16|2xfwd+bwd|8.22|7.58 TCd|7.10|1.16x ~|1.07x ~|
|H5120 plain T4 w BF16|fwd|2.14|2.02 Liger|2.18|0.99x ~|0.93x ~|
|H5120 plain T4 w BF16|bwd|3.87|4.29 TC|2.78|1.39x|1.54x|
|H5120 plain T4 w BF16|2xfwd+bwd|7.49|7.81 TC|6.27|1.19x ~|1.25x ~|
|H5120 plain T16 w BF16|fwd|2.18|2.11 Liger|2.18|1.00x ~|0.97x ~|
|H5120 plain T16 w BF16|bwd|3.94|4.58 TC|2.88|1.37x|1.59x|
|H5120 plain T16 w BF16|2xfwd+bwd|7.46|8.13 TC|6.50|1.15x|1.25x|
|H5120 plain T64 w BF16|fwd|2.24|2.16 Liger|2.24|1.00x|0.96x|
|H5120 plain T64 w BF16|bwd|4.10|5.09 TC|3.46|1.19x|1.47x|
|H5120 plain T64 w BF16|2xfwd+bwd|7.71|8.74 TC|6.98|1.11x|1.25x|
|H6144 plain T1 w BF16|fwd|2.53|2.06 Liger|2.53|1.00x ~|0.82x ~|
|H6144 plain T1 w BF16|bwd|4.00|2.75 TC|2.66|1.51x ~|1.04x ~|
|H6144 plain T1 w BF16|2xfwd+bwd|8.10|6.66 TC|6.78|1.19x ~|0.98x ~|
|H6144 plain T4 w BF16|fwd|2.50|2.08 Liger|2.46|1.01x ~|0.84x ~|
|H6144 plain T4 w BF16|bwd|3.94|4.26 TC|2.69|1.46x|1.58x ~|
|H6144 plain T4 w BF16|2xfwd+bwd|8.02|7.78 TCd|6.78|1.18x|1.15x ~|
|H6144 plain T16 w BF16|fwd|2.53|2.11 Liger|2.53|1.00x|0.84x|
|H6144 plain T16 w BF16|bwd|3.97|5.06 TCd|2.85|1.39x|1.78x|
|H6144 plain T16 w BF16|2xfwd+bwd|8.74|9.12 TCd|7.55|1.16x ~|1.21x ~|
|H6144 plain T64 w BF16|fwd|2.59|2.21 Liger|2.59|1.00x|0.85x|
|H6144 plain T64 w BF16|bwd|6.37|5.12 TC|4.16|1.53x|1.23x|
|H6144 plain T64 w BF16|2xfwd+bwd|10.37|8.86 TC|8.35|1.24x|1.06x|
|H6144 res T1 w BF16|fwd|2.85|2.54 Liger|2.88|0.99x ~|0.88x ~|
|H6144 res T1 w BF16|bwd|4.03|3.14 TC|2.72|1.48x|1.15x|
|H6144 res T1 w BF16|2xfwd+bwd|8.67|7.65 TC|7.42|1.17x ~|1.03x ~|
|H6144 res T4 w BF16|fwd|2.82|2.53 Liger|2.85|0.99x ~|0.89x ~|
|H6144 res T4 w BF16|bwd|3.97|4.86 TC|2.75|1.44x ~|1.77x ~|
|H6144 res T4 w BF16|2xfwd+bwd|8.70|9.31 TC|7.42|1.17x|1.25x ~|
|H6144 res T16 w BF16|fwd|2.91|2.66 Liger|2.91|1.00x|0.91x|
|H6144 res T16 w BF16|bwd|4.06|4.70 TC|2.85|1.43x|1.65x|
|H6144 res T16 w BF16|2xfwd+bwd|8.89|9.28 TC|7.58|1.17x|1.22x|
|H6144 res T64 w BF16|fwd|3.14|2.91 Liger|3.14|1.00x|0.93x|
|H6144 res T64 w BF16|bwd|7.17|5.92 TC|4.35|1.65x|1.36x ~|
|H6144 res T64 w BF16|2xfwd+bwd|12.13|10.78 TC|9.44|1.28x ~|1.14x ~|
|H1536 plain T1 w FP32|fwd|1.70|1.86 Liger|1.70|1.00x ~|1.09x ~|
|H1536 plain T1 w FP32|bwd|3.74|2.21 TC|2.40|1.56x|0.92x ~|
|H1536 plain T1 w FP32|2xfwd+bwd|7.20|6.30 TC|5.79|1.24x ~|1.09x ~|
|H1536 plain T4 w FP32|fwd|1.82|1.79 Liger|1.79|1.02x ~|1.00x ~|
|H1536 plain T4 w FP32|bwd|3.84|4.45 TCd|2.50|1.54x|1.78x ~|
|H1536 plain T4 w FP32|2xfwd+bwd|7.49|8.46 TC|6.08|1.23x ~|1.39x ~|
|H1536 plain T16 w FP32|fwd|1.86|1.82 Liger|1.85|1.00x ~|0.98x ~|
|H1536 plain T16 w FP32|bwd|3.90|3.87 TC|2.50|1.56x|1.55x ~|
|H1536 plain T16 w FP32|2xfwd+bwd|7.01|8.26 TC|5.50|1.27x ~|1.50x ~|
|H1536 plain T64 w FP32|fwd|1.89|1.82 Liger|1.89|1.00x|0.97x|
|H1536 plain T64 w FP32|bwd|3.90|4.26 TC|2.78|1.40x|1.53x ~|
|H1536 plain T64 w FP32|2xfwd+bwd|6.85|7.52 TC|5.63|1.22x|1.34x|
|H2048 plain T1 w FP32|fwd|1.66|1.89 Liger|1.66|1.00x ~|1.13x ~|
|H2048 plain T1 w FP32|bwd|3.71|2.05 TCd|2.50|1.49x|0.82x ~|
|H2048 plain T1 w FP32|2xfwd+bwd|6.69|5.47 TCd|5.31|1.26x ~|1.03x ~|
|H2048 plain T4 w FP32|fwd|1.79|1.79 Liger|1.82|0.98x ~|0.98x ~|
|H2048 plain T4 w FP32|bwd|3.71|3.30 TC|2.53|1.47x|1.30x ~|
|H2048 plain T4 w FP32|2xfwd+bwd|6.75|6.16 TC|5.50|1.23x ~|1.12x ~|
|H2048 plain T16 w FP32|fwd|1.89|1.82 Liger|1.89|1.00x ~|0.97x ~|
|H2048 plain T16 w FP32|bwd|3.71|4.29 TC|2.56|1.45x|1.68x ~|
|H2048 plain T16 w FP32|2xfwd+bwd|7.49|8.06 TC|6.18|1.21x ~|1.31x ~|
|H2048 plain T64 w FP32|fwd|1.95|1.86 Liger|1.95|1.00x|0.95x|
|H2048 plain T64 w FP32|bwd|4.03|4.51 TC|2.94|1.37x|1.53x ~|
|H2048 plain T64 w FP32|2xfwd+bwd|7.97|8.42 TC|6.66|1.20x ~|1.26x ~|
|H4096 plain T1 w FP32|fwd|1.98|1.98 QA|1.95|1.02x ~|1.02x ~|
|H4096 plain T1 w FP32|bwd|3.68|2.69 TCd|3.01|1.22x|0.89x ~|
|H4096 plain T1 w FP32|2xfwd+bwd|7.14|6.24 TCd|6.43|1.11x ~|0.97x ~|
|H4096 plain T4 w FP32|fwd|1.98|2.02 QA|1.98|1.00x ~|1.02x|
|H4096 plain T4 w FP32|bwd|3.68|3.74 TC|3.01|1.22x|1.24x ~|
|H4096 plain T4 w FP32|2xfwd+bwd|7.04|7.12 TCd|6.24|1.13x ~|1.14x ~|
|H4096 plain T16 w FP32|fwd|2.02|2.08 Liger|2.02|1.00x|1.03x|
|H4096 plain T16 w FP32|bwd|3.74|4.26 TC|3.17|1.18x|1.34x ~|
|H4096 plain T16 w FP32|2xfwd+bwd|7.10|7.52 TC|6.46|1.10x ~|1.16x ~|
|H4096 plain T64 w FP32|fwd|2.08|2.14 Liger|2.08|1.00x|1.03x|
|H4096 plain T64 w FP32|bwd|4.96|4.58 TCd|4.03|1.23x|1.13x|
|H4096 plain T64 w FP32|2xfwd+bwd|8.38|7.87 TCd|7.23|1.16x|1.09x|
|H5120 plain T1 w FP32|fwd|2.08|2.11 QA|2.11|0.98x ~|1.00x ~|
|H5120 plain T1 w FP32|bwd|3.78|2.82 TC|2.78|1.36x|1.01x ~|
|H5120 plain T1 w FP32|2xfwd+bwd|8.06|7.26 TC|6.91|1.17x ~|1.05x ~|
|H5120 plain T4 w FP32|fwd|2.08|2.08 QA|2.08|1.00x ~|1.00x ~|
|H5120 plain T4 w FP32|bwd|3.78|4.64 TC|2.78|1.36x|1.67x ~|
|H5120 plain T4 w FP32|2xfwd+bwd|7.97|8.64 TC|7.14|1.12x ~|1.21x ~|
|H5120 plain T16 w FP32|fwd|2.17|2.11 QA|2.14|1.01x|0.99x|
|H5120 plain T16 w FP32|bwd|3.81|4.69 TC|2.88|1.32x|1.63x ~|
|H5120 plain T16 w FP32|2xfwd+bwd|7.68|8.38 TC|6.56|1.17x ~|1.28x ~|
|H5120 plain T64 w FP32|fwd|2.21|2.21 QA|2.22|0.99x|0.99x|
|H5120 plain T64 w FP32|bwd|4.13|5.22 TC|3.36|1.23x|1.55x|
|H5120 plain T64 w FP32|2xfwd+bwd|7.78|9.09 TCd|7.07|1.10x|1.29x|
|H6144 plain T1 w FP32|fwd|2.43|2.08 QA|2.43|1.00x ~|0.86x ~|
|H6144 plain T1 w FP32|bwd|3.74|3.20 TC|2.66|1.41x|1.20x ~|
|H6144 plain T1 w FP32|2xfwd+bwd|7.87|7.17 TC|6.69|1.18x ~|1.07x ~|
|H6144 plain T4 w FP32|fwd|2.40|2.08 QA|2.43|0.99x ~|0.86x ~|
|H6144 plain T4 w FP32|bwd|3.74|4.48 TC|2.66|1.41x|1.69x ~|
|H6144 plain T4 w FP32|2xfwd+bwd|7.81|8.10 TC|6.69|1.17x ~|1.21x ~|
|H6144 plain T16 w FP32|fwd|2.46|2.14 QA|2.46|1.00x|0.87x|
|H6144 plain T16 w FP32|bwd|3.84|5.12 TC|2.78|1.38x|1.84x ~|
|H6144 plain T16 w FP32|2xfwd+bwd|8.51|9.70 TC|7.55|1.13x ~|1.28x ~|
|H6144 plain T64 w FP32|fwd|2.53|2.30 Liger|2.53|1.00x|0.91x|
|H6144 plain T64 w FP32|bwd|6.50|5.57 TC|4.16|1.56x|1.34x ~|
|H6144 plain T64 w FP32|2xfwd+bwd|10.98|9.89 TC|8.93|1.23x ~|1.11x ~|
|H6144 res T1 w FP32|fwd|2.85|2.69 Liger|2.85|1.00x ~|0.94x ~|
|H6144 res T1 w FP32|bwd|3.84|3.58 TC|2.72|1.41x|1.32x ~|
|H6144 res T1 w FP32|2xfwd+bwd|8.64|8.06 TC|7.39|1.17x ~|1.09x ~|
|H6144 res T4 w FP32|fwd|2.88|2.66 Liger|2.88|1.00x ~|0.92x ~|
|H6144 res T4 w FP32|bwd|3.87|5.63 TC|2.75|1.41x|2.05x ~|
|H6144 res T4 w FP32|2xfwd+bwd|8.74|10.02 TC|7.39|1.18x ~|1.35x ~|
|H6144 res T16 w FP32|fwd|2.98|2.72 Liger|2.98|1.00x|0.91x|
|H6144 res T16 w FP32|bwd|3.94|5.89 TC|2.82|1.40x|2.09x|
|H6144 res T16 w FP32|2xfwd+bwd|8.90|10.59 TC|7.74|1.15x|1.37x|
|H6144 res T64 w FP32|fwd|3.20|2.98 Liger|3.17|1.01x|0.94x|
|H6144 res T64 w FP32|bwd|7.07|6.14 TC|4.32|1.64x ~|1.42x ~|
|H6144 res T64 w FP32|2xfwd+bwd|12.19|11.30 TC|9.57|1.27x|1.18x|
**GB300 (sm_103a): new non-deterministic Cake vs deterministic Cake vs
non-Cake SOTA, rows inside this PR's windows (T <= cut)**
GPU GB300 sm_103a, 152 SMs, compute capability 10.3; driver 580.159.03;
CUDA toolkit Build cuda_13.4.r13.4/compiler.38855100_0; torch
2.14.1+cu130 (CUDA runtime 13.0). NVML SM clock during timing 2070-2070
MHz (median 2070, 100% of samples at max) in each of s1-s4; memory
3996-3996 MHz; foreign GPU processes at start / end 0 / 0. Residual
bias: Cake gains somewhat more from the context than the other
non-compile libraries on small rows: SOTA/new vs the non-compile
candidates is up to about 1.10x higher (GB300) / 1.14x higher (B200) in
the one-process measurement than in a process without torch.compile;
large rows <= 1.02x. So small-row ratios may overstate the new kernel's
lead by up to about that much. Ratios vs det are unaffected. All rows
here are small rows (T <= 65), and the kernels run 2-8 us.
Rows: 12 forms at T1, T16 and the top of each window (T64, T32 or T48),
x 3 ops = 108 rows.
|workload (H / form / T / w)|op|det us|SOTA us (name)|new
us|det/new|SOTA/new|
|---|---|---|---|---|---|---|
|H1536 plain T1 w BF16|fwd|1.82|1.76 Liger|1.82|1.00x ~|0.96x ~|
|H1536 plain T1 w BF16|bwd|4.13|2.02 TCd|2.82|1.47x|0.72x ~|
|H1536 plain T1 w BF16|2xfwd+bwd|7.97|6.38 TCd|6.50|1.23x ~|0.98x ~|
|H1536 plain T16 w BF16|fwd|1.82|1.73 Liger|1.82|1.00x ~|0.95x|
|H1536 plain T16 w BF16|bwd|4.45|4.64 TCd|3.01|1.48x|1.54x ~|
|H1536 plain T16 w BF16|2xfwd+bwd|7.39|8.51 TCd|5.89|1.26x|1.45x ~|
|H1536 plain T64 w BF16|fwd|1.86|1.79 Liger|1.82|1.02x|0.98x|
|H1536 plain T64 w BF16|bwd|5.12|3.90 TC|3.78|1.36x|1.03x|
|H1536 plain T64 w BF16|2xfwd+bwd|7.97|8.45 TCd|6.62|1.20x|1.28x|
|H2048 plain T1 w BF16|fwd|1.70|1.60 Liger|1.66|1.02x ~|0.96x ~|
|H2048 plain T1 w BF16|bwd|4.35|2.02 TCd|2.78|1.56x|0.72x ~|
|H2048 plain T1 w BF16|2xfwd+bwd|7.14|5.34 TCd|5.47|1.30x ~|0.98x ~|
|H2048 plain T16 w BF16|fwd|1.82|1.76 Liger|1.82|1.00x|0.96x|
|H2048 plain T16 w BF16|bwd|4.83|4.67 TCd|3.04|1.59x|1.54x ~|
|H2048 plain T16 w BF16|2xfwd+bwd|7.84|8.48 TCd|5.98|1.31x|1.42x ~|
|H2048 plain T64 w BF16|fwd|1.86|1.79 Liger|1.86|1.00x|0.97x|
|H2048 plain T64 w BF16|bwd|5.82|5.25 TCd|4.83|1.21x|1.09x ~|
|H2048 plain T64 w BF16|2xfwd+bwd|9.57|9.98 TCd|8.61|1.11x ~|1.16x ~|
|H4096 plain T1 w BF16|fwd|2.02|1.95 QA|1.95|1.03x ~|1.00x ~|
|H4096 plain T1 w BF16|bwd|4.03|2.91 TCd|3.01|1.34x|0.97x ~|
|H4096 plain T1 w BF16|2xfwd+bwd|7.26|7.33 TCd|6.24|1.16x|1.17x ~|
|H4096 plain T16 w BF16|fwd|1.95|1.95 Liger|1.95|1.00x ~|1.00x ~|
|H4096 plain T16 w BF16|bwd|4.26|4.67 TCd|3.20|1.33x ~|1.46x|
|H4096 plain T16 w BF16|2xfwd+bwd|7.46|8.74 TCd|6.30|1.18x|1.39x|
|H4096 plain T32 w BF16|fwd|1.95|1.95 Liger|1.95|1.00x|1.00x|
|H4096 plain T32 w BF16|bwd|5.70|4.83 TCd|3.87|1.47x|1.25x|
|H4096 plain T32 w BF16|2xfwd+bwd|8.77|8.90 TCd|6.94|1.26x|1.28x|
|H5120 plain T1 w BF16|fwd|2.21|1.98 QA|2.18|1.01x ~|0.91x ~|
|H5120 plain T1 w BF16|bwd|4.45|3.62 TCd|3.10|1.43x|1.16x ~|
|H5120 plain T1 w BF16|2xfwd+bwd|8.96|9.09 TCd|7.39|1.21x ~|1.23x ~|
|H5120 plain T16 w BF16|fwd|2.11|2.05 Liger|2.11|1.00x|0.97x ~|
|H5120 plain T16 w BF16|bwd|4.86|5.60 TC|3.81|1.28x|1.47x ~|
|H5120 plain T16 w BF16|2xfwd+bwd|8.54|10.62 TCd|7.39|1.16x|1.44x|
|H5120 plain T64 w BF16|fwd|2.18|2.11 Liger|2.18|1.00x|0.97x|
|H5120 plain T64 w BF16|bwd|5.95|6.08 TCd|5.57|1.07x|1.09x|
|H5120 plain T64 w BF16|2xfwd+bwd|9.12|10.91 TCd|9.02|1.01x|1.21x|
|H6144 plain T1 w BF16|fwd|2.53|2.05 QA|2.53|1.00x ~|0.81x ~|
|H6144 plain T1 w BF16|bwd|4.42|3.81 TCd|2.98|1.48x|1.28x ~|
|H6144 plain T1 w BF16|2xfwd+bwd|8.58|9.25 TCd|7.01|1.22x|1.32x ~|
|H6144 plain T16 w BF16|fwd|2.46|2.08 Liger|2.46|1.00x|0.84x|
|H6144 plain T16 w BF16|bwd|5.06|5.57 TCd|3.78|1.34x|1.47x|
|H6144 plain T16 w BF16|2xfwd+bwd|9.18|10.94 TCd|7.78|1.18x|1.41x|
|H6144 plain T32 w BF16|fwd|2.50|2.11 Liger|2.50|1.00x|0.85x|
|H6144 plain T32 w BF16|bwd|6.08|6.08 TCd|4.96|1.23x|1.23x ~|
|H6144 plain T32 w BF16|2xfwd+bwd|10.59|12.35 TCd|9.50|1.11x|1.30x|
|H6144 res T1 w BF16|fwd|2.78|2.53 Liger|2.82|0.99x ~|0.90x ~|
|H6144 res T1 w BF16|bwd|4.45|5.22 TCd|3.01|1.48x|1.73x ~|
|H6144 res T1 w BF16|2xfwd+bwd|9.09|14.40 TCd|7.62|1.19x ~|1.89x ~|
|H6144 res T16 w BF16|fwd|2.85|2.62 Liger|2.85|1.00x|0.92x|
|H6144 res T16 w BF16|bwd|5.12|5.15 TC|3.84|1.33x|1.34x|
|H6144 res T16 w BF16|2xfwd+bwd|9.86|16.03 Liger|8.38|1.18x|1.91x|
|H6144 res T48 w BF16|fwd|2.98|2.78 Liger|2.98|1.00x|0.94x|
|H6144 res T48 w BF16|bwd|6.82|7.55 TCd|6.18|1.10x|1.22x ~|
|H6144 res T48 w BF16|2xfwd+bwd|11.46|11.94 TCd|11.17|1.03x ~|1.07x ~|
|H1536 plain T1 w FP32|fwd|1.70|1.89 Liger|1.76|0.96x ~|1.07x ~|
|H1536 plain T1 w FP32|bwd|4.19|2.21 TC|2.85|1.47x ~|0.78x ~|
|H1536 plain T1 w FP32|2xfwd+bwd|7.81|6.37 TCd|6.30|1.24x ~|1.01x ~|
|H1536 plain T16 w FP32|fwd|1.82|1.76 Liger|1.82|1.00x ~|0.96x|
|H1536 plain T16 w FP32|bwd|4.45|4.83 TCd|3.01|1.48x|1.61x ~|
|H1536 plain T16 w FP32|2xfwd+bwd|8.00|9.15 TCd|6.72|1.19x|1.36x ~|
|H1536 plain T64 w FP32|fwd|1.82|1.79 Liger|1.82|1.00x ~|0.98x ~|
|H1536 plain T64 w FP32|bwd|5.12|4.86 TCd|3.78|1.36x|1.29x|
|H1536 plain T64 w FP32|2xfwd+bwd|8.48|8.90 TCd|7.07|1.20x ~|1.26x ~|
|H2048 plain T1 w FP32|fwd|1.79|1.86 Liger|1.73|1.04x ~|1.07x ~|
|H2048 plain T1 w FP32|bwd|4.32|1.98 TCd|2.85|1.52x|0.70x ~|
|H2048 plain T1 w FP32|2xfwd+bwd|7.17|5.44 TCd|5.79|1.24x ~|0.94x ~|
|H2048 plain T16 w FP32|fwd|1.86|1.79 Liger|1.86|1.00x ~|0.97x ~|
|H2048 plain T16 w FP32|bwd|4.80|4.77 TCd|3.10|1.55x|1.54x ~|
|H2048 plain T16 w FP32|2xfwd+bwd|7.87|8.45 TCd|6.02|1.31x|1.40x ~|
|H2048 plain T64 w FP32|fwd|1.89|1.82 Liger|1.89|1.00x|0.97x|
|H2048 plain T64 w FP32|bwd|6.08|4.83 TCd|4.90|1.24x|0.99x|
|H2048 plain T64 w FP32|2xfwd+bwd|8.99|8.51 TCd|7.68|1.17x|1.11x|
|H4096 plain T1 w FP32|fwd|1.92|1.95 QA|1.98|0.97x ~|0.98x ~|
|H4096 plain T1 w FP32|bwd|4.13|2.66 TCd|3.01|1.37x|0.88x ~|
|H4096 plain T1 w FP32|2xfwd+bwd|8.29|7.81 TCd|7.01|1.18x ~|1.11x ~|
|H4096 plain T16 w FP32|fwd|2.02|2.02 QA|1.98|1.02x ~|1.02x ~|
|H4096 plain T16 w FP32|bwd|4.29|4.96 TCd|3.26|1.31x|1.52x|
|H4096 plain T16 w FP32|2xfwd+bwd|8.38|9.89 TCd|7.04|1.19x ~|1.40x ~|
|H4096 plain T32 w FP32|fwd|1.98|2.05 Liger|1.95|1.02x|1.05x|
|H4096 plain T32 w FP32|bwd|5.57|5.18 TCd|3.87|1.44x|1.34x ~|
|H4096 plain T32 w FP32|2xfwd+bwd|8.99|9.15 TCd|7.14|1.26x ~|1.28x ~|
|H5120 plain T1 w FP32|fwd|2.27|2.02 QA|2.27|1.00x ~|0.89x ~|
|H5120 plain T1 w FP32|bwd|4.48|3.39 TCd|3.07|1.46x|1.10x ~|
|H5120 plain T1 w FP32|2xfwd+bwd|8.03|7.97 TCd|6.66|1.21x ~|1.20x ~|
|H5120 plain T16 w FP32|fwd|2.21|2.08 Liger|2.18|1.01x|0.96x ~|
|H5120 plain T16 w FP32|bwd|4.93|5.44 TCd|3.81|1.29x|1.43x|
|H5120 plain T16 w FP32|2xfwd+bwd|8.58|10.18 TCd|7.36|1.17x|1.38x|
|H5120 plain T48 w FP32|fwd|2.24|2.14 Liger|2.21|1.01x|0.97x|
|H5120 plain T48 w FP32|bwd|5.76|5.15 TC|5.44|1.06x|0.95x|
|H5120 plain T48 w FP32|2xfwd+bwd|9.06|10.46 TCd|9.09|1.00x|1.15x|
|H6144 plain T1 w FP32|fwd|2.43|1.86 QA|2.46|0.99x ~|0.75x ~|
|H6144 plain T1 w FP32|bwd|4.54|3.52 TCd|2.98|1.53x|1.18x ~|
|H6144 plain T1 w FP32|2xfwd+bwd|9.22|9.54 TCd|7.78|1.19x ~|1.23x ~|
|H6144 plain T16 w FP32|fwd|2.40|2.08 QA|2.40|1.00x|0.87x|
|H6144 plain T16 w FP32|bwd|5.06|5.38 TC|3.74|1.35x|1.44x ~|
|H6144 plain T16 w FP32|2xfwd+bwd|9.89|11.62 TCd|8.22|1.20x|1.41x ~|
|H6144 plain T48 w FP32|fwd|2.43|2.21 Liger|2.43|1.00x|0.91x|
|H6144 plain T48 w FP32|bwd|6.21|5.70 TCd|5.92|1.05x|0.96x ~|
|H6144 plain T48 w FP32|2xfwd+bwd|9.79|11.04 TCd|10.18|0.96x|1.08x|
|H6144 res T1 w FP32|fwd|2.78|2.56 Liger|2.75|1.01x ~|0.93x ~|
|H6144 res T1 w FP32|bwd|4.54|4.74 TCd|2.94|1.54x|1.61x ~|
|H6144 res T1 w FP32|2xfwd+bwd|9.22|12.83 Liger|7.49|1.23x ~|1.71x ~|
|H6144 res T16 w FP32|fwd|2.85|2.62 Liger|2.85|1.00x|0.92x|
|H6144 res T16 w FP32|bwd|5.06|6.85 TCd|3.71|1.36x|1.84x|
|H6144 res T16 w FP32|2xfwd+bwd|9.98|13.06 Liger|8.38|1.19x|1.56x|
|H6144 res T48 w FP32|fwd|3.01|2.78 Liger|3.01|1.00x|0.93x|
|H6144 res T48 w FP32|bwd|6.62|7.90 TCd|5.92|1.12x|1.34x ~|
|H6144 res T48 w FP32|2xfwd+bwd|11.42|13.22 Liger|11.01|1.04x|1.20x|
**R200 (sm_107a): backward-only timing, not the two-baseline record**
GPU R200 sm_107a, 212 SMs; CUDA [GAP: the CUDA toolkit and runtime of
this R200 run are not in the source; fill in before posting].
Backward-only timing: latencies are span medians in us, pooled over 4
sessions (CUDA graph, L2 flush, CUPTI span; balanced implementation
order).
- new = this PR's program.
- det = `deterministic=True`.
- FI main ND = the `deterministic=False` program that FlashInfer main
runs at that T, which this PR replaces.
- best torch.compile = the faster of the torch.compile default and
max-autotune-no-cudagraphs modes on the FP32 reference backward. It is
the only non-Cake baseline in this timing.
The measured programs are text-identical (CUDA and host code) to this
PR's. The R200 rows with det, SOTA from the full candidate set, and new
are pending. New is faster than det and than FI main ND on every row.
Rows: 12 forms at T 1, 2, 3, 4, 8, 12, 16, 24, 32, 48 and 64, all inside
the R200 windows.
|workload (H / form / T / w)|det us|best torch.compile us|FI main ND
us|new us|det/new|best torch.compile/new|FI main ND/new|
|---|---|---|---|---|---|---|---|
|H1536 plain T1 w BF16|4.74|1.55|7.28|3.23|1.465x|0.480x|2.252x|
|H1536 plain T2 w BF16|4.86|3.57|7.62|3.28|1.483x|1.088x|2.322x|
|H1536 plain T3 w BF16|4.77|3.60|7.55|3.30|1.447x|1.091x|2.291x|
|H1536 plain T4 w BF16|4.86|3.87|7.46|3.30|1.476x|1.173x|2.262x|
|H1536 plain T8 w BF16|4.67|3.76|7.76|3.41|1.371x|1.103x|2.277x|
|H1536 plain T12 w BF16|4.94|3.73|7.66|3.39|1.458x|1.100x|2.259x|
|H1536 plain T16 w BF16|4.75|3.78|4.85|3.54|1.344x|1.068x|1.371x|
|H1536 plain T24 w BF16|4.86|3.84|4.88|3.57|1.363x|1.076x|1.368x|
|H1536 plain T32 w BF16|4.70|3.58|4.69|3.44|1.367x|1.041x|1.363x|
|H1536 plain T48 w BF16|4.70|3.60|4.74|3.66|1.284x|0.984x|1.293x|
|H1536 plain T64 w BF16|4.74|3.71|4.75|3.79|1.249x|0.979x|1.253x|
|H2048 plain T1 w BF16|4.59|1.58|8.91|3.02|1.519x|0.523x|2.947x|
|H2048 plain T2 w BF16|4.59|3.20|9.25|3.09|1.487x|1.036x|2.995x|
|H2048 plain T3 w BF16|4.56|3.36|8.91|3.09|1.477x|1.087x|2.886x|
|H2048 plain T4 w BF16|4.46|3.09|8.77|3.17|1.409x|0.975x|2.768x|
|H2048 plain T8 w BF16|4.61|3.36|4.61|3.23|1.426x|1.040x|1.426x|
|H2048 plain T12 w BF16|4.62|3.38|4.62|3.28|1.410x|1.030x|1.410x|
|H2048 plain T16 w BF16|4.64|3.41|4.59|3.26|1.422x|1.046x|1.407x|
|H2048 plain T24 w BF16|4.59|3.52|4.48|3.33|1.380x|1.057x|1.346x|
|H2048 plain T32 w BF16|4.61|3.50|4.61|3.47|1.327x|1.009x|1.327x|
|H2048 plain T48 w BF16|4.66|3.66|4.61|3.81|1.223x|0.961x|1.210x|
|H2048 plain T64 w BF16|4.64|3.79|4.64|4.06|1.142x|0.933x|1.142x|
|H4096 plain T1 w BF16|5.17|1.79|12.24|3.74|1.381x|0.479x|3.269x|
|H4096 plain T2 w BF16|5.30|3.34|12.29|3.82|1.385x|0.874x|3.213x|
|H4096 plain T3 w BF16|5.18|3.41|12.14|3.82|1.356x|0.893x|3.176x|
|H4096 plain T4 w BF16|5.22|3.44|11.81|3.86|1.353x|0.891x|3.063x|
|H4096 plain T8 w BF16|5.30|3.62|5.30|3.79|1.397x|0.955x|1.397x|
|H4096 plain T12 w BF16|5.28|3.58|5.28|3.84|1.375x|0.932x|1.375x|
|H4096 plain T16 w BF16|5.31|3.62|5.31|4.05|1.312x|0.894x|1.312x|
|H4096 plain T24 w BF16|5.33|3.84|5.31|4.14|1.286x|0.928x|1.282x|
|H4096 plain T32 w BF16|5.31|3.81|5.30|4.22|1.257x|0.903x|1.254x|
|H4096 plain T48 w BF16|5.34|3.90|5.36|4.58|1.168x|0.852x|1.171x|
|H4096 plain T64 w BF16|5.38|3.89|5.39|4.74|1.135x|0.821x|1.139x|
|H5120 plain T1 w BF16|5.92|2.05|12.32|3.46|1.713x|0.592x|3.565x|
|H5120 plain T2 w BF16|6.03|4.02|12.13|3.54|1.706x|1.136x|3.430x|
|H5120 plain T3 w BF16|6.02|4.10|12.62|3.60|1.671x|1.139x|3.507x|
|H5120 plain T4 w BF16|6.03|4.24|12.59|3.63|1.661x|1.168x|3.467x|
|H5120 plain T8 w BF16|6.27|4.24|12.53|3.68|1.704x|1.152x|3.404x|
|H5120 plain T12 w BF16|6.26|4.32|12.45|3.78|1.657x|1.143x|3.297x|
|H5120 plain T16 w BF16|6.00|4.24|12.43|3.86|1.556x|1.098x|3.224x|
|H5120 plain T24 w BF16|6.13|4.37|6.22|4.16|1.473x|1.050x|1.496x|
|H5120 plain T32 w BF16|5.94|4.08|5.92|4.18|1.421x|0.976x|1.418x|
|H5120 plain T48 w BF16|5.92|4.06|5.90|4.58|1.294x|0.886x|1.290x|
|H5120 plain T64 w BF16|5.97|4.26|5.97|4.62|1.291x|0.922x|1.291x|
|H6144 plain T1 w BF16|5.41|2.10|13.25|3.39|1.594x|0.619x|3.906x|
|H6144 plain T2 w BF16|5.23|3.63|13.35|3.30|1.587x|1.100x|4.049x|
|H6144 plain T3 w BF16|5.44|3.70|13.68|3.41|1.596x|1.085x|4.015x|
|H6144 plain T4 w BF16|5.34|3.84|5.28|3.46|1.547x|1.110x|1.528x|
|H6144 plain T8 w BF16|5.26|3.97|5.25|3.58|1.469x|1.109x|1.464x|
|H6144 plain T12 w BF16|5.26|3.95|5.23|3.55|1.482x|1.113x|1.473x|
|H6144 plain T16 w BF16|5.28|3.95|5.31|3.70|1.429x|1.068x|1.437x|
|H6144 plain T24 w BF16|5.33|4.10|5.36|3.95|1.348x|1.038x|1.356x|
|H6144 plain T32 w BF16|5.44|4.13|5.42|4.11|1.323x|1.005x|1.319x|
|H6144 plain T48 w BF16|5.42|4.13|5.42|4.53|1.198x|0.912x|1.198x|
|H6144 plain T64 w BF16|5.55|4.40|5.58|4.66|1.192x|0.944x|1.199x|
|H6144 res T1 w BF16|5.17|2.53|5.28|3.39|1.524x|0.746x|1.557x|
|H6144 res T2 w BF16|5.14|3.78|5.15|3.46|1.486x|1.092x|1.491x|
|H6144 res T3 w BF16|5.41|3.92|5.14|3.55|1.523x|1.104x|1.446x|
|H6144 res T4 w BF16|5.31|4.43|5.33|3.50|1.516x|1.266x|1.521x|
|H6144 res T8 w BF16|5.25|4.54|5.25|3.58|1.464x|1.268x|1.464x|
|H6144 res T12 w BF16|5.33|4.13|5.31|3.62|1.473x|1.141x|1.469x|
|H6144 res T16 w BF16|5.25|4.13|5.22|3.76|1.396x|1.098x|1.387x|
|H6144 res T24 w BF16|5.33|4.27|5.30|4.02|1.327x|1.062x|1.319x|
|H6144 res T32 w BF16|5.60|4.34|5.62|4.24|1.321x|1.024x|1.325x|
|H6144 res T48 w BF16|5.65|4.38|5.73|4.59|1.230x|0.954x|1.247x|
|H6144 res T64 w BF16|6.02|4.72|6.08|4.80|1.253x|0.983x|1.267x|
|H1536 plain T1 w FP32|4.67|1.68|7.28|3.28|1.424x|0.512x|2.220x|
|H1536 plain T2 w FP32|4.70|3.60|7.41|3.30|1.427x|1.091x|2.248x|
|H1536 plain T3 w FP32|4.50|3.63|7.52|3.42|1.313x|1.061x|2.196x|
|H1536 plain T4 w FP32|4.75|3.76|7.42|3.31|1.435x|1.136x|2.242x|
|H1536 plain T8 w FP32|4.51|3.78|7.68|3.42|1.318x|1.105x|2.243x|
|H1536 plain T12 w FP32|4.83|3.90|7.65|3.57|1.354x|1.092x|2.143x|
|H1536 plain T16 w FP32|4.70|3.92|7.55|3.46|1.361x|1.133x|2.185x|
|H1536 plain T24 w FP32|4.77|3.94|7.58|3.62|1.319x|1.088x|2.097x|
|H1536 plain T32 w FP32|4.64|3.55|7.50|3.55|1.306x|1.000x|2.113x|
|H1536 plain T48 w FP32|4.69|3.57|7.52|3.74|1.252x|0.955x|2.009x|
|H1536 plain T64 w FP32|4.70|3.62|7.58|3.70|1.273x|0.978x|2.052x|
|H2048 plain T1 w FP32|4.40|1.62|8.88|3.18|1.382x|0.509x|2.789x|
|H2048 plain T2 w FP32|4.42|3.04|9.09|3.18|1.387x|0.956x|2.854x|
|H2048 plain T3 w FP32|4.32|3.18|8.85|3.14|1.378x|1.013x|2.822x|
|H2048 plain T4 w FP32|4.32|3.26|8.91|3.15|1.371x|1.035x|2.827x|
|H2048 plain T8 w FP32|4.43|3.38|9.25|3.31|1.338x|1.021x|2.792x|
|H2048 plain T12 w FP32|4.43|3.38|9.30|3.33|1.332x|1.015x|2.793x|
|H2048 plain T16 w FP32|4.42|3.39|9.23|3.39|1.302x|1.000x|2.722x|
|H2048 plain T24 w FP32|4.40|3.44|9.26|3.41|1.291x|1.009x|2.719x|
|H2048 plain T32 w FP32|4.45|3.49|9.23|3.54|1.258x|0.986x|2.611x|
|H2048 plain T48 w FP32|4.45|3.71|9.20|3.84|1.158x|0.966x|2.396x|
|H2048 plain T64 w FP32|4.48|3.81|9.25|4.10|1.094x|0.929x|2.258x|
|H4096 plain T1 w FP32|5.28|1.94|12.38|3.74|1.410x|0.519x|3.307x|
|H4096 plain T2 w FP32|5.17|3.28|11.92|3.76|1.374x|0.872x|3.170x|
|H4096 plain T3 w FP32|5.10|3.42|12.16|3.81|1.341x|0.898x|3.193x|
|H4096 plain T4 w FP32|5.17|3.44|11.74|3.74|1.380x|0.920x|3.137x|
|H4096 plain T8 w FP32|5.25|3.60|11.74|3.89|1.350x|0.925x|3.021x|
|H4096 plain T12 w FP32|5.26|3.60|11.68|3.90|1.348x|0.923x|2.992x|
|H4096 plain T16 w FP32|5.25|3.63|11.76|3.94|1.333x|0.921x|2.988x|
|H4096 plain T24 w FP32|5.30|3.81|11.81|4.10|1.293x|0.929x|2.883x|
|H4096 plain T32 w FP32|5.23|3.84|11.84|4.21|1.243x|0.912x|2.813x|
|H4096 plain T48 w FP32|5.33|3.90|11.92|4.46|1.194x|0.874x|2.670x|
|H4096 plain T64 w FP32|5.34|4.00|11.84|4.74|1.128x|0.844x|2.500x|
|H5120 plain T1 w FP32|5.81|2.16|12.38|3.54|1.643x|0.610x|3.503x|
|H5120 plain T2 w FP32|6.06|3.94|12.30|3.55|1.707x|1.110x|3.463x|
|H5120 plain T3 w FP32|5.86|4.24|12.51|3.63|1.612x|1.168x|3.445x|
|H5120 plain T4 w FP32|6.03|4.27|12.54|3.60|1.676x|1.186x|3.484x|
|H5120 plain T8 w FP32|6.21|4.22|12.48|3.71|1.672x|1.137x|3.362x|
|H5120 plain T12 w FP32|6.14|4.38|12.50|3.86|1.593x|1.135x|3.241x|
|H5120 plain T16 w FP32|5.98|4.10|5.92|3.89|1.539x|1.054x|1.523x|
|H5120 plain T24 w FP32|5.97|4.58|6.08|4.29|1.392x|1.068x|1.418x|
|H5120 plain T32 w FP32|5.82|4.06|5.79|4.29|1.358x|0.946x|1.351x|
|H5120 plain T48 w FP32|5.86|4.11|5.86|4.66|1.258x|0.882x|1.258x|
|H5120 plain T64 w FP32|5.92|4.27|5.95|4.67|1.267x|0.914x|1.274x|
|H6144 plain T1 w FP32|5.31|2.10|13.58|3.38|1.573x|0.621x|4.024x|
|H6144 plain T2 w FP32|5.07|3.81|13.31|3.57|1.422x|1.067x|3.731x|
|H6144 plain T3 w FP32|5.28|3.87|13.68|3.42|1.542x|1.132x|3.994x|
|H6144 plain T4 w FP32|5.22|3.89|13.57|3.52|1.482x|1.105x|3.855x|
|H6144 plain T8 w FP32|5.23|4.05|13.63|3.54|1.479x|1.144x|3.854x|
|H6144 plain T12 w FP32|5.17|4.11|13.54|3.60|1.436x|1.142x|3.760x|
|H6144 plain T16 w FP32|5.18|4.19|13.63|3.70|1.403x|1.132x|3.688x|
|H6144 plain T24 w FP32|5.22|4.24|5.20|3.94|1.325x|1.076x|1.321x|
|H6144 plain T32 w FP32|5.33|4.26|5.30|4.18|1.276x|1.019x|1.268x|
|H6144 plain T48 w FP32|5.52|4.38|5.39|4.56|1.210x|0.961x|1.182x|
|H6144 plain T64 w FP32|5.46|4.51|5.50|4.62|1.180x|0.976x|1.190x|
|H6144 res T1 w FP32|5.06|2.67|5.09|3.52|1.436x|0.759x|1.445x|
|H6144 res T2 w FP32|4.99|3.86|5.12|3.58|1.393x|1.078x|1.429x|
|H6144 res T3 w FP32|5.18|4.03|5.06|3.70|1.403x|1.089x|1.368x|
|H6144 res T4 w FP32|5.17|4.06|5.14|3.50|1.475x|1.160x|1.466x|
|H6144 res T8 w FP32|5.14|4.16|5.12|3.71|1.384x|1.121x|1.379x|
|H6144 res T12 w FP32|5.12|4.27|5.17|3.76|1.362x|1.136x|1.374x|
|H6144 res T16 w FP32|5.20|4.26|5.12|3.79|1.371x|1.124x|1.350x|
|H6144 res T24 w FP32|5.15|4.32|5.15|4.00|1.288x|1.080x|1.288x|
|H6144 res T32 w FP32|5.44|4.38|5.47|4.14|1.313x|1.058x|1.320x|
|H6144 res T48 w FP32|5.57|4.54|5.55|4.59|1.212x|0.989x|1.209x|
|H6144 res T64 w FP32|5.84|4.66|5.86|4.66|1.255x|1.000x|1.258x|
Cake ships the forms listed above in both modes on sm_100a, sm_103a and
sm_107a. This PR changes only the `deterministic=False` backward of the
12 routes above, on all three architectures.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->
## Summary by CodeRabbit
* **New Features**
* Expanded RMSNorm training coverage across additional GPU architectures
and input widths.
* Added an alternative low-token processing path designed to produce the
same input and weight gradients as the existing approach.
* **Updates**
* Revised GPU-specific workload ranges and execution settings for
RMSNorm training.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Co-authored-by: Yingyi Huang <averyh@nvidia.com>
…8827)
Related issue: https://github.com/flashinfer-ai/flashinfer/issues/8066
`flashinfer.cake_rmsnorm_train` is RMSNorm for training. The forward
returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and
dw (FP32). The two modes, `deterministic=True` and
`deterministic=False`, share the forward and differ only in the
backward. Both modes support:
- [T, H] rows with BF16 or FP32 weight
- a fused residual add (H6144)
- per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a
shared w[256]
This PR follows #8577, #8626, #8655, #8704 and #8810. It changes only
the `deterministic=False` backward routes on sm_100a (B200) and sm_107a
(R200). The changed windows are listed below by reason. Rows where new
is slower than det are shown as measured.
Against det, bwd ranges from 0.77x to 1.36x on B200 and from 0.57x to
1.10x on R200. new is slower than det on 56 of 93 B200 bwd rows and on
69 of 86 R200 bwd rows. Against the strongest in-budget non-Cake SOTA,
bwd ranges from 0.43x to 1.58x on B200 and from 0.23x to 1.58x on R200.
These ranges and counts cover the rows whose program this PR changes.
They leave out the [kept] rows and include rows marked ~.
**Changed windows by reason**
The dw budget is per element atol 1e-4 + rtol 1e-4 vs FP64. The FP32
eager reference is the dw error of the FP32 eager computation on the
same inputs. A program fails it when its worst dw error exceeds the
reference's (FP32 weight) or its median dw error exceeds the reference's
median (either weight dtype). Accuracy data from 60 inputs of each of
three input families are not yet qualifying evidence.
B200 (sm_100a):
- Budget breach:
- H1536 above T 78784 with BF16 weight and above T 75776 with FP32
weight ([fix]): FlashInfer main's program is faster but exceeds the dw
budget at a larger T. Below each cut, that program has qualifying dw
evidence. The TwoSum twin runs above the cut.
- H5120 above T 56256 ([twin]): the deterministic=True backward's pair
that this PR replaces up to T 65536 exceeds the dw budget. The
FP32-weight pair computes the same dw. Above T 65536, the replaced
program exceeds the budget (measured with FP32 weight; the BF16-weight
program computes the same dw).
- H6144 above T 65536: the replaced program exceeds the dw budget
(measured with FP32 weight; the BF16-weight program computes the same
dw).
- Eager-reference failure ([route], FP32 weight): H1536 T 1025 to 4096,
H2048 T 337 to 4096, H4096 T 1793 to 3584 and H5120 T 65 to 6144. The
pair this PR replaces exceeds the FP32 eager reference's dw error at
some T of each window. det runs the same pair there (H5120: from T 122;
at T 65 to 121, det runs new's kernel with the same launch).
- H1536 T 65 to 384 ([edge], [cell]) runs a column-task kernel that is
faster than the kernels it replaces. With FP32 weight, the pair it
replaces exceeds the FP32 eager reference's dw error at T 64, and at T
65 to 109 det runs the same column-task kernel, so det and new run the
same backward program at T 65. With BF16 weight, at T 65 to 109 det runs
its own column-task kernel; the one slower row there (2xfwd+bwd at T 65,
0.98x) is marked ~.
- Neutral re-routes:
- H2048 T 45057 to 58880 ([neutral], both weight dtypes) runs the single
kernel of the windows on either side, within 2 % of the time of the
deterministic=True backward's pair that FlashInfer main ran there. The
pair's worst dw element reached 0.969 (BF16 weight) and 0.900 (FP32
weight) of the dw budget, the single kernel's 0.800 and 0.777 (every
64th T, 120 inputs per family).
- H4096 and residual H6144 above T 65536 run the TwoSum program at 1.00x
to 1.01x the time of the replaced program. The replaced program's worst
dw element reached 0.987 of the dw budget for H4096 (FP32 weight 0.907)
and 0.903 for residual H6144 (FP32 weight 0.835), on 60 inputs per
family.
- Qualified-evidence choice with a follow-up PR: H2048 above T 65536
(both weight dtypes) runs the TwoSum program at 1.10x to 1.14x the time
of the program it replaces. The TwoSum program's worst dw error is at
most 0.0006 of the budget. The replaced program stays within the dw
budget at every sampled T, but only on 60 inputs per family, so its data
do not yet qualify.
- Where new is slower than det in the H1536, H2048, H4096, H5120 and
H6144 top windows, the program det runs there is faster but exceeds the
dw budget at some measured T.
- A follow-up PR is planned to replace the TwoSum program in the B200
top windows up to T 131072.
- H512 T 1 to 227 runs the deterministic=True backward's column-task
kernel, so det and new run the same backward program there.
R200 (sm_107a):
- Top windows (TwoSum program): H512 above T 49152, H1536 and H4096
above T 65536, H6144 above T 30720 (FP32 weight: T 28672) and residual
H6144 above T 7168.
- Budget breach: the replaced program exceeds the dw budget at some
measured T of the window for H512, H1536 and H4096 (both weight dtypes),
and for H6144 and residual H6144 (measured with FP32 weight; the
BF16-weight program computes the same dw).
- In the H6144 and residual H6144 top windows, new is also faster than
the replaced program.
- Where new is slower than det, the deterministic=True backward's
program is faster there but exceeds the dw budget at some measured T of
the window. Each window keeps its start, and a follow-up PR re-checks
the faster program above it up to its budget edge.
- Eager-reference failure (FP32 weight):
- Windows: H256 T 1297 to 1536, H1536 T 65 to 256 and T 1025 to 16384,
H2048 T 1025 to 16384, H5120 T 4097 to 19968, H6144 T 65 to 28672 and
residual H6144 T 65 to 96. In each, the pair this PR replaces exceeds
the FP32 eager reference's dw error at some T (residual H6144: at T 8,
below the window). The deterministic=True backward runs the same pair
there (H6144: up to T 19456), except on H256, where it runs its own
pair, which also exceeds the reference.
- H6144 T 19457 to 28672 runs the deterministic=True backward's pair of
that range.
- On H6144 T 97 to 2304, the window uses the column-task program, whose
worst and median dw error stay at or below the reference's at every
sampled T. The deterministic=True backward's pair is faster there; its
data on this range come from 60 inputs per family, and a follow-up PR
re-checks it.
- FP32-weight H2048 T 65 to 127 and H4096 T 65 to 256 run column-task
cells that replace FlashInfer main's slower single kernel. The
deterministic=True backward's pair is faster there but exceeds the FP32
eager reference's dw error in these windows.
- H512 T 1 to 77 (FP32 weight: T 1 to 229) runs the TwoSum column-task
kernel of that hidden size, which is faster than the single kernel it
replaces. The deterministic=True backward's pair is faster there, but
its dw error exceeds the FP32 eager reference's: the median at T 77 with
BF16 weight, the worst at T 229 with FP32 weight.
- Neutral re-route: H6144 T 19457 to 30720 with BF16 weight runs the
deterministic=True backward's pair for now. The pair it replaces stays
within the dw budget there and takes 0.97x to 0.99x of the det pair's
time. A follow-up PR moves the range back to it.
**Summary over all rows this PR changes**
For each GPU and op, the summary gives:
- the min-max of the printed ratios
- the number of rows < 1 by printed ratio ('1.00x' does not count as <
1)
The [kept] rows are not counted. SOTA/new covers only rows with an
in-budget SOTA (! rows are left out). The ranges include rows marked ~.
|GPU / arch|op|rows|det/new min-max|det/new < 1|SOTA/new
min-max|SOTA/new < 1 of n|! rows|
|---|---|---|---|---|---|---|---|
|B200 sm_100a|fwd|93|0.98x-1.02x|9|0.84x-1.26x|38 of 93|0|
|B200 sm_100a|bwd|93|0.77x-1.36x|56|0.43x-1.58x|35 of 86|7|
|B200 sm_100a|2xfwd+bwd|93|0.82x-1.18x|57|0.65x-1.37x|23 of 86|7|
|R200 sm_107a, 212 SMs|fwd|86|0.91x-1.04x|10|0.74x-1.21x|36 of 86|0|
|R200 sm_107a, 212 SMs|bwd|86|0.57x-1.10x|69|0.23x-1.58x|53 of 85|1|
|R200 sm_107a, 212 SMs|2xfwd+bwd|86|0.70x-1.04x|70|0.46x-1.38x|46 of
85|1|
**Measurement**
- Latency: the median of the CUPTI span of one CUDA-graph replay with
cold L2. 4 sessions run in parallel, one per GPU; each session times 201
replays per implementation and op, split over 3 interleaved rounds; the
median pools the samples of the 4 sessions. The p10-p90 ranges are
omitted here. Each implementation's own setup, reduction, cast and copy
kernels are inside the span.
- Context (figures measured on GB300 and B200, not on R200): Cake gains
somewhat more from the context than the other non-compile libraries on
small rows: SOTA/new vs the non-compile candidates is up to about 1.10x
higher (GB300) / 1.14x higher (B200) in the one-process measurement than
in a process without torch.compile; large rows <= 1.02x. So small-row
ratios may overstate the new kernel's lead by up to about that much.
Ratios vs det are unaffected.
- GPU, CUDA version and build are given in the caption of each table.
- Workload: eps 1e-6. For every implementation, y, dx and h_new are
BF16, and rstd and dw are FP32.
- det = deterministic Cake: `deterministic=True` of the same tree, which
this PR does not change. new = this PR's non-deterministic Cake
(`deterministic=False`).
- SOTA = the fastest non-Cake candidate that has these outputs and meets
the accuracy budget, so it is the strongest in-budget comparison. The
budget is dx atol = rtol = 1e-2, and dw per element atol 1e-4 + rtol
1e-4 vs FP64, taking the worst of 5 seeds.
- Candidates: QuACK 0.6.5 (heuristic config and per-shape autotuned),
torch.compile (max-autotune-no-cudagraphs and default mode),
Liger-Kernel 0.8.5, and Transformer Engine (run as the FP32 problem).
- Not candidates: PyTorch's fused rms_norm (dtype error), Apex (not
installed) and FlashInfer's inference rmsnorm (forward only).
- In a B200 record, the torch.compile max-autotune candidate failed with
a disk-quota error on 2 rows. The record does not name them. On those
rows, SOTA is the fastest of the remaining candidates.
- Accuracy: new and det pass on every row.
- Speedup = baseline latency / new latency (> 1: new is faster). Ratios
are computed from unrounded medians, so a few differ in the last digit
from the ratio of the printed latencies. Both modes run the same forward
programs, and the printed fwd ratios are measured values.
- Rows: every measured row whose program this PR changes, plus the
[kept] rows. Rows where new is slower than det are shown as measured.
Key:
- form: plain, or res = fused residual add (h_new = BF16(h + u),
combined grad). op 2xfwd+bwd = checkpointed step (forward, forward
recompute, backward).
- Markers (B200):
- [fix] = window where the FlashInfer main program this PR replaces
exceeds the dw budget at a measured T (H1536 above T 78784 with BF16
weight, above T 75776 with FP32 weight). The TwoSum twin program runs
there.
- [neutral] = neutral re-route (H2048 T 45057 to 58880, both weight
dtypes), not a dw-budget fix. det runs the pair this PR replaces, so
det/new compares new with that pair.
- [twin] = the TwoSum twin program takes the top of the T range (H5120
above T 56256, residual H6144 above T 65536)
- [route] = FP32-weight re-route (eager-reference failure of the
replaced pair)
- [edge] = new T 65 edge (H1536; FP32-weight H5120)
- [cell] = new (64, 384] H1536 column-task cell
- [kept] = row in a window where this PR keeps FlashInfer main's
program, so the program at that T is unchanged. These rows are shown for
reference and are not counted in the summary or the ranges above.
- [cut] = row from a second record (H5120 T 56256 and 56257, H1536 T
75648)
- [carried] = row from an earlier record (same GPU model and software
stack) of a program that has not changed since
- R200: every row is a route change (at that T, this PR runs a different
program than FlashInfer main). For H6144 with FP32 weight (plain), every
measured T of the (64, 2304] window is shown.
- ! = no non-Cake candidate meets the dw budget. The fastest matched
candidate is shown, with its dw err/tol in parentheses. On B200,
SOTA/new is computed against it. On R200, no SOTA/new is given (-).
- ~ = the ratio fails the spread rule: the session medians of its
baseline (det or SOTA) or of new differ by > 5 % (largest minus smallest
of the 4 per-session medians, divided by the smallest). Session spread
(largest minus smallest of the 4 per-session medians, divided by the
smallest; largest of det / SOTA / new per row): B200 median 1.75 % (min
0.1 %, max 29.7 %), 40 of 300 rows marked ~; R200 median 1.6 % (min 0.2
%, max 18.3 %), 45 of 258 rows marked ~.
- SOTA names:
- TC = torch.compile (torch 2.14.1+cu130, inductor mode
max-autotune-no-cudagraphs, emulate_precision_casts=True) of the
explicit FP32-statistics reference
- TCd = the same with torch.compile inductor default mode
- Q = QuACK 0.6.5 (quack-kernels, nvidia-cutlass-dsl 4.8.1), heuristic
config
- QA = QuACK 0.6.5 rmsnorm_fwd_tuned / rmsnorm_bwd_tuned (exhaustive
per-shape autotune, cached before timing)
- Liger = Liger-Kernel 0.8.5 (Triton 3.8.0), casting_mode gemma (FP32
math)
- QA+add = torch add + QA (residual form)
**B200 (sm_100a): new non-deterministic Cake vs deterministic Cake vs
strongest non-Cake SOTA. Rows whose program this PR changes, plus [kept]
rows**
- GPU: B200 sm_100a, 148 SMs, compute capability 10.0; driver 580.82.07.
- Software: CUDA toolkit 13.4 (Build
cuda_13.4.r13.4/compiler.38855100_0); torch 2.14.1+cu130 (CUDA runtime
13.0).
- Sessions: 4 sessions (s1/s3 forward, s2/s4 reverse) run in parallel on
the node's 4 GPUs. Each row group runs in one process with det, all
non-Cake candidates and new, in a balanced implementation order.
- Builds: all rows were timed on earlier builds of this branch. On these
routes, those builds run the same programs as this PR's head in both
modes.
|workload [T, H] x, w, form|op|det us|SOTA us (name)|new
us|det/new|SOTA/new|
|---|---|---|---|---|---|---|
|[1, 512] BF16, w BF16, plain [carried]|fwd|1.89|1.76 TC|1.86|1.02x
~|0.95x ~|
|[1, 512] BF16, w BF16, plain [carried]|bwd|4.64|1.95 TCd|4.58|1.01x
~|0.43x ~|
|[1, 512] BF16, w BF16, plain [carried]|2xfwd+bwd|8.06|5.41
TCd|7.97|1.01x ~|0.68x ~|
|[2, 512] BF16, w BF16, plain [carried]|fwd|1.82|1.66 TCd|1.82|1.00x
~|0.91x|
|[2, 512] BF16, w BF16, plain [carried]|bwd|4.42|3.23 TCd|4.38|1.01x
~|0.74x ~|
|[2, 512] BF16, w BF16, plain [carried]|2xfwd+bwd|7.17|5.82
TCd|7.23|0.99x ~|0.81x ~|
|[8, 512] BF16, w BF16, plain [carried]|fwd|1.79|1.70 TCd|1.82|0.98x
~|0.93x ~|
|[8, 512] BF16, w BF16, plain [carried]|bwd|4.48|3.52 TCd|4.51|0.99x
~|0.78x ~|
|[8, 512] BF16, w BF16, plain [carried]|2xfwd+bwd|7.17|6.08
TCd|7.17|1.00x|0.85x|
|[16, 512] BF16, w BF16, plain [carried]|fwd|1.79|1.70 TCd|1.82|0.98x
~|0.93x ~|
|[16, 512] BF16, w BF16, plain [carried]|bwd|4.48|3.84 TCd|4.38|1.02x
~|0.88x ~|
|[16, 512] BF16, w BF16, plain [carried]|2xfwd+bwd|7.17|6.43
TCd|7.07|1.01x|0.91x|
|[64, 512] BF16, w BF16, plain [carried]|fwd|1.86|1.73
Liger|1.86|1.00x|0.93x|
|[64, 512] BF16, w BF16, plain [carried]|bwd|4.32|3.90 TC|4.61|0.94x
~|0.85x ~|
|[64, 512] BF16, w BF16, plain [carried]|2xfwd+bwd|7.23|6.56
TCd|7.26|1.00x|0.90x|
|[128, 512] BF16, w BF16, plain [carried]|fwd|1.89|1.79
TCd|1.89|1.00x|0.95x|
|[128, 512] BF16, w BF16, plain [carried]|bwd|4.45|4.13
TCd|4.38|1.01x|0.94x ~|
|[128, 512] BF16, w BF16, plain [carried]|2xfwd+bwd|7.33|6.75
TCd|7.36|1.00x|0.92x ~|
|[227, 512] BF16, w BF16, plain [carried]|fwd|1.95|1.82
TC|1.95|1.00x|0.93x|
|[227, 512] BF16, w BF16, plain [carried]|bwd|4.48|7.01
TC|4.51|0.99x|1.55x ~|
|[227, 512] BF16, w BF16, plain [carried]|2xfwd+bwd|8.35|10.85
TC|8.19|1.02x|1.32x|
|[65, 1536] BF16, w BF16, plain [edge]|fwd|1.89|1.82
Liger|1.89|1.00x|0.97x|
|[65, 1536] BF16, w BF16, plain [edge]|bwd|3.94|4.80
TC|3.87|1.02x|1.24x|
|[65, 1536] BF16, w BF16, plain [edge]|2xfwd+bwd|7.52|9.25 TC|7.65|0.98x
~|1.21x ~|
|[128, 1536] BF16, w BF16, plain [cell]|fwd|1.95|1.89
Liger|1.95|1.00x|0.97x|
|[128, 1536] BF16, w BF16, plain [cell]|bwd|5.76|4.35
TC|4.22|1.36x|1.03x|
|[128, 1536] BF16, w BF16, plain [cell]|2xfwd+bwd|8.51|7.58
TC|7.20|1.18x|1.05x|
|[256, 1536] BF16, w BF16, plain [cell]|fwd|2.05|2.02
Liger|2.05|1.00x|0.98x|
|[256, 1536] BF16, w BF16, plain [cell]|bwd|5.89|5.95
TC|4.58|1.29x|1.30x ~|
|[256, 1536] BF16, w BF16, plain [cell]|2xfwd+bwd|8.80|9.41
TC|7.78|1.13x|1.21x|
|[384, 1536] BF16, w BF16, plain [cell]|fwd|2.14|2.14
Liger|2.11|1.02x|1.02x|
|[384, 1536] BF16, w BF16, plain [cell]|bwd|6.02|6.21
TC|4.87|1.24x|1.28x|
|[384, 1536] BF16, w BF16, plain [cell]|2xfwd+bwd|9.09|10.56
TC|8.29|1.10x|1.27x|
|[75648, 1536] BF16, w BF16, plain [kept] [cut]|fwd|66.46|70.18
QA|66.46|1.00x|1.06x|
|[75648, 1536] BF16, w BF16, plain [kept] [cut]|bwd|117.79|133.02
TC|120.80|0.98x|1.10x|
|[75648, 1536] BF16, w BF16, plain [kept] [cut]|2xfwd+bwd|250.90|272.35
TC|254.91|0.98x|1.07x|
|[78784, 1536] BF16, w BF16, plain [kept]|fwd|69.41|72.19
TC|69.44|1.00x|1.04x|
|[78784, 1536] BF16, w BF16, plain [kept]|bwd|122.40|134.38
TC|120.99|1.01x|1.11x|
|[78784, 1536] BF16, w BF16, plain [kept]|2xfwd+bwd|262.18|277.66
TC|261.81|1.00x|1.06x|
|[78785, 1536] BF16, w BF16, plain [fix]|fwd|69.12|72.77
QA|69.28|1.00x|1.05x|
|[78785, 1536] BF16, w BF16, plain [fix]|bwd|122.50|134.02
TC|133.50|0.92x|1.00x|
|[78785, 1536] BF16, w BF16, plain [fix]|2xfwd+bwd|261.04|277.31
TC|272.48|0.96x|1.02x|
|[86272, 1536] BF16, w BF16, plain [fix]|fwd|75.65|78.97
TC|75.71|1.00x|1.04x|
|[86272, 1536] BF16, w BF16, plain [fix]|bwd|132.25|141.20
QA|144.35|0.92x|0.98x|
|[86272, 1536] BF16, w BF16, plain [fix]|2xfwd+bwd|284.19|296.69
QA|297.37|0.96x|1.00x|
|[98304, 1536] BF16, w BF16, plain [fix]|fwd|85.95|90.34
QA|85.98|1.00x|1.05x|
|[98304, 1536] BF16, w BF16, plain [fix]|bwd|149.02|157.44
QA|163.62|0.91x|0.96x|
|[98304, 1536] BF16, w BF16, plain [fix]|2xfwd+bwd|321.18|334.62
QA|336.89|0.95x|0.99x|
|[131072, 1536] BF16, w BF16, plain [fix]|fwd|114.88|119.26
QA|115.01|1.00x|1.04x|
|[131072, 1536] BF16, w BF16, plain [fix]|bwd|193.57|207.21
TC|214.21|0.90x|0.97x|
|[131072, 1536] BF16, w BF16, plain [fix]|2xfwd+bwd|423.55|443.44
TC|446.14|0.95x|0.99x|
|[262144, 1536] BF16, w BF16, plain [fix]|fwd|232.99|235.49
TC|233.28|1.00x|1.01x|
|[262144, 1536] BF16, w BF16, plain [fix]|bwd|371.26|385.60
TC|415.04|0.89x|0.93x|
|[262144, 1536] BF16, w BF16, plain [fix]|2xfwd+bwd|831.94|852.70
TC|881.49|0.94x|0.97x|
|[45057, 2048] BF16, w BF16, plain [neutral]|fwd|52.62|52.83
Q|52.58|1.00x|1.00x|
|[45057, 2048] BF16, w BF16, plain [neutral]|bwd|95.55|105.31
TC|94.94|1.01x|1.11x|
|[45057, 2048] BF16, w BF16, plain [neutral]|2xfwd+bwd|201.47|213.76
QA|201.47|1.00x|1.06x|
|[51968, 2048] BF16, w BF16, plain [neutral]|fwd|60.64|60.96
Q|60.67|1.00x|1.00x|
|[51968, 2048] BF16, w BF16, plain [neutral]|bwd|107.14|118.14
QA|107.76|0.99x|1.10x|
|[51968, 2048] BF16, w BF16, plain [neutral]|2xfwd+bwd|228.25|241.25
QA|229.66|0.99x|1.05x|
|[58880, 2048] BF16, w BF16, plain [neutral]|fwd|68.96|69.34
Q|68.96|1.00x|1.01x|
|[58880, 2048] BF16, w BF16, plain [neutral]|bwd|119.81|130.46
QA|121.44|0.99x|1.07x|
|[58880, 2048] BF16, w BF16, plain [neutral]|2xfwd+bwd|256.69|269.54
QA|259.62|0.99x|1.04x|
|[65537, 2048] BF16, w BF16, plain [carried]|fwd|76.77|76.98
Q|76.48|1.00x|1.01x|
|[65537, 2048] BF16, w BF16, plain [carried]|bwd|133.31|150.46
TC|149.66|0.89x|1.01x|
|[65537, 2048] BF16, w BF16, plain [carried]|2xfwd+bwd|285.82|308.16
TC|303.58|0.94x|1.02x|
|[98304, 2048] BF16, w BF16, plain [carried]|fwd|114.51|114.94
Q|114.56|1.00x|1.00x|
|[98304, 2048] BF16, w BF16, plain [carried]|bwd|192.29|206.75
TC|218.94|0.88x|0.94x|
|[98304, 2048] BF16, w BF16, plain [carried]|2xfwd+bwd|421.53|437.09
TC|451.44|0.93x|0.97x|
|[131072, 2048] BF16, w BF16, plain [carried]|fwd|155.94|154.69
QA|156.26|1.00x|0.99x|
|[131072, 2048] BF16, w BF16, plain [carried]|bwd|251.07|259.94
TC|288.35|0.87x|0.90x|
|[131072, 2048] BF16, w BF16, plain [carried]|2xfwd+bwd|555.85|565.63
TC|598.75|0.93x|0.94x|
|[65537, 4096] BF16, w BF16, plain [carried]|fwd|154.72|156.45
QA|155.01|1.00x|1.01x|
|[65537, 4096] BF16, w BF16, plain [carried]|bwd|252.94|267.07
TC|259.74|0.97x|1.03x|
|[65537, 4096] BF16, w BF16, plain [carried]|2xfwd+bwd|555.20|578.56
TC|567.46|0.98x|1.02x|
|[98304, 4096] BF16, w BF16, plain [carried]|fwd|231.28|233.63
QA|231.13|1.00x|1.01x|
|[98304, 4096] BF16, w BF16, plain [carried]|bwd|371.33|395.44
TC|385.44|0.96x|1.03x|
|[98304, 4096] BF16, w BF16, plain [carried]|2xfwd+bwd|826.72|859.90
TC|844.89|0.98x|1.02x|
|[131072, 4096] BF16, w BF16, plain [carried]|fwd|306.82|308.83
TC|307.26|1.00x|1.01x|
|[131072, 4096] BF16, w BF16, plain [carried]|bwd|486.49|496.67
TC|511.20|0.95x|0.97x|
|[131072, 4096] BF16, w BF16, plain [carried]|2xfwd+bwd|1094.30|1114.77
TC|1122.66|0.97x|0.99x|
|[56256, 5120] BF16, w BF16, plain [kept] [cut]|fwd|170.37|168.61
TC|170.43|1.00x|0.99x|
|[56256, 5120] BF16, w BF16, plain [kept] [cut]|bwd|266.18|402.80
TCd|265.90|1.00x|1.51x|
|[56256, 5120] BF16, w BF16, plain [kept] [cut]|2xfwd+bwd|601.95|747.00
TCd|601.63|1.00x|1.24x|
|[56257, 5120] BF16, w BF16, plain [twin] [cut]|fwd|170.37|168.74
TC|170.50|1.00x|0.99x|
|[56257, 5120] BF16, w BF16, plain [twin] [cut]|bwd|265.89|402.08
TCd|332.57|0.80x|1.21x|
|[56257, 5120] BF16, w BF16, plain [twin] [cut]|2xfwd+bwd|601.15|745.92
TCd|671.07|0.90x|1.11x|
|[57345, 5120] BF16, w BF16, plain [twin]|fwd|173.86|172.03
TC|173.98|1.00x|0.99x|
|[57345, 5120] BF16, w BF16, plain [twin]|bwd|273.73|357.44
TC|338.16|0.81x|1.06x|
|[57345, 5120] BF16, w BF16, plain [twin]|2xfwd+bwd|616.03|699.10
TC|684.27|0.90x|1.02x|
|[98304, 5120] BF16, w BF16, plain [twin]|fwd|294.53|292.25
QA|294.56|1.00x|0.99x|
|[98304, 5120] BF16, w BF16, plain [twin]|bwd|454.94|582.48
TC|570.24|0.80x|1.02x|
|[98304, 5120] BF16, w BF16, plain [twin]|2xfwd+bwd|1040.56|1190.38
TC|1157.50|0.90x|1.03x|
|[131072, 5120] BF16, w BF16, plain [twin]|fwd|391.30|388.25
QA|391.04|1.00x|0.99x|
|[131072, 5120] BF16, w BF16, plain [twin]|bwd|613.34|604.35 QA!
(2.650)|754.75|0.81x|0.80x!|
|[131072, 5120] BF16, w BF16, plain [twin]|2xfwd+bwd|1405.58|1381.82 QA!
(2.650)|1547.68|0.91x|0.89x!|
|[65537, 6144] BF16, w BF16, plain [carried]|fwd|232.00|231.90
QA|232.16|1.00x|1.00x|
|[65537, 6144] BF16, w BF16, plain [carried]|bwd|366.72|411.01
TC|394.11|0.93x|1.04x|
|[65537, 6144] BF16, w BF16, plain [carried]|2xfwd+bwd|823.87|872.73
TC|853.85|0.96x|1.02x|
|[98304, 6144] BF16, w BF16, plain [carried]|fwd|346.30|345.31
QA|346.43|1.00x|1.00x|
|[98304, 6144] BF16, w BF16, plain [carried]|bwd|548.32|602.40
TC|584.06|0.94x|1.03x|
|[98304, 6144] BF16, w BF16, plain [carried]|2xfwd+bwd|1235.84|1344.33
TC|1288.17|0.96x|1.04x|
|[131072, 6144] BF16, w BF16, plain [carried]|fwd|459.97|458.53
Q|460.16|1.00x|1.00x|
|[131072, 6144] BF16, w BF16, plain [carried]|bwd|718.64|708.06 QA!
(2.786)|775.77|0.93x|0.91x!|
|[131072, 6144] BF16, w BF16, plain [carried]|2xfwd+bwd|1634.81|1646.78
QA! (2.786)|1740.48|0.94x|0.95x!|
|[65537, 6144] BF16, w BF16, res [twin]|fwd|468.38|464.75
Liger|468.46|1.00x|0.99x|
|[65537, 6144] BF16, w BF16, res [twin]|bwd|580.00|499.13
TC|512.46|1.13x|0.97x|
|[65537, 6144] BF16, w BF16, res [twin]|2xfwd+bwd|1519.81|1443.58
TC|1439.36|1.06x|1.00x|
|[98304, 6144] BF16, w BF16, res [twin]|fwd|700.38|696.43
Liger|700.16|1.00x|0.99x|
|[98304, 6144] BF16, w BF16, res [twin]|bwd|866.05|762.62
TCd|766.88|1.13x|0.99x|
|[98304, 6144] BF16, w BF16, res [twin]|2xfwd+bwd|2282.25|2198.98
TCd|2159.67|1.06x|1.02x|
|[131072, 6144] BF16, w BF16, res [twin]|fwd|932.45|929.57
Liger|932.32|1.00x|1.00x|
|[131072, 6144] BF16, w BF16, res [twin]|bwd|1146.88|982.43 TC!
(1.241)|1005.97|1.14x|0.98x!|
|[131072, 6144] BF16, w BF16, res [twin]|2xfwd+bwd|3048.06|2931.69 TCd!
(1.241)|2882.10|1.06x|1.02x!|
|[1, 512] BF16, w FP32, plain [carried]|fwd|2.02|1.73
TCd|2.05|0.98x|0.84x|
|[1, 512] BF16, w FP32, plain [carried]|bwd|4.10|1.92
TC|4.06|1.01x|0.47x|
|[1, 512] BF16, w FP32, plain [carried]|2xfwd+bwd|7.30|4.80
TC|7.39|0.99x ~|0.65x ~|
|[2, 512] BF16, w FP32, plain [carried]|fwd|1.79|1.63 Liger|1.79|1.00x
~|0.91x ~|
|[2, 512] BF16, w FP32, plain [carried]|bwd|4.10|3.20 TCd|4.09|1.00x
~|0.78x ~|
|[2, 512] BF16, w FP32, plain [carried]|2xfwd+bwd|7.04|5.79
TCd|7.01|1.00x|0.83x|
|[8, 512] BF16, w FP32, plain [carried]|fwd|1.79|1.70
Liger|1.79|1.00x|0.95x|
|[8, 512] BF16, w FP32, plain [carried]|bwd|4.26|3.42 TC|4.06|1.05x
~|0.84x ~|
|[8, 512] BF16, w FP32, plain [carried]|2xfwd+bwd|7.07|6.02
TC|7.01|1.01x|0.86x|
|[16, 512] BF16, w FP32, plain [carried]|fwd|1.79|1.70
TCd|1.79|1.00x|0.95x|
|[16, 512] BF16, w FP32, plain [carried]|bwd|4.16|3.74
TC|4.06|1.02x|0.92x ~|
|[16, 512] BF16, w FP32, plain [carried]|2xfwd+bwd|7.14|6.43
TCd|7.04|1.01x|0.91x|
|[64, 512] BF16, w FP32, plain [carried]|fwd|1.82|1.73
TCd|1.82|1.00x|0.95x|
|[64, 512] BF16, w FP32, plain [carried]|bwd|4.06|3.81 TC|4.19|0.97x
~|0.91x ~|
|[64, 512] BF16, w FP32, plain [carried]|2xfwd+bwd|7.01|6.50
TCd|7.04|1.00x|0.92x ~|
|[128, 512] BF16, w FP32, plain [carried]|fwd|1.89|1.76
TC|1.89|1.00x|0.93x|
|[128, 512] BF16, w FP32, plain [carried]|bwd|4.29|4.38 TCd|4.22|1.02x
~|1.04x ~|
|[128, 512] BF16, w FP32, plain [carried]|2xfwd+bwd|7.71|7.84
TC|7.81|0.99x ~|1.00x|
|[227, 512] BF16, w FP32, plain [carried]|fwd|1.92|1.82
TC|1.95|0.98x|0.93x|
|[227, 512] BF16, w FP32, plain [carried]|bwd|4.38|6.82 TC|4.32|1.01x
~|1.58x ~|
|[227, 512] BF16, w FP32, plain [carried]|2xfwd+bwd|7.84|10.72
TC|8.13|0.96x ~|1.32x ~|
|[65, 1536] BF16, w FP32, plain [edge]|fwd|1.86|1.82
Liger|1.89|0.98x|0.97x|
|[65, 1536] BF16, w FP32, plain [edge]|bwd|3.87|4.32 TC|3.94|0.98x|1.10x
~|
|[65, 1536] BF16, w FP32, plain [edge]|2xfwd+bwd|6.88|7.17
TC|6.82|1.01x|1.05x ~|
|[128, 1536] BF16, w FP32, plain [cell]|fwd|1.92|1.89
Liger|1.95|0.98x|0.97x|
|[128, 1536] BF16, w FP32, plain [cell]|bwd|5.70|4.77
TC|4.22|1.35x|1.13x|
|[128, 1536] BF16, w FP32, plain [cell]|2xfwd+bwd|8.48|7.93
TC|7.26|1.17x|1.09x|
|[256, 1536] BF16, w FP32, plain [cell]|fwd|2.02|2.05
Liger|2.02|1.00x|1.02x|
|[256, 1536] BF16, w FP32, plain [cell]|bwd|5.86|6.78
TC|4.58|1.28x|1.48x ~|
|[256, 1536] BF16, w FP32, plain [cell]|2xfwd+bwd|9.54|11.78
TC|8.80|1.08x|1.34x|
|[384, 1536] BF16, w FP32, plain [cell]|fwd|2.11|2.24
Liger|2.11|1.00x|1.06x|
|[384, 1536] BF16, w FP32, plain [cell]|bwd|5.92|6.67
TC|4.93|1.20x|1.35x ~|
|[384, 1536] BF16, w FP32, plain [cell]|2xfwd+bwd|9.54|12.32
TC|9.28|1.03x ~|1.33x ~|
|[1025, 1536] BF16, w FP32, plain [route]|fwd|2.59|2.72
Q|2.59|1.00x|1.05x|
|[1025, 1536] BF16, w FP32, plain [route]|bwd|7.55|10.13
TC|9.44|0.80x|1.07x|
|[1025, 1536] BF16, w FP32, plain [route]|2xfwd+bwd|12.06|15.94
TC|14.40|0.84x|1.11x|
|[2048, 1536] BF16, w FP32, plain [route]|fwd|3.36|3.55
Q|3.36|1.00x|1.06x|
|[2048, 1536] BF16, w FP32, plain [route]|bwd|10.34|10.75
TC|11.42|0.90x|0.94x ~|
|[2048, 1536] BF16, w FP32, plain [route]|2xfwd+bwd|16.13|18.24
TC|17.63|0.91x|1.03x|
|[3136, 1536] BF16, w FP32, plain [route]|fwd|4.45|4.48
QA|4.45|1.00x|1.01x|
|[3136, 1536] BF16, w FP32, plain [route]|bwd|12.71|16.64 TC|12.70|1.00x
~|1.31x|
|[3136, 1536] BF16, w FP32, plain [route]|2xfwd+bwd|20.13|24.96
TC|20.35|0.99x|1.23x|
|[4096, 1536] BF16, w FP32, plain [route]|fwd|4.86|5.47
TC|4.86|1.00x|1.13x ~|
|[4096, 1536] BF16, w FP32, plain [route]|bwd|13.25|17.87
TC|14.66|0.90x|1.22x|
|[4096, 1536] BF16, w FP32, plain [route]|2xfwd+bwd|20.90|26.85
TC|23.04|0.91x|1.17x|
|[65536, 1536] BF16, w FP32, plain [kept]|fwd|59.18|61.73
TC|59.17|1.00x|1.04x|
|[65536, 1536] BF16, w FP32, plain [kept]|bwd|103.26|108.96
QA|100.18|1.03x|1.09x|
|[65536, 1536] BF16, w FP32, plain [kept]|2xfwd+bwd|222.66|236.89
TC|219.36|1.02x|1.08x|
|[75648, 1536] BF16, w FP32, plain [kept] [cut]|fwd|68.26|70.75
TC|68.38|1.00x|1.03x|
|[75648, 1536] BF16, w FP32, plain [kept] [cut]|bwd|117.28|128.90
TC|117.86|1.00x|1.09x|
|[75648, 1536] BF16, w FP32, plain [kept] [cut]|2xfwd+bwd|254.78|268.61
TC|255.90|1.00x|1.05x|
|[75776, 1536] BF16, w FP32, plain [kept]|fwd|68.45|71.39
TC|68.38|1.00x|1.04x|
|[75776, 1536] BF16, w FP32, plain [kept]|bwd|117.89|122.78
QA|118.46|1.00x|1.04x|
|[75776, 1536] BF16, w FP32, plain [kept]|2xfwd+bwd|255.14|265.98
TC|255.94|1.00x|1.04x|
|[75777, 1536] BF16, w FP32, plain [fix]|fwd|68.38|71.33
TC|68.29|1.00x|1.04x|
|[75777, 1536] BF16, w FP32, plain [fix]|bwd|118.75|123.14
QA|128.19|0.93x|0.96x|
|[75777, 1536] BF16, w FP32, plain [fix]|2xfwd+bwd|255.81|266.01
TC|264.83|0.97x|1.00x|
|[98304, 1536] BF16, w FP32, plain [fix]|fwd|88.74|91.36
TC|88.80|1.00x|1.03x|
|[98304, 1536] BF16, w FP32, plain [fix]|bwd|148.58|154.37
QA|162.85|0.91x|0.95x|
|[98304, 1536] BF16, w FP32, plain [fix]|2xfwd+bwd|326.11|343.26
TC|340.51|0.96x|1.01x|
|[131072, 1536] BF16, w FP32, plain [fix]|fwd|118.08|120.77
TC|117.95|1.00x|1.02x|
|[131072, 1536] BF16, w FP32, plain [fix]|bwd|193.63|207.20
TC|213.12|0.91x|0.97x|
|[131072, 1536] BF16, w FP32, plain [fix]|2xfwd+bwd|429.92|446.08
TC|449.57|0.96x|0.99x|
|[337, 2048] BF16, w FP32, plain [route]|fwd|2.21|2.30
QA|2.21|1.00x|1.04x|
|[337, 2048] BF16, w FP32, plain [route]|bwd|6.59|7.68
TC|5.63|1.17x|1.36x|
|[337, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|9.89|11.65
TC|9.12|1.08x|1.28x|
|[640, 2048] BF16, w FP32, plain [route]|fwd|2.50|2.59
QA|2.50|1.00x|1.04x|
|[640, 2048] BF16, w FP32, plain [route]|bwd|8.32|7.70
TC|8.64|0.96x|0.89x|
|[640, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|12.22|12.26
TC|12.54|0.97x|0.98x|
|[1024, 2048] BF16, w FP32, plain [route]|fwd|2.78|2.94
QA|2.78|1.00x|1.06x|
|[1024, 2048] BF16, w FP32, plain [route]|bwd|8.51|9.44
TC|9.95|0.86x|0.95x|
|[1024, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|13.79|15.33
TC|15.26|0.90x|1.00x|
|[1025, 2048] BF16, w FP32, plain [route]|fwd|3.20|2.98
TC|3.20|1.00x|0.93x|
|[1025, 2048] BF16, w FP32, plain [route]|bwd|8.51|10.34 TC|10.08|0.84x
~|1.03x ~|
|[1025, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|14.34|16.29
TC|15.62|0.92x|1.04x|
|[1728, 2048] BF16, w FP32, plain [route]|fwd|4.06|3.75
QA|4.10|0.99x|0.91x|
|[1728, 2048] BF16, w FP32, plain [route]|bwd|11.65|11.78
TC|11.39|1.02x|1.03x|
|[1728, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|18.46|18.66
TC|18.40|1.00x|1.01x|
|[2773, 2048] BF16, w FP32, plain [route]|fwd|4.58|5.31
TC|4.61|0.99x|1.15x|
|[2773, 2048] BF16, w FP32, plain [route]|bwd|14.05|19.87
TC|12.90|1.09x|1.54x|
|[2773, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|22.21|28.74
TC|21.28|1.04x|1.35x|
|[3920, 2048] BF16, w FP32, plain [route]|fwd|5.79|6.56
QA|5.79|1.00x|1.13x|
|[3920, 2048] BF16, w FP32, plain [route]|bwd|17.82|24.48
QA|15.58|1.14x|1.57x|
|[3920, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|29.31|35.71
QA|27.26|1.08x|1.31x|
|[4096, 2048] BF16, w FP32, plain [route]|fwd|5.94|6.75
QA|5.95|1.00x|1.13x|
|[4096, 2048] BF16, w FP32, plain [route]|bwd|18.14|24.96
Q|16.16|1.12x|1.54x|
|[4096, 2048] BF16, w FP32, plain [route]|2xfwd+bwd|28.26|36.51
Q|26.59|1.06x|1.37x|
|[45057, 2048] BF16, w FP32, plain [neutral]|fwd|55.34|57.33
QA|55.30|1.00x|1.04x|
|[45057, 2048] BF16, w FP32, plain [neutral]|bwd|95.84|103.17
QA|94.46|1.01x|1.09x|
|[45057, 2048] BF16, w FP32, plain [neutral]|2xfwd+bwd|203.07|213.95
QA|201.85|1.01x|1.06x|
|[51968, 2048] BF16, w FP32, plain [neutral]|fwd|63.39|65.95
TC|63.33|1.00x|1.04x|
|[51968, 2048] BF16, w FP32, plain [neutral]|bwd|107.42|115.30
QA|106.94|1.00x|1.08x|
|[51968, 2048] BF16, w FP32, plain [neutral]|2xfwd+bwd|230.34|242.43
QA|230.32|1.00x|1.05x|
|[58880, 2048] BF16, w FP32, plain [neutral]|fwd|71.26|71.33
TC|71.41|1.00x|1.00x|
|[58880, 2048] BF16, w FP32, plain [neutral]|bwd|120.16|127.74
QA|120.88|0.99x|1.06x|
|[58880, 2048] BF16, w FP32, plain [neutral]|2xfwd+bwd|258.59|270.97
QA|260.48|0.99x|1.04x|
|[65537, 2048] BF16, w FP32, plain [carried]|fwd|78.82|81.12
TC|78.78|1.00x|1.03x|
|[65537, 2048] BF16, w FP32, plain [carried]|bwd|133.41|149.98
TC|150.02|0.89x|1.00x|
|[65537, 2048] BF16, w FP32, plain [carried]|2xfwd+bwd|287.09|308.19
TC|305.41|0.94x|1.01x|
|[98304, 2048] BF16, w FP32, plain [carried]|fwd|116.38|118.02
TC|116.67|1.00x|1.01x|
|[98304, 2048] BF16, w FP32, plain [carried]|bwd|192.35|206.34
TC|220.00|0.87x|0.94x|
|[98304, 2048] BF16, w FP32, plain [carried]|2xfwd+bwd|421.02|437.79
TC|450.65|0.93x|0.97x|
|[131072, 2048] BF16, w FP32, plain [carried]|fwd|155.04|156.00
TC|154.91|1.00x|1.01x|
|[131072, 2048] BF16, w FP32, plain [carried]|bwd|250.65|262.85
TC|289.60|0.87x|0.91x|
|[131072, 2048] BF16, w FP32, plain [carried]|2xfwd+bwd|554.30|570.75
TC|596.57|0.93x|0.96x|
|[1793, 4096] BF16, w FP32, plain [route]|fwd|6.50|6.50 QA|6.46|1.00x
~|1.00x ~|
|[1793, 4096] BF16, w FP32, plain [route]|bwd|15.36|18.91
QA|15.97|0.96x|1.18x|
|[1793, 4096] BF16, w FP32, plain [route]|2xfwd+bwd|28.06|31.10
QA|29.06|0.97x|1.07x|
|[2569, 4096] BF16, w FP32, plain [route]|fwd|8.19|7.91
TC|8.19|1.00x|0.96x ~|
|[2569, 4096] BF16, w FP32, plain [route]|bwd|18.85|22.21
QA|19.01|0.99x|1.17x|
|[2569, 4096] BF16, w FP32, plain [route]|2xfwd+bwd|35.07|38.05
QA|35.55|0.99x|1.07x|
|[3072, 4096] BF16, w FP32, plain [route]|fwd|9.57|9.50
TC|9.54|1.00x|1.00x ~|
|[3072, 4096] BF16, w FP32, plain [route]|bwd|21.06|24.16
QA|20.83|1.01x|1.16x|
|[3072, 4096] BF16, w FP32, plain [route]|2xfwd+bwd|39.78|41.34
QA|39.90|1.00x|1.04x|
|[3584, 4096] BF16, w FP32, plain [route]|fwd|11.26|11.39
QA|11.13|1.01x|1.02x|
|[3584, 4096] BF16, w FP32, plain [route]|bwd|23.17|26.37
QA|22.05|1.05x|1.20x|
|[3584, 4096] BF16, w FP32, plain [route]|2xfwd+bwd|44.19|45.44
QA|43.49|1.02x|1.04x|
|[65537, 4096] BF16, w FP32, plain [carried]|fwd|159.23|157.28
TC|159.38|1.00x|0.99x|
|[65537, 4096] BF16, w FP32, plain [carried]|bwd|252.27|269.63
TC|259.71|0.97x|1.04x|
|[65537, 4096] BF16, w FP32, plain [carried]|2xfwd+bwd|563.97|580.49
TC|575.18|0.98x|1.01x|
|[98304, 4096] BF16, w FP32, plain [carried]|fwd|236.22|234.98
TCd|236.38|1.00x|0.99x|
|[98304, 4096] BF16, w FP32, plain [carried]|bwd|370.54|467.01
TCd|385.18|0.96x|1.21x|
|[98304, 4096] BF16, w FP32, plain [carried]|2xfwd+bwd|838.14|934.96
TCd|857.39|0.98x|1.09x|
|[131072, 4096] BF16, w FP32, plain [carried]|fwd|313.09|311.23
TC|313.26|1.00x|0.99x|
|[131072, 4096] BF16, w FP32, plain [carried]|bwd|486.53|495.36
TC|510.24|0.95x|0.97x|
|[131072, 4096] BF16, w FP32, plain [carried]|2xfwd+bwd|1115.98|1121.12
TC|1141.13|0.98x|0.98x|
|[65, 5120] BF16, w FP32, plain [edge]|fwd|2.21|2.21
Liger|2.21|1.00x|1.00x|
|[65, 5120] BF16, w FP32, plain [edge]|bwd|4.13|4.99
TC|4.13|1.00x|1.21x|
|[65, 5120] BF16, w FP32, plain [edge]|2xfwd+bwd|7.81|8.67
TC|7.71|1.01x|1.12x|
|[77, 5120] BF16, w FP32, plain [route]|fwd|2.24|2.24
Liger|2.21|1.01x|1.01x|
|[77, 5120] BF16, w FP32, plain [route]|bwd|5.12|5.02
TC|5.12|1.00x|0.98x|
|[77, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|8.83|8.80
TC|8.86|1.00x|0.99x|
|[256, 5120] BF16, w FP32, plain [route]|fwd|2.72|2.78
QA|2.72|1.00x|1.02x|
|[256, 5120] BF16, w FP32, plain [route]|bwd|7.78|10.14
TCd|6.94|1.12x|1.46x|
|[256, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|12.58|15.55
TCd|12.00|1.05x|1.30x ~|
|[1024, 5120] BF16, w FP32, plain [route]|fwd|4.54|4.77
TCd|4.54|1.00x|1.05x|
|[1024, 5120] BF16, w FP32, plain [route]|bwd|12.86|16.99
QA|13.73|0.94x|1.24x|
|[1024, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|20.83|25.70
QA|22.14|0.94x|1.16x|
|[1792, 5120] BF16, w FP32, plain [route]|fwd|6.18|7.78
QA|6.18|1.00x|1.26x|
|[1792, 5120] BF16, w FP32, plain [route]|bwd|17.15|20.45
QA|22.27|0.77x|0.92x|
|[1792, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|28.02|34.56
QA|34.21|0.82x|1.01x|
|[1793, 5120] BF16, w FP32, plain [route]|fwd|6.21|7.81
QA|6.21|1.00x|1.26x|
|[1793, 5120] BF16, w FP32, plain [route]|bwd|17.28|20.61
QA|21.44|0.81x|0.96x|
|[1793, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|28.06|34.62
QA|33.18|0.85x|1.04x|
|[2752, 5120] BF16, w FP32, plain [route]|fwd|8.61|10.82
TCd|8.67|0.99x|1.25x|
|[2752, 5120] BF16, w FP32, plain [route]|bwd|22.21|25.38
QA|26.95|0.82x|0.94x|
|[2752, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|38.46|45.25
QA|43.87|0.88x|1.03x|
|[4352, 5120] BF16, w FP32, plain [route]|fwd|13.09|16.35
QA|13.06|1.00x|1.25x|
|[4352, 5120] BF16, w FP32, plain [route]|bwd|29.82|33.47
QA|35.33|0.84x|0.95x|
|[4352, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|54.75|63.71
QA|60.70|0.90x|1.05x|
|[6144, 5120] BF16, w FP32, plain [route]|fwd|18.21|21.63
QA|18.18|1.00x|1.19x|
|[6144, 5120] BF16, w FP32, plain [route]|bwd|38.94|42.37
QA|44.64|0.87x|0.95x|
|[6144, 5120] BF16, w FP32, plain [route]|2xfwd+bwd|74.94|85.70
QA|80.86|0.93x|1.06x|
|[56256, 5120] BF16, w FP32, plain [kept] [cut]|fwd|169.02|165.86
TCd|169.09|1.00x|0.98x|
|[56256, 5120] BF16, w FP32, plain [kept] [cut]|bwd|267.07|485.25
TCd|266.94|1.00x|1.82x|
|[56256, 5120] BF16, w FP32, plain [kept] [cut]|2xfwd+bwd|603.94|814.46
TCd|604.10|1.00x|1.35x|
|[56257, 5120] BF16, w FP32, plain [twin] [cut]|fwd|169.15|165.85
TCd|169.25|1.00x|0.98x|
|[56257, 5120] BF16, w FP32, plain [twin] [cut]|bwd|267.01|485.62
TCd|337.95|0.79x|1.44x|
|[56257, 5120] BF16, w FP32, plain [twin] [cut]|2xfwd+bwd|604.19|814.69
TCd|677.01|0.89x|1.20x|
|[57345, 5120] BF16, w FP32, plain [twin]|fwd|173.41|169.50
TC|173.31|1.00x|0.98x|
|[57345, 5120] BF16, w FP32, plain [twin]|bwd|274.62|384.58
TC|343.97|0.80x|1.12x|
|[57345, 5120] BF16, w FP32, plain [twin]|2xfwd+bwd|618.75|720.74
TC|690.53|0.90x|1.04x|
|[98304, 5120] BF16, w FP32, plain [twin]|fwd|295.71|290.08
QA|295.71|1.00x|0.98x|
|[98304, 5120] BF16, w FP32, plain [twin]|bwd|456.61|602.45
TC|579.53|0.79x|1.04x|
|[98304, 5120] BF16, w FP32, plain [twin]|2xfwd+bwd|1044.99|1226.13
TC|1172.63|0.89x|1.05x|
|[131072, 5120] BF16, w FP32, plain [twin]|fwd|394.02|384.32
QA|393.57|1.00x|0.98x|
|[131072, 5120] BF16, w FP32, plain [twin]|bwd|614.94|601.22 QA!
(2.352)|767.33|0.80x|0.78x!|
|[131072, 5120] BF16, w FP32, plain [twin]|2xfwd+bwd|1407.79|1381.87 QA!
(2.352)|1568.92|0.90x|0.88x!|
|[65537, 6144] BF16, w FP32, plain [carried]|fwd|233.63|230.08
TC|233.41|1.00x|0.99x|
|[65537, 6144] BF16, w FP32, plain [carried]|bwd|368.77|425.04
TC|397.18|0.93x|1.07x|
|[65537, 6144] BF16, w FP32, plain [carried]|2xfwd+bwd|833.02|886.33
TC|868.03|0.96x|1.02x|
|[98304, 6144] BF16, w FP32, plain [carried]|fwd|347.23|345.57
QA|347.17|1.00x|1.00x|
|[98304, 6144] BF16, w FP32, plain [carried]|bwd|550.62|532.54 QA!
(1.771)|589.85|0.93x|0.90x!|
|[98304, 6144] BF16, w FP32, plain [carried]|2xfwd+bwd|1251.48|1229.73
QA! (1.771)|1318.96|0.95x|0.93x!|
|[131072, 6144] BF16, w FP32, plain [carried]|fwd|461.73|458.56
QA|463.42|1.00x|0.99x|
|[131072, 6144] BF16, w FP32, plain [carried]|bwd|718.24|706.66 QA!
(2.585)|788.05|0.91x|0.90x!|
|[131072, 6144] BF16, w FP32, plain [carried]|2xfwd+bwd|1664.94|1653.58
QA! (2.585)|1790.40|0.93x|0.92x!|
|[65537, 6144] BF16, w FP32, res [twin]|fwd|466.57|467.55
Liger|466.62|1.00x|1.00x|
|[65537, 6144] BF16, w FP32, res [twin]|bwd|570.02|648.54
TC|513.17|1.11x|1.26x|
|[65537, 6144] BF16, w FP32, res [twin]|2xfwd+bwd|1507.90|1613.57
TC|1433.57|1.05x|1.13x|
|[98304, 6144] BF16, w FP32, res [twin]|fwd|698.21|702.02
Liger|698.21|1.00x|1.01x|
|[98304, 6144] BF16, w FP32, res [twin]|bwd|851.31|974.59
TC|766.37|1.11x|1.27x|
|[98304, 6144] BF16, w FP32, res [twin]|2xfwd+bwd|2264.11|2461.61
TC|2148.09|1.05x|1.15x|
|[131072, 6144] BF16, w FP32, res [twin]|fwd|929.28|936.35
Liger|929.50|1.00x|1.01x|
|[131072, 6144] BF16, w FP32, res [twin]|bwd|1124.24|985.60 QA+add!
(3.690)|1009.71|1.11x|0.98x!|
|[131072, 6144] BF16, w FP32, res [twin]|2xfwd+bwd|3019.15|3256.93 TC!
(1.472)|2887.38|1.05x|1.13x!|
**R200 (sm_107a, 212 SMs): new non-deterministic Cake vs deterministic
Cake vs strongest non-Cake SOTA. Every route-change row**
- GPU: R200 sm_107a, 212 SMs, compute capability 10.7; driver 620.62.
- Software: CUDA toolkit 13.4 (Build
cuda_13.4.r13.4/compiler.38855100_0); torch 2.14.1+cu130 (CUDA runtime
13.0).
- Sessions: 4 sessions run in parallel on the node's 4 GPUs, timed as
described under Measurement.
- Builds: all rows were timed on an earlier build of this branch. On
these routes, that build runs the same programs as this PR's head in
both modes.
|workload [T, H] x, w, form|op|det us|SOTA us (name)|new
us|det/new|SOTA/new|
|---|---|---|---|---|---|---|
|[1, 512] BF16, w BF16, plain|fwd|1.63|1.44 TC|1.66|0.98x ~|0.87x ~|
|[1, 512] BF16, w BF16, plain|bwd|4.83|1.54 TCd|5.92|0.82x|0.26x ~|
|[1, 512] BF16, w BF16, plain|2xfwd+bwd|7.71|4.38 TC|9.18|0.84x|0.48x ~|
|[4, 512] BF16, w BF16, plain|fwd|1.31|1.31 TC|1.34|0.98x ~|0.98x ~|
|[4, 512] BF16, w BF16, plain|bwd|4.38|3.01 TC|5.66|0.77x ~|0.53x ~|
|[4, 512] BF16, w BF16, plain|2xfwd+bwd|6.66|5.07 TC|8.42|0.79x ~|0.60x
~|
|[16, 512] BF16, w BF16, plain|fwd|1.63|1.47 Liger|1.63|1.00x ~|0.90x ~|
|[16, 512] BF16, w BF16, plain|bwd|4.51|3.33 TC|5.73|0.79x ~|0.58x ~|
|[16, 512] BF16, w BF16, plain|2xfwd+bwd|7.30|6.02 TCd|8.64|0.84x
~|0.70x ~|
|[64, 512] BF16, w BF16, plain|fwd|1.63|1.54 Liger|1.60|1.02x ~|0.96x ~|
|[64, 512] BF16, w BF16, plain|bwd|4.70|3.65 TCd|5.76|0.82x|0.63x ~|
|[64, 512] BF16, w BF16, plain|2xfwd+bwd|7.26|6.24 TC|8.74|0.83x|0.71x
~|
|[77, 512] BF16, w BF16, plain|fwd|1.63|1.55 Liger|1.63|1.00x ~|0.95x ~|
|[77, 512] BF16, w BF16, plain|bwd|4.61|3.58 TC|5.76|0.80x|0.62x ~|
|[77, 512] BF16, w BF16, plain|2xfwd+bwd|7.33|6.27 TC|8.77|0.84x|0.72x
~|
|[49153, 512] BF16, w BF16, plain|fwd|10.37|11.04 TC|10.37|1.00x|1.06x|
|[49153, 512] BF16, w BF16, plain|bwd|19.46|26.88 QA|24.74|0.79x|1.09x|
|[49153, 512] BF16, w BF16, plain|2xfwd+bwd|39.14|49.44
TC|44.35|0.88x|1.11x|
|[98304, 512] BF16, w BF16, plain|fwd|18.34|18.34 TC|18.40|1.00x|1.00x|
|[98304, 512] BF16, w BF16, plain|bwd|31.01|39.65 TC|39.46|0.79x|1.00x|
|[98304, 512] BF16, w BF16, plain|2xfwd+bwd|66.75|75.60
TC|75.65|0.88x|1.00x|
|[131072, 512] BF16, w BF16, plain|fwd|23.81|23.74 TC|23.78|1.00x|1.00x|
|[131072, 512] BF16, w BF16, plain|bwd|38.91|48.00 QA|48.77|0.80x|0.98x|
|[131072, 512] BF16, w BF16, plain|2xfwd+bwd|85.47|97.63
TC|95.84|0.89x|1.02x|
|[65537, 1536] BF16, w BF16, plain|fwd|33.06|33.86 TC|32.96|1.00x|1.03x|
|[65537, 1536] BF16, w BF16, plain|bwd|60.61|69.54 TC|68.06|0.89x|1.02x|
|[65537, 1536] BF16, w BF16, plain|2xfwd+bwd|126.40|135.84
TC|133.73|0.95x|1.02x|
|[98304, 1536] BF16, w BF16, plain|fwd|48.53|48.45 TC|48.42|1.00x|1.00x|
|[98304, 1536] BF16, w BF16, plain|bwd|85.12|92.32 TC|94.02|0.91x|0.98x|
|[98304, 1536] BF16, w BF16, plain|2xfwd+bwd|181.92|187.39
TC|192.99|0.94x|0.97x|
|[131072, 1536] BF16, w BF16, plain|fwd|64.00|63.30
TC|63.84|1.00x|0.99x|
|[131072, 1536] BF16, w BF16, plain|bwd|109.18|114.08
TC|120.86|0.90x|0.94x|
|[131072, 1536] BF16, w BF16, plain|2xfwd+bwd|236.64|237.92
TC|252.00|0.94x|0.94x|
|[65537, 4096] BF16, w BF16, plain|fwd|86.30|80.58 TC|86.30|1.00x|0.93x|
|[65537, 4096] BF16, w BF16, plain|bwd|132.32|141.02
QA|175.97|0.75x|0.80x|
|[65537, 4096] BF16, w BF16, plain|2xfwd+bwd|303.17|304.83
QA|348.05|0.87x|0.88x|
|[98304, 4096] BF16, w BF16, plain|fwd|127.27|119.04
TC|127.30|1.00x|0.94x|
|[98304, 4096] BF16, w BF16, plain|bwd|192.18|220.16
TC|260.48|0.74x|0.85x|
|[98304, 4096] BF16, w BF16, plain|2xfwd+bwd|445.47|456.48
TC|514.77|0.87x|0.89x|
|[131072, 4096] BF16, w BF16, plain|fwd|168.38|156.32
TC|168.34|1.00x|0.93x|
|[131072, 4096] BF16, w BF16, plain|bwd|254.24|266.05 QA!
(2.185)|343.81|0.74x|-|
|[131072, 4096] BF16, w BF16, plain|2xfwd+bwd|589.28|585.79 QA!
(2.185)|679.84|0.87x|-|
|[19457, 6144] BF16, w BF16, plain|fwd|38.56|39.10 TC|38.56|1.00x|1.01x|
|[19457, 6144] BF16, w BF16, plain|bwd|66.18|68.70 QA|65.92|1.00x|1.04x|
|[19457, 6144] BF16, w BF16, plain|2xfwd+bwd|142.50|145.47
QA|142.27|1.00x|1.02x|
|[22528, 6144] BF16, w BF16, plain|fwd|44.48|44.74 QA|44.48|1.00x|1.01x|
|[22528, 6144] BF16, w BF16, plain|bwd|74.85|77.18 QA|74.75|1.00x|1.03x|
|[22528, 6144] BF16, w BF16, plain|2xfwd+bwd|163.10|165.10
QA|162.98|1.00x|1.01x|
|[24576, 6144] BF16, w BF16, plain|fwd|48.10|48.77 QA|48.16|1.00x|1.01x|
|[24576, 6144] BF16, w BF16, plain|bwd|80.58|83.28 QA|80.55|1.00x|1.03x|
|[24576, 6144] BF16, w BF16, plain|2xfwd+bwd|176.67|179.97
QA|176.51|1.00x|1.02x|
|[28416, 6144] BF16, w BF16, plain|fwd|55.39|55.49 QA|55.39|1.00x|1.00x|
|[28416, 6144] BF16, w BF16, plain|bwd|90.72|93.57 QA|90.53|1.00x|1.03x|
|[28416, 6144] BF16, w BF16, plain|2xfwd+bwd|200.66|203.14
QA|200.38|1.00x|1.01x|
|[30720, 6144] BF16, w BF16, plain|fwd|59.78|59.65 QA|59.84|1.00x|1.00x|
|[30720, 6144] BF16, w BF16, plain|bwd|96.93|99.74 QA|96.83|1.00x|1.03x|
|[30720, 6144] BF16, w BF16, plain|2xfwd+bwd|215.71|217.95
QA|215.33|1.00x|1.01x|
|[30721, 6144] BF16, w BF16, plain|fwd|59.84|58.98 TC|59.84|1.00x|0.99x|
|[30721, 6144] BF16, w BF16, plain|bwd|96.58|99.97
QA|124.34|0.78x|0.80x|
|[30721, 6144] BF16, w BF16, plain|2xfwd+bwd|216.10|218.53
QA|244.29|0.88x|0.89x|
|[65536, 6144] BF16, w BF16, plain|fwd|124.42|121.79
QA|124.51|1.00x|0.98x|
|[65536, 6144] BF16, w BF16, plain|bwd|192.86|245.12
TC|246.18|0.78x|1.00x|
|[65536, 6144] BF16, w BF16, plain|2xfwd+bwd|441.03|489.22
TC|494.88|0.89x|0.99x|
|[131072, 6144] BF16, w BF16, plain|fwd|246.94|238.08
QA|247.12|1.00x|0.96x|
|[131072, 6144] BF16, w BF16, plain|bwd|371.23|451.09
TC|474.63|0.78x|0.95x|
|[131072, 6144] BF16, w BF16, plain|2xfwd+bwd|864.29|926.82
TC|968.88|0.89x|0.96x|
|[7169, 6144] BF16, w BF16, res|fwd|29.73|32.38 TC|29.64|1.00x|1.09x|
|[7169, 6144] BF16, w BF16, res|bwd|48.13|40.91
QA+add|46.50|1.04x|0.88x|
|[7169, 6144] BF16, w BF16, res|2xfwd+bwd|106.40|121.97
QA+add|104.93|1.01x|1.16x|
|[32768, 6144] BF16, w BF16, res|fwd|122.11|137.17
TC|122.18|1.00x|1.12x|
|[32768, 6144] BF16, w BF16, res|bwd|173.11|133.15
QA+add|158.19|1.09x|0.84x|
|[32768, 6144] BF16, w BF16, res|2xfwd+bwd|415.87|426.02
TC|401.06|1.04x|1.06x|
|[131072, 6144] BF16, w BF16, res|fwd|475.94|532.37
TC|476.15|1.00x|1.12x|
|[131072, 6144] BF16, w BF16, res|bwd|648.51|550.00
TC|590.15|1.10x|0.93x|
|[131072, 6144] BF16, w BF16, res|2xfwd+bwd|1599.11|1619.06
TC|1540.71|1.04x|1.05x|
|[1297, 256] BF16, w FP32, plain|fwd|2.05|2.05 TC|2.08|0.98x|0.98x|
|[1297, 256] BF16, w FP32, plain|bwd|5.87|6.05 TC|6.43|0.91x|0.94x|
|[1297, 256] BF16, w FP32, plain|2xfwd+bwd|9.09|9.38
TC|9.82|0.93x|0.95x|
|[1312, 256] BF16, w FP32, plain|fwd|2.05|2.02 TC|2.05|1.00x ~|0.98x ~|
|[1312, 256] BF16, w FP32, plain|bwd|5.89|6.62 TC|6.40|0.92x|1.03x|
|[1312, 256] BF16, w FP32, plain|2xfwd+bwd|9.06|9.82
TC|9.79|0.92x|1.00x|
|[1408, 256] BF16, w FP32, plain|fwd|2.05|2.08 TCd|2.05|1.00x|1.02x|
|[1408, 256] BF16, w FP32, plain|bwd|5.92|5.73 TC|6.46|0.92x|0.89x|
|[1408, 256] BF16, w FP32, plain|2xfwd+bwd|9.12|9.09
TC|9.82|0.93x|0.93x|
|[1536, 256] BF16, w FP32, plain|fwd|2.18|2.21 TC|2.18|1.00x|1.01x ~|
|[1536, 256] BF16, w FP32, plain|bwd|6.45|7.42 TC|6.59|0.98x ~|1.13x|
|[1536, 256] BF16, w FP32, plain|2xfwd+bwd|10.37|11.68
TC|10.69|0.97x|1.09x ~|
|[1, 512] BF16, w FP32, plain|fwd|1.22|1.09 TCd|1.22|1.00x ~|0.89x ~|
|[1, 512] BF16, w FP32, plain|bwd|4.26|1.25 TC|5.31|0.80x ~|0.23x ~|
|[1, 512] BF16, w FP32, plain|2xfwd+bwd|6.62|3.55 TCd|7.78|0.85x
~|0.46x|
|[4, 512] BF16, w FP32, plain|fwd|1.38|1.38 QA|1.50|0.91x ~|0.91x ~|
|[4, 512] BF16, w FP32, plain|bwd|4.32|3.20 TC|5.41|0.80x|0.59x ~|
|[4, 512] BF16, w FP32, plain|2xfwd+bwd|7.23|6.02 TC|8.80|0.82x|0.68x|
|[16, 512] BF16, w FP32, plain|fwd|1.70|1.60 TCd|1.70|1.00x|0.94x|
|[16, 512] BF16, w FP32, plain|bwd|4.54|3.78 TC|5.41|0.84x|0.70x ~|
|[16, 512] BF16, w FP32, plain|2xfwd+bwd|7.78|6.85 TCd|8.77|0.89x
~|0.78x ~|
|[64, 512] BF16, w FP32, plain|fwd|1.63|1.57 TCd|1.60|1.02x ~|0.98x ~|
|[64, 512] BF16, w FP32, plain|bwd|4.51|3.58 TCd|5.50|0.82x|0.65x ~|
|[64, 512] BF16, w FP32, plain|2xfwd+bwd|7.07|6.21 TCd|8.51|0.83x|0.73x
~|
|[229, 512] BF16, w FP32, plain|fwd|1.63|1.63 Q|1.66|0.98x|0.98x|
|[229, 512] BF16, w FP32, plain|bwd|4.51|5.54 TC|5.44|0.83x|1.02x ~|
|[229, 512] BF16, w FP32, plain|2xfwd+bwd|7.33|8.13 TC|8.61|0.85x|0.94x|
|[49153, 512] BF16, w FP32, plain|fwd|10.37|11.14 TC|10.40|1.00x|1.07x|
|[49153, 512] BF16, w FP32, plain|bwd|19.52|26.69 QA|24.64|0.79x|1.08x|
|[49153, 512] BF16, w FP32, plain|2xfwd+bwd|38.88|54.75
TC|44.26|0.88x|1.24x|
|[98304, 512] BF16, w FP32, plain|fwd|18.40|18.56 TC|18.40|1.00x|1.01x|
|[98304, 512] BF16, w FP32, plain|bwd|31.10|40.45 QA|39.46|0.79x|1.03x|
|[98304, 512] BF16, w FP32, plain|2xfwd+bwd|67.23|83.46
TC|76.00|0.88x|1.10x|
|[131072, 512] BF16, w FP32, plain|fwd|23.90|24.32 TC|23.94|1.00x|1.02x|
|[131072, 512] BF16, w FP32, plain|bwd|38.88|48.13 QA|48.70|0.80x|0.99x|
|[131072, 512] BF16, w FP32, plain|2xfwd+bwd|85.28|100.33
TC|95.79|0.89x|1.05x|
|[65, 1536] BF16, w FP32, plain|fwd|1.66|1.66 Liger|1.63|1.02x ~|1.02x
~|
|[65, 1536] BF16, w FP32, plain|bwd|4.80|3.62 TC|5.18|0.93x ~|0.70x|
|[65, 1536] BF16, w FP32, plain|2xfwd+bwd|7.62|6.62 TC|8.26|0.92x|0.80x|
|[68, 1536] BF16, w FP32, plain|fwd|1.66|1.63 Liger|1.70|0.98x ~|0.96x
~|
|[68, 1536] BF16, w FP32, plain|bwd|4.86|3.62 TC|5.15|0.94x|0.70x|
|[68, 1536] BF16, w FP32, plain|2xfwd+bwd|7.65|6.72 TC|8.29|0.92x|0.81x|
|[128, 1536] BF16, w FP32, plain|fwd|1.82|1.76 Liger|1.76|1.04x|1.00x|
|[128, 1536] BF16, w FP32, plain|bwd|4.86|3.87 TC|5.38|0.90x|0.72x|
|[128, 1536] BF16, w FP32, plain|2xfwd+bwd|7.68|7.33
TC|8.51|0.90x|0.86x|
|[256, 1536] BF16, w FP32, plain|fwd|2.08|1.98 Q|2.08|1.00x|0.95x|
|[256, 1536] BF16, w FP32, plain|bwd|5.34|6.29 TC|6.14|0.87x|1.02x ~|
|[256, 1536] BF16, w FP32, plain|2xfwd+bwd|9.28|10.72
TC|10.46|0.89x|1.02x ~|
|[1025, 1536] BF16, w FP32, plain|fwd|2.34|2.37 Q|2.24|1.04x ~|1.06x ~|
|[1025, 1536] BF16, w FP32, plain|bwd|6.46|9.06 TC|9.60|0.67x|0.94x|
|[1025, 1536] BF16, w FP32, plain|2xfwd+bwd|10.24|13.41
TC|13.57|0.75x|0.99x|
|[1264, 1536] BF16, w FP32, plain|fwd|2.43|2.46 Q|2.40|1.01x ~|1.03x ~|
|[1264, 1536] BF16, w FP32, plain|bwd|6.56|9.47 TC|9.82|0.67x|0.96x ~|
|[1264, 1536] BF16, w FP32, plain|2xfwd+bwd|10.62|13.98
TC|13.98|0.76x|1.00x|
|[2048, 1536] BF16, w FP32, plain|fwd|3.01|2.91 Q|3.10|0.97x|0.94x|
|[2048, 1536] BF16, w FP32, plain|bwd|8.00|8.86 TC|11.14|0.72x|0.80x|
|[2048, 1536] BF16, w FP32, plain|2xfwd+bwd|13.25|15.01
TC|16.48|0.80x|0.91x|
|[4096, 1536] BF16, w FP32, plain|fwd|4.00|4.19 TC|4.00|1.00x|1.05x|
|[4096, 1536] BF16, w FP32, plain|bwd|10.02|20.38 TC|12.90|0.78x|1.58x|
|[4096, 1536] BF16, w FP32, plain|2xfwd+bwd|17.54|27.81
TC|20.22|0.87x|1.38x|
|[8192, 1536] BF16, w FP32, plain|fwd|5.47|6.18 TC|5.47|1.00x|1.13x|
|[8192, 1536] BF16, w FP32, plain|bwd|14.37|23.33 TC|17.54|0.82x|1.33x|
|[8192, 1536] BF16, w FP32, plain|2xfwd+bwd|24.38|34.37
TC|27.17|0.90x|1.27x|
|[16384, 1536] BF16, w FP32, plain|fwd|9.86|11.68 TC|9.82|1.00x|1.19x|
|[16384, 1536] BF16, w FP32, plain|bwd|20.58|33.31 QA|26.66|0.77x|1.25x|
|[16384, 1536] BF16, w FP32, plain|2xfwd+bwd|39.10|55.02
QA|44.93|0.87x|1.22x|
|[65537, 1536] BF16, w FP32, plain|fwd|33.02|34.21 TC|33.09|1.00x|1.03x|
|[65537, 1536] BF16, w FP32, plain|bwd|60.42|68.93 QA|68.03|0.89x|1.01x|
|[65537, 1536] BF16, w FP32, plain|2xfwd+bwd|126.13|147.01
QA|134.08|0.94x|1.10x|
|[98304, 1536] BF16, w FP32, plain|fwd|48.51|49.22 TC|48.51|1.00x|1.01x|
|[98304, 1536] BF16, w FP32, plain|bwd|84.96|92.70 QA|94.50|0.90x|0.98x|
|[98304, 1536] BF16, w FP32, plain|2xfwd+bwd|181.65|206.94
QA|193.82|0.94x|1.07x|
|[131072, 1536] BF16, w FP32, plain|fwd|63.97|63.78
TC|64.03|1.00x|1.00x|
|[131072, 1536] BF16, w FP32, plain|bwd|108.83|163.20
TC|120.99|0.90x|1.35x|
|[131072, 1536] BF16, w FP32, plain|2xfwd+bwd|236.83|290.46
TC|252.87|0.94x|1.15x|
|[127, 2048] BF16, w FP32, plain|fwd|1.82|1.79 Liger|1.89|0.97x ~|0.95x
~|
|[127, 2048] BF16, w FP32, plain|bwd|4.96|4.13 TC|5.63|0.88x|0.73x|
|[127, 2048] BF16, w FP32, plain|2xfwd+bwd|8.16|7.30
TC|8.96|0.91x|0.81x|
|[1025, 2048] BF16, w FP32, plain|fwd|3.42|2.53 QA|3.42|1.00x|0.74x|
|[1025, 2048] BF16, w FP32, plain|bwd|6.75|9.47 TC|11.74|0.57x|0.81x|
|[1025, 2048] BF16, w FP32, plain|2xfwd+bwd|12.35|14.30
TC|17.57|0.70x|0.81x|
|[1280, 2048] BF16, w FP32, plain|fwd|3.46|2.78 Q|3.42|1.01x|0.81x|
|[1280, 2048] BF16, w FP32, plain|bwd|6.91|7.62 TC|11.94|0.58x|0.64x|
|[1280, 2048] BF16, w FP32, plain|2xfwd+bwd|12.70|12.99
TC|17.95|0.71x|0.72x|
|[1984, 2048] BF16, w FP32, plain|fwd|3.87|3.33 QA|3.84|1.01x|0.87x|
|[1984, 2048] BF16, w FP32, plain|bwd|8.61|10.62 TC|13.22|0.65x|0.80x|
|[1984, 2048] BF16, w FP32, plain|2xfwd+bwd|15.81|17.22
TC|20.55|0.77x|0.84x|
|[4096, 2048] BF16, w FP32, plain|fwd|5.15|4.99 TC|5.15|1.00x|0.97x|
|[4096, 2048] BF16, w FP32, plain|bwd|11.07|20.86 TC|16.51|0.67x|1.26x|
|[4096, 2048] BF16, w FP32, plain|2xfwd+bwd|20.67|29.60
TC|26.02|0.79x|1.14x|
|[8192, 2048] BF16, w FP32, plain|fwd|8.48|8.13 TC|8.45|1.00x|0.96x|
|[8192, 2048] BF16, w FP32, plain|bwd|16.03|24.35 TC|21.25|0.75x|1.15x|
|[8192, 2048] BF16, w FP32, plain|2xfwd+bwd|33.47|38.56
TC|38.50|0.87x|1.00x|
|[16384, 2048] BF16, w FP32, plain|fwd|13.41|13.60 TC|13.38|1.00x|1.02x|
|[16384, 2048] BF16, w FP32, plain|bwd|23.81|35.10 TC|30.85|0.77x|1.14x|
|[16384, 2048] BF16, w FP32, plain|2xfwd+bwd|49.66|61.06
TC|57.06|0.87x|1.07x|
|[256, 4096] BF16, w FP32, plain|fwd|2.30|2.27 QA|2.30|1.00x|0.99x ~|
|[256, 4096] BF16, w FP32, plain|bwd|6.21|6.50 TC|6.85|0.91x|0.95x ~|
|[256, 4096] BF16, w FP32, plain|2xfwd+bwd|10.08|10.85
TC|11.10|0.91x|0.98x|
|[65537, 4096] BF16, w FP32, plain|fwd|87.46|83.94 TC|87.39|1.00x|0.96x|
|[65537, 4096] BF16, w FP32, plain|bwd|132.45|185.28
TC|180.93|0.73x|1.02x|
|[65537, 4096] BF16, w FP32, plain|2xfwd+bwd|302.79|351.68
TC|354.94|0.85x|0.99x|
|[98304, 4096] BF16, w FP32, plain|fwd|131.20|123.51
TC|131.39|1.00x|0.94x|
|[98304, 4096] BF16, w FP32, plain|bwd|192.74|221.73
TC|268.94|0.72x|0.82x|
|[98304, 4096] BF16, w FP32, plain|2xfwd+bwd|449.44|465.75
TC|530.87|0.85x|0.88x|
|[131072, 4096] BF16, w FP32, plain|fwd|175.23|162.88
TC|175.33|1.00x|0.93x|
|[131072, 4096] BF16, w FP32, plain|bwd|255.06|287.07
TC|354.53|0.72x|0.81x|
|[131072, 4096] BF16, w FP32, plain|2xfwd+bwd|598.59|608.96
TC|703.57|0.85x|0.87x|
|[4097, 5120] BF16, w FP32, plain|fwd|8.74|10.56 QA|8.74|1.00x|1.21x|
|[4097, 5120] BF16, w FP32, plain|bwd|19.10|24.00 QA|27.42|0.70x|0.88x|
|[4097, 5120] BF16, w FP32, plain|2xfwd+bwd|35.87|43.20
QA|44.10|0.81x|0.98x|
|[8192, 5120] BF16, w FP32, plain|fwd|16.38|18.24 QA|16.35|1.00x|1.12x|
|[8192, 5120] BF16, w FP32, plain|bwd|29.09|33.92 QA|44.74|0.65x|0.76x|
|[8192, 5120] BF16, w FP32, plain|2xfwd+bwd|60.51|69.30
QA|76.51|0.79x|0.91x|
|[13024, 5120] BF16, w FP32, plain|fwd|23.84|26.08 QA|23.90|1.00x|1.09x|
|[13024, 5120] BF16, w FP32, plain|bwd|40.54|45.86 QA|68.10|0.60x|0.67x|
|[13024, 5120] BF16, w FP32, plain|2xfwd+bwd|87.14|97.70
QA|115.81|0.75x|0.84x|
|[16384, 5120] BF16, w FP32, plain|fwd|29.47|31.87 QA|29.44|1.00x|1.08x|
|[16384, 5120] BF16, w FP32, plain|bwd|48.22|53.47 QA|82.75|0.58x|0.65x|
|[16384, 5120] BF16, w FP32, plain|2xfwd+bwd|105.70|115.81
QA|141.44|0.75x|0.82x|
|[19968, 5120] BF16, w FP32, plain|fwd|35.07|37.38 QA|34.88|1.01x|1.07x|
|[19968, 5120] BF16, w FP32, plain|bwd|56.90|61.79 QA|98.46|0.58x|0.63x|
|[19968, 5120] BF16, w FP32, plain|2xfwd+bwd|125.76|135.17
QA|168.26|0.75x|0.80x|
|[256, 6144] BF16, w FP32, plain|fwd|2.78|2.69 QA|2.82|0.99x|0.95x|
|[256, 6144] BF16, w FP32, plain|bwd|6.56|7.10 TC|8.10|0.81x|0.88x ~|
|[256, 6144] BF16, w FP32, plain|2xfwd+bwd|11.30|12.06
TC|12.96|0.87x|0.93x|
|[1024, 6144] BF16, w FP32, plain|fwd|4.32|4.10 QA|4.32|1.00x ~|0.95x|
|[1024, 6144] BF16, w FP32, plain|bwd|10.62|12.32 TC|13.54|0.78x|0.91x|
|[1024, 6144] BF16, w FP32, plain|2xfwd+bwd|18.34|20.45
TC|21.54|0.85x|0.95x|
|[2304, 6144] BF16, w FP32, plain|fwd|6.43|7.49 QA|6.43|1.00x ~|1.16x ~|
|[2304, 6144] BF16, w FP32, plain|bwd|16.16|21.34 QA|22.30|0.72x|0.96x|
|[2304, 6144] BF16, w FP32, plain|2xfwd+bwd|28.16|34.94
QA|35.04|0.80x|1.00x|
|[2305, 6144] BF16, w FP32, plain|fwd|6.50|7.58 QA|6.50|1.00x|1.17x|
|[2305, 6144] BF16, w FP32, plain|bwd|16.03|21.18 QA|22.75|0.70x|0.93x|
|[2305, 6144] BF16, w FP32, plain|2xfwd+bwd|27.94|34.53
QA|34.85|0.80x|0.99x|
|[4096, 6144] BF16, w FP32, plain|fwd|10.08|11.10 QA|10.14|0.99x|1.09x|
|[4096, 6144] BF16, w FP32, plain|bwd|22.91|25.76 QA|29.28|0.78x|0.88x|
|[4096, 6144] BF16, w FP32, plain|2xfwd+bwd|42.02|46.14
QA|48.96|0.86x|0.94x|
|[8192, 6144] BF16, w FP32, plain|fwd|17.82|18.94 QA|17.82|1.00x|1.06x|
|[8192, 6144] BF16, w FP32, plain|bwd|33.98|37.41 QA|46.75|0.73x|0.80x|
|[8192, 6144] BF16, w FP32, plain|2xfwd+bwd|69.19|74.30
QA|82.21|0.84x|0.90x|
|[16384, 6144] BF16, w FP32, plain|fwd|33.33|34.85 QA|33.25|1.00x|1.05x|
|[16384, 6144] BF16, w FP32, plain|bwd|56.77|60.42 QA|86.27|0.66x|0.70x|
|[16384, 6144] BF16, w FP32, plain|2xfwd+bwd|122.69|128.71
QA|152.51|0.80x|0.84x|
|[19456, 6144] BF16, w FP32, plain|fwd|38.59|40.61 QA|38.56|1.00x|1.05x|
|[19456, 6144] BF16, w FP32, plain|bwd|65.82|69.18 QA|99.87|0.66x|0.69x|
|[19456, 6144] BF16, w FP32, plain|2xfwd+bwd|142.29|149.23
QA|176.77|0.80x|0.84x|
|[19457, 6144] BF16, w FP32, plain|fwd|38.66|40.67 QA|38.58|1.00x|1.05x|
|[19457, 6144] BF16, w FP32, plain|bwd|66.43|69.14 QA|66.56|1.00x|1.04x|
|[19457, 6144] BF16, w FP32, plain|2xfwd+bwd|142.85|149.44
QA|142.78|1.00x|1.05x|
|[22528, 6144] BF16, w FP32, plain|fwd|44.61|46.58 QA|44.67|1.00x|1.04x|
|[22528, 6144] BF16, w FP32, plain|bwd|75.30|77.86 QA|75.46|1.00x|1.03x|
|[22528, 6144] BF16, w FP32, plain|2xfwd+bwd|164.13|170.34
QA|164.51|1.00x|1.04x|
|[24576, 6144] BF16, w FP32, plain|fwd|48.27|50.46 QA|48.22|1.00x|1.05x|
|[24576, 6144] BF16, w FP32, plain|bwd|81.22|83.62 QA|81.19|1.00x|1.03x|
|[24576, 6144] BF16, w FP32, plain|2xfwd+bwd|177.12|183.49
QA|177.15|1.00x|1.04x|
|[28416, 6144] BF16, w FP32, plain|fwd|55.60|57.66 QA|55.49|1.00x|1.04x|
|[28416, 6144] BF16, w FP32, plain|bwd|91.17|93.82 QA|90.78|1.00x|1.03x|
|[28416, 6144] BF16, w FP32, plain|2xfwd+bwd|201.41|207.91
QA|201.06|1.00x|1.03x|
|[28672, 6144] BF16, w FP32, plain|fwd|55.84|58.21 QA|55.78|1.00x|1.04x|
|[28672, 6144] BF16, w FP32, plain|bwd|91.54|94.43 QA|91.39|1.00x|1.03x|
|[28672, 6144] BF16, w FP32, plain|2xfwd+bwd|202.62|209.50
QA|202.24|1.00x|1.04x|
|[28673, 6144] BF16, w FP32, plain|fwd|55.87|58.06 QA|55.81|1.00x|1.04x|
|[28673, 6144] BF16, w FP32, plain|bwd|91.17|94.27
QA|117.47|0.78x|0.80x|
|[28673, 6144] BF16, w FP32, plain|2xfwd+bwd|202.21|209.41
QA|229.19|0.88x|0.91x|
|[65536, 6144] BF16, w FP32, plain|fwd|124.48|125.38
TC|124.53|1.00x|1.01x|
|[65536, 6144] BF16, w FP32, plain|bwd|193.79|252.16
TC|246.69|0.79x|1.02x|
|[65536, 6144] BF16, w FP32, plain|2xfwd+bwd|441.98|500.22
TC|495.52|0.89x|1.01x|
|[131072, 6144] BF16, w FP32, plain|fwd|247.09|243.68
TC|247.30|1.00x|0.99x|
|[131072, 6144] BF16, w FP32, plain|bwd|372.23|487.87
TCd|475.39|0.78x|1.03x|
|[131072, 6144] BF16, w FP32, plain|2xfwd+bwd|865.83|1113.54
TCd|969.32|0.89x|1.15x|
|[96, 6144] BF16, w FP32, res|fwd|3.10|2.72 Liger|3.10|1.00x|0.88x|
|[96, 6144] BF16, w FP32, res|bwd|5.92|5.09 TC|5.82|1.02x|0.87x|
|[96, 6144] BF16, w FP32, res|2xfwd+bwd|11.20|10.53 TC|11.49|0.97x|0.92x
~|
|[7169, 6144] BF16, w FP32, res|fwd|30.18|33.66 TC|30.14|1.00x|1.12x|
|[7169, 6144] BF16, w FP32, res|bwd|47.97|40.87
QA+add|46.08|1.04x|0.89x|
|[7169, 6144] BF16, w FP32, res|2xfwd+bwd|107.36|121.09
QA+add|105.44|1.02x|1.15x|
|[32768, 6144] BF16, w FP32, res|fwd|123.90|134.82
TC|123.84|1.00x|1.09x|
|[32768, 6144] BF16, w FP32, res|bwd|172.86|133.54
QA+add|160.38|1.08x|0.83x|
|[32768, 6144] BF16, w FP32, res|2xfwd+bwd|419.30|451.90
QA+add|406.69|1.03x|1.11x|
|[131072, 6144] BF16, w FP32, res|fwd|478.98|537.51
TC|478.94|1.00x|1.12x|
|[131072, 6144] BF16, w FP32, res|bwd|647.68|785.59
TC|592.96|1.09x|1.32x|
|[131072, 6144] BF16, w FP32, res|2xfwd+bwd|1604.36|1860.10
TC|1548.95|1.04x|1.20x|
Cake ships both modes of `cake_rmsnorm_train`, with all the forms above,
on sm_100a, sm_103a and sm_107a. This PR does not change the forward,
`deterministic=True` or the sm_103a routes.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->
## Summary by CodeRabbit
* **New Features**
* Expanded RMSNorm training kernel coverage across additional GPU
architectures and token-count ranges.
* **Improvements**
* Updated backward-kernel selection and accumulation to improve
numerical accuracy across supported routes.
* Added support for additional weight and partial-gradient precision
combinations.
* **Bug Fixes**
* Updated accuracy checks and route boundaries to match the revised
kernel behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Co-authored-by: Yingyi Huang <averyh@nvidia.com>
…of the sm_103a H4096 / H5120 backward (#8849)
Related issue: #8066
`flashinfer.cake_rmsnorm_train` is RMSNorm for training. The forward
returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and
dw (FP32). The two modes, `deterministic=True` and
`deterministic=False`, share the forward and differ only in the
backward. Both modes support:
- [T, H] rows with BF16 or FP32 weight
- a fused residual add (H6144)
- per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a
shared w[256]
This PR follows #8577, #8626, #8655, #8704, #8810 and #8827. On sm_103a
(GB300), it changes only the `deterministic=False` backward of plain [T,
H] rows with H4096 and H5120 at the largest token counts. At those token
counts, FlashInfer main runs the single kernel from #8655. That kernel
exceeds the dw budget at every sampled grid T above a cut set for each
route. Above the cut, this PR runs the TwoSum single kernels that these
routes ran before #8655 (as of #8626). At and below the cut, FlashInfer
main's kernel is unchanged.
This is an accuracy fix, so the backward is slower on these rows. On
every bwd and 2xfwd+bwd row, new is slower than det, than SOTA and than
FlashInfer main's program. FlashInfer main's program was timed in
separate runs. The bwd ranges are:
- Against FlashInfer main's program, bwd is 0.9904x to 0.9971x on H4096
and 0.8426x to 0.8576x on H5120.
- Against det, bwd is 0.9767x to 0.9796x on H4096 and 0.8526x to 0.8741x
on H5120.
- Against SOTA, bwd is 0.9391x to 0.9718x on H4096 and 0.8079x to
0.8301x on H5120. No non-Cake candidate meets the accuracy budget on any
bwd or 2xfwd+bwd row. On those rows, SOTA is therefore the fastest
out-of-budget candidate (marked !).
**Changed windows**
The dw budget is per element atol 1e-4 + rtol 1e-4 vs FP64. In the
table, (a, b] means a < T <= b. A breach value is the worst dw error
divided by the budget.
|route (sm_103a, `deterministic=False` backward)|FlashInfer main's
kernel, unchanged|TwoSum kernel (this PR)|FlashInfer main's kernel:
first sampled breach above the cut|its worst dw error / budget at T
1048576|
|---|---|---|---|---|
|H4096, BF16 weight|(50176, 262144]|(262144, inf)|1.050 at T
393216|1.576|
|H4096, FP32 weight|(9216, 262144]|(262144, inf)|1.038 at T
393216|1.474|
|H5120, BF16 weight|(61952, 393216]|(393216, inf)|1.100 at T
524288|1.486|
|H5120, FP32 weight|(61952, 196608]|(196608, inf)|1.004 at T
262144|1.809|
- Budget breach: above each cut, FlashInfer main's kernel exceeds the dw
budget at every sampled grid T. Each grid T was sampled with 120 inputs
from each of three input families. H5120 with FP32 weight at T 262144
and T 393216 was also sampled with 240 inputs per family. Where a T has
two runs, the worse run counts.
- Above each cut, the `deterministic=True` backward's programs also
exceed the dw budget at every sampled grid T. Of the measured
candidates, only the TwoSum kernel stays within the budget there.
- Cut rule: each cut is the largest grid T (131072, 196608, 262144,
393216, 524288, 786432, 1048576) at or below which FlashInfer main's
kernel stays within the dw budget at every sampled point. Two of those
points are above 0.95 of the budget: H5120 with BF16 weight at T 262144
and T 393216 (worst 0.993). Both stayed within the budget on a rerun
with 240 inputs per family.
- Accuracy: every changed window meets the dw budget (per element atol
1e-4 + rtol 1e-4 vs FP64). In each window, the TwoSum kernel's worst dw
error is 0.0006 of the budget at every sampled T, from the first T above
the cut up to T 1048576. Each T was sampled with 120 inputs from each of
three input families.
**Summary over all rows**
Each ratio cell is the min-max of the printed ratios over 6 rows (both
weight dtypes, 3 T each). The "! rows" column counts the rows where no
non-Cake candidate meets the accuracy budget. On bwd and 2xfwd+bwd, that
is every row, so those SOTA/new ranges are against out-of-budget
candidates.
|GPU / arch|H|op|rows|det/new|SOTA/new|FI main/new|! rows|
|---|---|---|---|---|---|---|---|
|GB300
sm_103a|4096|fwd|6|0.9990x-1.0001x|0.9878x-1.0029x|0.9999x-1.0002x|0|
|GB300
sm_103a|4096|bwd|6|0.9767x-0.9796x|0.9391x-0.9718x!|0.9904x-0.9971x|6|
|GB300
sm_103a|4096|2xfwd+bwd|6|0.9890x-0.9909x|0.9683x-0.9966x!|0.9949x-0.9987x|6|
|GB300
sm_103a|5120|fwd|6|0.9970x-1.0009x|0.9830x-0.9904x|0.9999x-1.0000x|0|
|GB300
sm_103a|5120|bwd|6|0.8526x-0.8741x|0.8079x-0.8301x!|0.8426x-0.8576x|6|
|GB300
sm_103a|5120|2xfwd+bwd|6|0.9274x-0.9396x|0.8981x-0.9128x!|0.9245x-0.9311x|6|
**Measurement**
- Latency: the median CUPTI span of a CUDA-graph replay with L2 flushed.
Each median pools 4 sessions, one per GPU of a 4-GPU GB300 node. Each
session times 201 replays per implementation and op, split over 3
interleaved rounds. Each implementation's own setup, reduction, cast and
copy kernels are inside the span.
- det, new and the non-Cake candidates of a row group run in one
process. Their order follows a balanced design over the 4 sessions x 3
rounds, so each implementation takes each position about equally often.
- det = deterministic Cake (`deterministic=True`). This PR does not
change it, and on these rows it runs the same programs as on FlashInfer
main.
- new = this PR's `deterministic=False`.
- FI main = FlashInfer main's `deterministic=False`. In bwd, this is the
kernel this PR replaces.
- FI main was timed with the same harness, but in separate runs on other
nodes, so FI main/new compares two runs. FI main was timed on an earlier
FlashInfer main that runs the same programs on these routes.
- det/new and SOTA/new compare implementations timed in the same
process.
- SOTA = the fastest non-Cake candidate that has these outputs and meets
the accuracy budget. The budget is dx atol = rtol = 1e-2 and dw per
element atol 1e-4 + rtol 1e-4 vs FP64, taking the worst of 5 seeds.
- Candidates: QuACK 0.6.5 (quack-kernels, nvidia-cutlass-dsl 4.8.1) with
its heuristic config, and QuACK 0.6.5 rmsnorm_fwd_tuned /
rmsnorm_bwd_tuned (exhaustive per-shape autotune, cached before timing);
torch.compile (torch 2.14.1+cu130, inductor mode
max-autotune-no-cudagraphs, emulate_precision_casts=True) and
torch.compile (inductor default mode, emulate_precision_casts=True) of
the explicit FP32-statistics reference; Liger-Kernel 0.8.5 (Triton
3.8.0), casting_mode gemma (FP32 math), offset 0; Transformer Engine
2.20.2 (transformer_engine_torch 2.20.2 built against torch
2.14.1+cu130), run as the FP32 problem.
- PyTorch's fused rms_norm has no result on these rows (dtype error).
- On BF16-weight rows, library APIs that return dw in BF16 are not
matched outputs and are never chosen as SOTA.
- A/A: each process also times a second copy of det and of new.
- Over every session, row and op, the largest |1 - A/A| is 0.0060.
- A row's A/A noise is the larger |1 - A/A| of its det and new copies.
On every bwd and 2xfwd+bwd row, det/new differs from 1 by more than the
row's A/A noise in each of the 4 sessions. No fwd row does so in all 4
sessions.
- Speedup = baseline latency / new latency (> 1: new is faster). Ratios
are computed from unrounded medians, so a few differ in the last digit
from the ratio of the printed latencies.
- Builds: det, new and SOTA were timed on an earlier build of this
branch. At the timed rows, that build runs the same programs as this
PR's head in both modes.
- Rows: for each changed route, the first T above its cut and two larger
T. Each row has fwd, bwd and 2xfwd+bwd.
Key:
- form: plain = no residual add.
- op 2xfwd+bwd = checkpointed step (forward, forward recompute,
backward).
- [rerun] = row timed in a second record run of the same harness on
another GB300 node. At these shapes, Transformer Engine fails with a
CUDA illegal memory access, also when run alone. Its failing
configurations were therefore not run, and SOTA is chosen from the
remaining candidates.
- ! (x) = no non-Cake candidate meets the accuracy budget on this row
and op. The SOTA column shows the fastest matched candidate with its dw
err/tol (dw error divided by the dw budget) in parentheses, and SOTA/new
is computed against that candidate.
- SOTA names:
- QA = QuACK, per-shape tuned configuration
- Q = QuACK, heuristic configuration
- TC = torch.compile, max-autotune-no-cudagraphs mode
- TCd = torch.compile, default mode
**GB300 (sm_103a): new non-deterministic Cake vs deterministic Cake vs
fastest in-budget non-Cake candidate (!: none) vs FlashInfer main's
program**
- GPU: NVIDIA GB300, sm_103a, 152 SMs, compute capability 10.3.
- Software: NVIDIA driver 580.159.03; CUDA toolkit 13.4 (Build
cuda_13.4.r13.4/compiler.38855100_0); torch 2.14.1+cu130 (CUDA runtime
13.0).
- Timing: see Measurement.
|workload [T, H] x, w, form|op|det us|SOTA us (name)|FI main us|new
us|det/new|SOTA/new|FI main/new|
|---|---|---|---|---|---|---|---|---|
|[262145, 4096] BF16, w BF16, plain|fwd|598.56|600.40
Q|598.72|598.66|0.9998x|1.0029x|1.0001x|
|[262145, 4096] BF16, w BF16, plain|bwd|961.51|954.69 TC!
(1.168)|977.99|982.42|0.9787x|0.9718x!|0.9955x|
|[262145, 4096] BF16, w BF16, plain|2xfwd+bwd|2155.56|2168.16 QA!
(3.789)|2172.15|2175.47|0.9909x|0.9966x!|0.9985x|
|[524288, 4096] BF16, w BF16, plain|fwd|1193.19|1178.98
TCd|1193.03|1193.11|1.0001x|0.9882x|0.9999x|
|[524288, 4096] BF16, w BF16, plain|bwd|1907.51|1850.54 TC!
(1.855)|1941.32|1948.51|0.9790x|0.9497x!|0.9963x|
|[524288, 4096] BF16, w BF16, plain|2xfwd+bwd|4290.49|4204.44 TCd!
(1.190)|4324.73|4331.69|0.9905x|0.9706x!|0.9984x|
|[1048576, 4096] BF16, w BF16, plain|fwd|2382.04|2353.13
TCd|2382.34|2382.19|0.9999x|0.9878x|1.0001x|
|[1048576, 4096] BF16, w BF16, plain|bwd|3793.97|3671.83 TCd!
(2.194)|3873.27|3884.52|0.9767x|0.9452x!|0.9971x|
|[1048576, 4096] BF16, w BF16, plain|2xfwd+bwd|8556.83|8372.39 TCd!
(2.194)|8634.82|8646.11|0.9897x|0.9683x!|0.9987x|
|[262145, 4096] BF16, w FP32, plain|fwd|597.73|598.21
QA|598.27|598.31|0.9990x|0.9998x|0.9999x|
|[262145, 4096] BF16, w FP32, plain|bwd|964.77|956.07 TCd!
(1.508)|977.46|986.98|0.9775x|0.9687x!|0.9904x|
|[262145, 4096] BF16, w FP32, plain|2xfwd+bwd|2158.22|2151.72 TCd!
(1.508)|2171.12|2182.28|0.9890x|0.9860x!|0.9949x|
|[524288, 4096] BF16, w FP32, plain|fwd|1189.80|1185.06
TC|1190.66|1190.73|0.9992x|0.9952x|0.9999x|
|[524288, 4096] BF16, w FP32, plain|bwd|1913.00|1856.63 TCd!
(1.190)|1941.46|1952.94|0.9796x|0.9507x!|0.9941x|
|[524288, 4096] BF16, w FP32, plain|2xfwd+bwd|4290.29|4220.85 TC!
(1.190)|4319.46|4331.25|0.9905x|0.9745x!|0.9973x|
|[1048576, 4096] BF16, w FP32, plain [rerun]|fwd|2374.86|2364.54
TC|2374.89|2374.53|1.0001x|0.9958x|1.0002x|
|[1048576, 4096] BF16, w FP32, plain [rerun]|bwd|3807.10|3656.42 TC!
(2.423)|3864.42|3893.68|0.9778x|0.9391x!|0.9925x|
|[1048576, 4096] BF16, w FP32, plain [rerun]|2xfwd+bwd|8553.94|8379.61
TC! (2.423)|8610.93|8639.44|0.9901x|0.9699x!|0.9967x|
|[393217, 5120] BF16, w BF16, plain|fwd|1122.44|1109.48
QA|1122.24|1122.40|1.0000x|0.9885x|0.9999x|
|[393217, 5120] BF16, w BF16, plain|bwd|1817.88|1728.04 QA!
(6.592)|1787.01|2083.72|0.8724x|0.8293x!|0.8576x|
|[393217, 5120] BF16, w BF16, plain|2xfwd+bwd|4059.67|3945.59 QA!
(6.592)|4028.10|4326.29|0.9384x|0.9120x!|0.9311x|
|[524288, 5120] BF16, w BF16, plain [rerun]|fwd|1494.53|1480.13
QA|1494.49|1494.43|1.0001x|0.9904x|1.0000x|
|[524288, 5120] BF16, w BF16, plain [rerun]|bwd|2424.14|2302.02 QA!
(8.578)|2376.19|2773.28|0.8741x|0.8301x!|0.8568x|
|[524288, 5120] BF16, w BF16, plain [rerun]|2xfwd+bwd|5411.65|5257.39
QA! (8.578)|5360.80|5759.78|0.9396x|0.9128x!|0.9307x|
|[786432, 5120] BF16, w BF16, plain [rerun]|fwd|2239.84|2214.11
QA|2237.89|2237.86|1.0009x|0.9894x|1.0000x|
|[786432, 5120] BF16, w BF16, plain [rerun]|bwd|3621.84|3440.00 QA!
(13.934)|3551.34|4148.26|0.8731x|0.8293x!|0.8561x|
|[786432, 5120] BF16, w BF16, plain [rerun]|2xfwd+bwd|8100.58|7866.21
QA! (13.934)|8023.62|8620.96|0.9396x|0.9125x!|0.9307x|
|[196609, 5120] BF16, w FP32, plain|fwd|566.88|559.17
QA|568.58|568.59|0.9970x|0.9834x|1.0000x|
|[196609, 5120] BF16, w FP32, plain|bwd|913.09|868.45 QA!
(5.207)|909.52|1070.96|0.8526x|0.8109x!|0.8493x|
|[196609, 5120] BF16, w FP32, plain|2xfwd+bwd|2047.16|1983.88 QA!
(5.207)|2046.36|2207.47|0.9274x|0.8987x!|0.9270x|
|[524288, 5120] BF16, w FP32, plain [rerun]|fwd|1505.76|1481.34
QA|1506.08|1506.14|0.9997x|0.9835x|1.0000x|
|[524288, 5120] BF16, w FP32, plain [rerun]|bwd|2414.94|2284.48 QA!
(11.463)|2389.78|2826.24|0.8545x|0.8083x!|0.8456x|
|[524288, 5120] BF16, w FP32, plain [rerun]|2xfwd+bwd|5425.63|5243.68
QA! (11.463)|5400.60|5836.59|0.9296x|0.8984x!|0.9253x|
|[786432, 5120] BF16, w FP32, plain [rerun]|fwd|2258.37|2219.30
QA|2257.77|2257.70|1.0003x|0.9830x|1.0000x|
|[786432, 5120] BF16, w FP32, plain [rerun]|bwd|3610.51|3418.46 QA!
(11.435)|3565.13|4231.30|0.8533x|0.8079x!|0.8426x|
|[786432, 5120] BF16, w FP32, plain [rerun]|2xfwd+bwd|8129.10|7852.54
QA! (11.435)|8083.15|8743.18|0.9298x|0.8981x!|0.9245x|
Cake ships both modes of `cake_rmsnorm_train`, with all the forms above,
on sm_100a, sm_103a and sm_107a. This PR does not change the forward,
`deterministic=True`, the other sm_103a routes, or any sm_100a or
sm_107a route.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->
## Summary by CodeRabbit
* **New Features**
* Added training-time RMSNorm support for newer SM103 GPUs, including
hidden widths of 4096 and 5120.
* Expanded supported token-count ranges for BF16 and FP32 workloads,
with kernel selection adjusted for different workload sizes.
* Added coverage for the updated SM103 routing behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Co-authored-by: Yingyi Huang <averyh@nvidia.com>
…d windows (#8848)
Related issue: #8066
`flashinfer.cake_rmsnorm_train` is RMSNorm for training. The forward
returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and
dw (FP32). The two modes, `deterministic=True` and
`deterministic=False`, share the forward and differ only in the
backward. Both modes support:
- [T, H] rows with BF16 or FP32 weight
- a fused residual add (H6144)
- per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a
shared w[256]
This PR follows #8577, #8626, #8655, #8704, #8810 and #8827. It is the
follow-up that #8827 announced for the B200 top windows. It changes only
the `deterministic=False` backward on sm_100a (B200), in the 12 windows
below. Since #8827, FlashInfer main runs the TwoSum twin program in
these windows. With this PR, the backward there runs a single kernel
with blocked FP32 dw sums, which is the kernel kind that #8655 ships on
GB300 and R200. Above T 131072, the TwoSum twin stays.
- H1536: T 78785 to 131072 with BF16 weight, T 75777 to 131072 with FP32
weight
- H2048, H4096, H6144 and residual H6144: T 65537 to 131072, both weight
dtypes
- H5120: T 56257 to 131072, both weight dtypes
Against the TwoSum twin it replaces (main), the new backward is 1.029x
to 1.266x faster on every bwd row (det-anchored, see Measurement).
Against det, bwd ranges from 0.99x to 1.23x. Against the strongest
in-budget non-Cake SOTA, bwd ranges from 1.01x to 2.29x. Wherever
SOTA/new is below 1 on bwd or 2xfwd+bwd, the candidate shown exceeds the
dw budget (!).
On 9 rows (marked -), det is faster than new in all 4 sessions, each
time by more than the row's A/A floor in that session. On these rows
det/new is 0.9876 to 0.9982 before rounding. They are:
- bwd with FP32 weight at H4096 T 130304 and 131072, and at H5120 T
56257, 98304 and 122304
- bwd with BF16 weight at H4096 T 131072 and H5120 T 98304
- 2xfwd+bwd with FP32 weight at H4096 T 130304 and 131072
These rows come from one measurement run and are shown as measured. A
follow-up PR re-checks them.
**Summary over all rows (39 workloads x 3 ops)**
|op|det/new geomean [min, max]|det/new < 1|SOTA/new geomean [min,
max]|SOTA/new < 1|
|---|---|---|---|---|
|fwd|1.000x [0.99, 1.00]|18 of 39|1.001x [0.97, 1.05]|20 of 39|
|bwd|1.037x [0.99, 1.23]|13 of 39|1.186x [1.01, 2.29]|0 of 29|
|2xfwd+bwd|1.014x [0.99, 1.08]|15 of 39|1.091x [1.00, 1.47]|0 of 29|
The "< 1" counts use unrounded ratios, so a row printed as 1.00x can
count as < 1. SOTA/new covers only rows with an in-budget SOTA, so the
20 rows marked ! are left out. This PR does not change the forward: det,
new and main run the same forward programs, and the fwd ratios are
measured values.
**Measurement**
- Latency: the median of the CUPTI span of one CUDA-graph replay with L2
flushed (cold L2).
- 4 sessions run in parallel, one per GPU: s1 and s3 in forward order,
s2 and s4 in reverse order.
- Each session times 201 replays per implementation and op, split over 3
interleaved rounds. The median pools the 4 sessions.
- Each implementation's own setup, reduction, cast and copy kernels are
inside the span.
- The p10-p90 ranges are omitted here.
- Process: each row group runs in one process that holds det, new and
every non-Cake candidate. The implementation order rotates, so each
implementation takes each position about equally often over sessions x
rounds.
- Context (measured on GB300 and B200): on small rows, Cake gains
somewhat more than the other non-compile libraries from sharing the
process with torch.compile. SOTA/new vs the non-compile candidates is up
to about 1.10x higher (GB300) / 1.14x higher (B200) in the one-process
measurement than in a process without torch.compile; on large rows the
difference is <= 1.02x. Every row here is a large row (T >= 56257).
Ratios vs det are unaffected.
- main/new is det-anchored. main was timed in its own process with only
det and main, so main/new = (main / det of that process) / (new / det of
this PR's process). det runs byte-identical code in both processes, so
this removes the drift between processes. The raw ratio, main us / new
us, is shown in parentheses.
- Workload: eps 1e-6. For every implementation, y, dx and h_new are
BF16, and rstd and dw are FP32.
- det = deterministic Cake (`deterministic=True`), which this PR does
not change. new = this PR's non-deterministic Cake
(`deterministic=False`).
- SOTA = the fastest non-Cake candidate that has these outputs and meets
the accuracy budget, taking the worst of 5 seeds. The budget is dx atol
= rtol = 1e-2, and dw per element atol 1e-4 + rtol 1e-4 vs FP64. new and
det meet it on every row.
- Candidates: QuACK 0.6.5 (heuristic config and per-shape autotuned),
torch.compile (max-autotune-no-cudagraphs and default mode),
Liger-Kernel 0.8.5, and Transformer Engine 2.20.2 (run as the FP32
problem).
- Not candidates: PyTorch's fused rms_norm (dtype error), Apex (not
installed) and FlashInfer's inference rmsnorm (forward only).
- Accuracy: in every changed window, no measured point of new's dw
exceeds the dw budget (per element atol 1e-4 + rtol 1e-4 vs FP64), and
with FP32 weight no point exceeds the dw error of FP32 eager on the same
inputs. The worst point, 0.9988 of the budget (FP32-weight H5120, T
113536, 120 seeds), is confirmed <= 1 with 240 seeds. The seed sets are:
- 3 x 120 seeds (three input families) in all 12 windows
- fresh seeds 120..239 in the H4096, H5120, H6144 and residual H6144
windows
- 3 x 240 seeds in all 12 windows: 8 T per window, and in the
FP32-weight H5120 window the 32 T of the dense every-64 grid (112512,
114560] around the worst point
- Speedup = baseline latency / new latency (> 1: new is faster). Ratios
are computed from unrounded medians, so a few differ in the last digit
from the ratio of the printed latencies.
- Rows: for each window, the first T above its lower edge, T 98304 and T
131072. Three more rows come from an earlier screen that showed det
faster: H5120 T 128576 (BF16 weight), and H4096 T 130304 and H5120 T
122304 (FP32 weight).
Key:
- form: plain, or res = fused residual add (h_new = BF16(h + u),
combined grad). op 2xfwd+bwd = checkpointed step (forward, forward
recompute, backward).
- main = this backend on FlashInfer main before this PR, which is the
program this PR replaces. Its `deterministic=False` backward runs the
TwoSum twin in these windows (#8827), and its forward is the same
program as new's.
- + / - after det/new: in each of the 4 sessions, det/new differs from 1
by more than the row's A/A floor in that session, and all 4 sessions
agree that new (+) or det (-) is faster. Unmarked rows are within the
floor in at least one session, or the sessions disagree on which is
faster.
- Each session also times a second copy of det and of new. A row's A/A
floor in a session is the larger |1 - A/A| of its det and new copies in
that session.
- The largest A/A deviation of any row in a session was 0.0125 (s1),
0.0117 (s2), 0.0119 (s3) and 0.0078 (s4).
- ! = no non-Cake candidate meets the dw budget. The table shows the
fastest matched candidate, with its dw err/tol in parentheses, and
computes SOTA/new against it.
- SOTA names:
- TC = torch.compile (torch 2.14.1+cu130, inductor mode
max-autotune-no-cudagraphs, emulate_precision_casts=True) of the
explicit FP32-statistics reference
- TCd = the same with torch.compile inductor default mode
- Q = QuACK 0.6.5 (quack-kernels, nvidia-cutlass-dsl 4.8.1), heuristic
config
- QA = QuACK 0.6.5 rmsnorm_fwd_tuned / rmsnorm_bwd_tuned (exhaustive
per-shape autotune, cached before timing)
- Liger = Liger-Kernel 0.8.5 (Triton 3.8.0), casting_mode gemma (FP32
math), offset 0
- QA+add = torch add + QA (residual form)
**B200 (sm_100a): new non-deterministic Cake vs deterministic Cake,
strongest non-Cake SOTA and main. Every measured row**
- GPU: B200 sm_100a, 148 SMs, compute capability 10.0; driver 580.82.07;
one 8-GPU node used exclusively.
- Software: CUDA toolkit Build cuda_13.4.r13.4/compiler.38855100_0;
torch 2.14.1+cu130 (CUDA runtime 13.0).
- Builds: new was timed on an earlier build of this branch that runs the
same programs on these routes as this PR's head. main was timed on a
pre-merge build of #8827 that runs the same programs on these routes as
its merged head.
|workload [T, H] x, w, form|op|det us|SOTA us (name)|new
us|det/new|SOTA/new|main us|main/new|
|---|---|---|---|---|---|---|---|---|
|[78785, 1536] BF16, w BF16, plain|fwd|69.92|72.86
TC|69.76|1.00x|1.04x|71.26|1.004 (1.022)|
|[78785, 1536] BF16, w BF16, plain|bwd|122.56|140.26 TC|117.54|1.04x
+|1.19x|133.25|1.135 (1.134)|
|[78785, 1536] BF16, w BF16, plain|2xfwd+bwd|263.55|284.29
TC|258.59|1.02x +|1.10x|274.62|1.063 (1.062)|
|[98304, 1536] BF16, w BF16, plain|fwd|86.69|90.91
QA|86.66|1.00x|1.05x|88.61|1.001 (1.023)|
|[98304, 1536] BF16, w BF16, plain|bwd|148.86|157.98 QA|143.90|1.03x
+|1.10x|163.68|1.139 (1.137)|
|[98304, 1536] BF16, w BF16, plain|2xfwd+bwd|322.50|336.37
QA|318.08|1.01x +|1.06x|339.58|1.063 (1.068)|
|[131072, 1536] BF16, w BF16, plain|fwd|115.52|120.06
QA|115.90|1.00x|1.04x|117.25|0.999 (1.012)|
|[131072, 1536] BF16, w BF16, plain|bwd|193.58|207.23 TC|189.06|1.02x
+|1.10x|214.11|1.131 (1.133)|
|[131072, 1536] BF16, w BF16, plain|2xfwd+bwd|425.21|444.77
TC|421.63|1.01x +|1.05x|447.74|1.059 (1.062)|
|[65537, 2048] BF16, w BF16, plain|fwd|76.64|76.96
Q|76.67|1.00x|1.00x|78.78|1.001 (1.028)|
|[65537, 2048] BF16, w BF16, plain|bwd|133.28|150.94 TC|129.28|1.03x
+|1.17x|149.82|1.159 (1.159)|
|[65537, 2048] BF16, w BF16, plain|2xfwd+bwd|287.58|309.70
TC|284.22|1.01x +|1.09x|305.25|1.074 (1.074)|
|[98304, 2048] BF16, w BF16, plain|fwd|115.23|115.34
Q|115.30|1.00x|1.00x|117.25|1.000 (1.017)|
|[98304, 2048] BF16, w BF16, plain|bwd|192.13|216.38 TC|190.53|1.01x
+|1.14x|219.52|1.152 (1.152)|
|[98304, 2048] BF16, w BF16, plain|2xfwd+bwd|422.37|448.89
TC|421.22|1.00x +|1.07x|451.76|1.071 (1.073)|
|[131072, 2048] BF16, w BF16, plain|fwd|156.86|155.55
TC|157.09|1.00x|0.99x|157.31|0.999 (1.001)|
|[131072, 2048] BF16, w BF16, plain|bwd|251.04|260.10
TC|251.17|1.00x|1.04x|289.12|1.153 (1.151)|
|[131072, 2048] BF16, w BF16, plain|2xfwd+bwd|557.39|567.10
TC|558.21|1.00x|1.02x|600.16|1.074 (1.075)|
|[65537, 4096] BF16, w BF16, plain|fwd|155.74|157.28
QA|156.00|1.00x|1.01x|155.87|0.998 (0.999)|
|[65537, 4096] BF16, w BF16, plain|bwd|252.90|267.04 TC|250.45|1.01x
+|1.07x|257.92|1.029 (1.030)|
|[65537, 4096] BF16, w BF16, plain|2xfwd+bwd|557.73|580.91
TC|555.93|1.00x|1.04x|567.84|1.021 (1.021)|
|[98304, 4096] BF16, w BF16, plain|fwd|232.19|234.46
QA|232.09|1.00x|1.01x|232.50|1.000 (1.002)|
|[98304, 4096] BF16, w BF16, plain|bwd|371.04|456.99 TCd|368.83|1.01x
+|1.24x|384.22|1.038 (1.042)|
|[98304, 4096] BF16, w BF16, plain|2xfwd+bwd|827.45|925.69
TCd|825.98|1.00x|1.12x|846.98|1.022 (1.025)|
|[131072, 4096] BF16, w BF16, plain|fwd|308.64|310.98
QA|308.10|1.00x|1.01x|309.09|1.002 (1.003)|
|[131072, 4096] BF16, w BF16, plain|bwd|486.50|494.72 TC|491.42|0.99x
-|1.01x|509.33|1.037 (1.036)|
|[131072, 4096] BF16, w BF16, plain|2xfwd+bwd|1097.09|1115.01
TC|1098.45|1.00x|1.02x|1124.61|1.024 (1.024)|
|[56257, 5120] BF16, w BF16, plain|fwd|172.03|169.76
TC|172.13|1.00x|0.99x|172.19|0.999 (1.000)|
|[56257, 5120] BF16, w BF16, plain|bwd|266.21|406.72
TCd|267.26|1.00x|1.52x|332.42|1.244 (1.244)|
|[56257, 5120] BF16, w BF16, plain|2xfwd+bwd|603.81|752.26
TCd|604.61|1.00x|1.24x|674.08|1.113 (1.115)|
|[98304, 5120] BF16, w BF16, plain|fwd|296.96|292.10
QA|296.93|1.00x|0.98x|297.06|1.001 (1.000)|
|[98304, 5120] BF16, w BF16, plain|bwd|454.05|584.62 TC|454.88|1.00x
-|1.29x|569.37|1.251 (1.252)|
|[98304, 5120] BF16, w BF16, plain|2xfwd+bwd|1042.24|1182.69
TC|1043.23|1.00x|1.13x|1161.63|1.112 (1.113)|
|[128576, 5120] BF16, w BF16, plain|fwd|386.98|381.69
TC|386.75|1.00x|0.99x|386.81|1.001 (1.000)|
|[128576, 5120] BF16, w BF16, plain|bwd|599.41|596.72 QA!
(2.259)|605.18|0.99x|0.99x!|740.40|1.226 (1.223)|
|[128576, 5120] BF16, w BF16, plain|2xfwd+bwd|1368.11|1358.61 QA!
(2.259)|1374.03|1.00x|0.99x!|1511.79|1.101 (1.100)|
|[131072, 5120] BF16, w BF16, plain|fwd|394.34|389.21
QA|394.18|1.00x|0.99x|394.43|1.000 (1.001)|
|[131072, 5120] BF16, w BF16, plain|bwd|620.90|607.90 QA!
(2.650)|621.01|1.00x|0.98x!|754.78|1.216 (1.215)|
|[131072, 5120] BF16, w BF16, plain|2xfwd+bwd|1404.16|1383.23 QA!
(2.650)|1404.64|1.00x|0.98x!|1541.53|1.097 (1.097)|
|[65537, 6144] BF16, w BF16, plain|fwd|232.86|232.94
TC|232.90|1.00x|1.00x|233.34|1.000 (1.002)|
|[65537, 6144] BF16, w BF16, plain|bwd|367.10|411.39
TC|366.75|1.00x|1.12x|393.92|1.074 (1.074)|
|[65537, 6144] BF16, w BF16, plain|2xfwd+bwd|825.89|874.74
TC|825.73|1.00x|1.06x|854.91|1.034 (1.035)|
|[98304, 6144] BF16, w BF16, plain|fwd|348.13|347.15
Q|348.03|1.00x|1.00x|348.19|1.000 (1.000)|
|[98304, 6144] BF16, w BF16, plain|bwd|548.72|604.70 TC|544.42|1.01x
+|1.11x|584.25|1.076 (1.073)|
|[98304, 6144] BF16, w BF16, plain|2xfwd+bwd|1239.23|1326.78
TC|1235.39|1.00x|1.07x|1276.89|1.034 (1.034)|
|[131072, 6144] BF16, w BF16, plain|fwd|462.74|460.56
QA|462.46|1.00x|1.00x|462.40|1.001 (1.000)|
|[131072, 6144] BF16, w BF16, plain|bwd|717.60|710.22 QA!
(2.919)|717.73|1.00x|0.99x!|774.43|1.080 (1.079)|
|[131072, 6144] BF16, w BF16, plain|2xfwd+bwd|1637.76|1629.36 QA!
(2.919)|1640.48|1.00x|0.99x!|1701.41|1.038 (1.037)|
|[65537, 6144] BF16, w BF16, res|fwd|470.14|467.14
Liger|470.18|1.00x|0.99x|469.89|1.000 (0.999)|
|[65537, 6144] BF16, w BF16, res|bwd|586.24|497.70 TC|480.18|1.22x
+|1.04x|511.23|1.064 (1.065)|
|[65537, 6144] BF16, w BF16, res|2xfwd+bwd|1523.92|1440.00
TC|1410.56|1.08x +|1.02x|1448.51|1.027 (1.027)|
|[98304, 6144] BF16, w BF16, res|fwd|703.01|697.09
Liger|702.96|1.00x|0.99x|703.04|1.000 (1.000)|
|[98304, 6144] BF16, w BF16, res|bwd|874.13|758.91 TC|709.81|1.23x
+|1.07x|770.27|1.081 (1.085)|
|[98304, 6144] BF16, w BF16, res|2xfwd+bwd|2279.94|2173.25
TC|2111.89|1.08x +|1.03x|2168.99|1.026 (1.027)|
|[131072, 6144] BF16, w BF16, res|fwd|935.78|928.87
Liger|935.87|1.00x|0.99x|935.74|1.000 (1.000)|
|[131072, 6144] BF16, w BF16, res|bwd|1154.37|975.91 TCd!
(1.241)|940.66|1.23x +|1.04x!|1010.21|1.075 (1.074)|
|[131072, 6144] BF16, w BF16, res|2xfwd+bwd|3027.64|2866.11 TCd!
(1.241)|2797.35|1.08x +|1.02x!|2860.21|1.023 (1.022)|
|[75777, 1536] BF16, w FP32, plain|fwd|68.99|71.36
TC|69.02|1.00x|1.03x|70.50|1.001 (1.021)|
|[75777, 1536] BF16, w FP32, plain|bwd|118.85|123.10 QA|113.95|1.04x
+|1.08x|128.29|1.125 (1.126)|
|[75777, 1536] BF16, w FP32, plain|2xfwd+bwd|257.22|269.06
TC|252.99|1.02x +|1.06x|267.74|1.054 (1.058)|
|[98304, 1536] BF16, w FP32, plain|fwd|89.44|91.68
TC|89.38|1.00x|1.03x|90.98|1.002 (1.018)|
|[98304, 1536] BF16, w FP32, plain|bwd|148.90|153.86 QA|144.96|1.03x
+|1.06x|163.01|1.127 (1.124)|
|[98304, 1536] BF16, w FP32, plain|2xfwd+bwd|327.84|344.48
TC|324.74|1.01x +|1.06x|343.12|1.055 (1.057)|
|[131072, 1536] BF16, w FP32, plain|fwd|119.10|121.18
TC|119.36|1.00x|1.02x|120.58|1.000 (1.010)|
|[131072, 1536] BF16, w FP32, plain|bwd|193.70|207.10 TC|190.69|1.02x
+|1.09x|213.38|1.117 (1.119)|
|[131072, 1536] BF16, w FP32, plain|2xfwd+bwd|432.35|447.01
TC|430.05|1.01x|1.04x|452.83|1.052 (1.053)|
|[65537, 2048] BF16, w FP32, plain|fwd|79.23|82.08
TC|79.04|1.00x|1.04x|80.22|1.003 (1.015)|
|[65537, 2048] BF16, w FP32, plain|bwd|133.52|150.30 TC|129.54|1.03x
+|1.16x|150.37|1.160 (1.161)|
|[65537, 2048] BF16, w FP32, plain|2xfwd+bwd|288.32|310.34
TC|284.67|1.01x +|1.09x|307.77|1.079 (1.081)|
|[98304, 2048] BF16, w FP32, plain|fwd|117.39|119.17
TC|117.18|1.00x|1.02x|118.11|1.004 (1.008)|
|[98304, 2048] BF16, w FP32, plain|bwd|191.97|207.07 TC|190.03|1.01x
+|1.09x|220.35|1.160 (1.160)|
|[98304, 2048] BF16, w FP32, plain|2xfwd+bwd|422.46|441.38
TC|421.22|1.00x +|1.05x|453.02|1.075 (1.076)|
|[131072, 2048] BF16, w FP32, plain|fwd|156.25|156.67
TC|155.90|1.00x|1.00x|156.99|1.003 (1.007)|
|[131072, 2048] BF16, w FP32, plain|bwd|250.78|270.13
TC|250.53|1.00x|1.08x|290.46|1.161 (1.159)|
|[131072, 2048] BF16, w FP32, plain|2xfwd+bwd|556.66|579.49
TC|556.83|1.00x|1.04x|599.87|1.078 (1.077)|
|[65537, 4096] BF16, w FP32, plain|fwd|160.35|157.47
TC|160.26|1.00x|0.98x|160.70|1.001 (1.003)|
|[65537, 4096] BF16, w FP32, plain|bwd|252.32|269.92 TC|249.89|1.01x
+|1.08x|258.56|1.036 (1.035)|
|[65537, 4096] BF16, w FP32, plain|2xfwd+bwd|566.21|581.25
TC|564.37|1.00x +|1.03x|577.31|1.023 (1.023)|
|[98304, 4096] BF16, w FP32, plain|fwd|237.95|236.51
TCd|237.86|1.00x|0.99x|237.92|1.001 (1.000)|
|[98304, 4096] BF16, w FP32, plain|bwd|369.92|469.54
TCd|369.09|1.00x|1.27x|384.22|1.035 (1.041)|
|[98304, 4096] BF16, w FP32, plain|2xfwd+bwd|839.07|939.74
TCd|837.63|1.00x +|1.12x|856.99|1.021 (1.023)|
|[130304, 4096] BF16, w FP32, plain|fwd|313.34|306.94
TC|313.31|1.00x|0.98x|313.86|1.000 (1.002)|
|[130304, 4096] BF16, w FP32, plain|bwd|483.10|515.15 QA!
(2.333)|489.18|0.99x -|1.05x!|508.38|1.039 (1.039)|
|[130304, 4096] BF16, w FP32, plain|2xfwd+bwd|1104.14|1131.49 TC!
(1.344)|1110.05|0.99x -|1.02x!|1133.31|1.021 (1.021)|
|[131072, 4096] BF16, w FP32, plain|fwd|315.33|313.41
TC|315.26|1.00x|0.99x|315.36|1.000 (1.000)|
|[131072, 4096] BF16, w FP32, plain|bwd|485.74|494.62 TC|489.86|0.99x
-|1.01x|509.28|1.040 (1.040)|
|[131072, 4096] BF16, w FP32, plain|2xfwd+bwd|1110.97|1118.27
TC|1114.33|1.00x -|1.00x|1137.85|1.022 (1.021)|
|[56257, 5120] BF16, w FP32, plain|fwd|170.59|167.04
TC|171.55|0.99x|0.97x|171.39|0.995 (0.999)|
|[56257, 5120] BF16, w FP32, plain|bwd|267.33|484.96 TCd|270.29|0.99x
-|1.79x|337.89|1.251 (1.250)|
|[56257, 5120] BF16, w FP32, plain|2xfwd+bwd|607.17|816.35
TCd|611.17|0.99x|1.34x|678.91|1.110 (1.111)|
|[98304, 5120] BF16, w FP32, plain|fwd|298.94|291.90
QA|299.06|1.00x|0.98x|299.58|1.000 (1.002)|
|[98304, 5120] BF16, w FP32, plain|bwd|457.28|591.65 TC|458.43|1.00x
-|1.29x|579.29|1.263 (1.264)|
|[98304, 5120] BF16, w FP32, plain|2xfwd+bwd|1050.96|1195.52
TC|1052.67|1.00x|1.14x|1177.35|1.118 (1.118)|
|[122304, 5120] BF16, w FP32, plain|fwd|369.04|357.38
TC|369.74|1.00x|0.97x|371.33|0.999 (1.004)|
|[122304, 5120] BF16, w FP32, plain|bwd|564.77|565.49 QA!
(3.632)|566.93|1.00x -|1.00x!|717.12|1.266 (1.265)|
|[122304, 5120] BF16, w FP32, plain|2xfwd+bwd|1299.26|1284.44 QA!
(3.632)|1303.56|1.00x|0.99x!|1458.62|1.117 (1.119)|
|[131072, 5120] BF16, w FP32, plain|fwd|397.41|387.30
QA|396.83|1.00x|0.98x|397.21|1.001 (1.001)|
|[131072, 5120] BF16, w FP32, plain|bwd|619.17|604.53 QA!
(2.352)|618.74|1.00x|0.98x!|767.20|1.232 (1.240)|
|[131072, 5120] BF16, w FP32, plain|2xfwd+bwd|1409.81|1376.67 QA!
(2.352)|1408.75|1.00x|0.98x!|1560.91|1.105 (1.108)|
|[65537, 6144] BF16, w FP32, plain|fwd|234.85|231.23
TC|234.98|1.00x|0.98x|234.91|0.999 (1.000)|
|[65537, 6144] BF16, w FP32, plain|bwd|368.61|425.98
TC|367.94|1.00x|1.16x|397.09|1.080 (1.079)|
|[65537, 6144] BF16, w FP32, plain|2xfwd+bwd|833.22|887.36
TC|833.50|1.00x|1.06x|863.84|1.036 (1.036)|
|[98304, 6144] BF16, w FP32, plain|fwd|349.79|347.55
QA|349.63|1.00x|0.99x|349.57|1.001 (1.000)|
|[98304, 6144] BF16, w FP32, plain|bwd|551.39|534.72 QA!
(1.771)|547.17|1.01x +|0.98x!|588.99|1.077 (1.076)|
|[98304, 6144] BF16, w FP32, plain|2xfwd+bwd|1248.38|1228.22 QA!
(1.771)|1245.65|1.00x|0.99x!|1288.50|1.035 (1.034)|
|[131072, 6144] BF16, w FP32, plain|fwd|463.97|464.64
QA|464.26|1.00x|1.00x|464.54|0.999 (1.001)|
|[131072, 6144] BF16, w FP32, plain|bwd|717.82|707.14 QA!
(2.585)|720.10|1.00x|0.98x!|780.25|1.083 (1.084)|
|[131072, 6144] BF16, w FP32, plain|2xfwd+bwd|1647.59|1637.78 QA!
(2.585)|1656.51|0.99x|0.99x!|1726.19|1.040 (1.042)|
|[65537, 6144] BF16, w FP32, res|fwd|468.48|470.72
Liger|468.77|1.00x|1.00x|468.32|0.999 (0.999)|
|[65537, 6144] BF16, w FP32, res|bwd|575.62|649.28 TC|478.05|1.20x
+|1.36x|513.58|1.074 (1.074)|
|[65537, 6144] BF16, w FP32, res|2xfwd+bwd|1511.01|1596.87
TC|1411.36|1.07x +|1.13x|1447.28|1.026 (1.025)|
|[98304, 6144] BF16, w FP32, res|fwd|700.43|703.78
Liger|700.74|1.00x|1.00x|701.28|1.000 (1.001)|
|[98304, 6144] BF16, w FP32, res|bwd|858.85|1625.99 TCd|709.06|1.21x
+|2.29x|768.85|1.083 (1.084)|
|[98304, 6144] BF16, w FP32, res|2xfwd+bwd|2260.73|3089.63
TCd|2108.70|1.07x +|1.47x|2164.19|1.026 (1.026)|
|[131072, 6144] BF16, w FP32, res|fwd|932.35|935.31
Liger|932.13|1.00x|1.00x|932.41|1.001 (1.000)|
|[131072, 6144] BF16, w FP32, res|bwd|1132.91|982.43 QA+add!
(3.690)|931.18|1.22x +|1.06x!|1013.98|1.088 (1.089)|
|[131072, 6144] BF16, w FP32, res|2xfwd+bwd|3003.07|3171.44 TC!
(1.618)|2792.39|1.08x +|1.14x!|2861.18|1.026 (1.025)|
Cake ships both modes of `cake_rmsnorm_train`, with all the forms above,
on sm_100a, sm_103a and sm_107a. This PR does not change the forward,
`deterministic=True`, the sm_103a or sm_107a routes, or any B200 window
other than the 12 above.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->
## Summary by CodeRabbit
* **New Features**
* Added RMSNorm training support for additional GPU configurations,
including optimized backward processing across more input sizes.
* Added checks for tensor layouts, dimensions, and launch settings, with
clearer errors for invalid inputs.
* **Performance**
* Improved scheduling for selected workloads, including more targeted
processing for larger inputs.
* Added intermediate processing paths for certain input sizes to better
match workloads to available kernels.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Co-authored-by: Yingyi Huang <averyh@nvidia.com>
…rd windows (#8899)
Related issue: #8066
`flashinfer.cake_rmsnorm_train` is RMSNorm for training. The forward
returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and
dw (FP32). The two modes, `deterministic=True` and
`deterministic=False`, share the forward and differ only in the
backward. Both modes support:
- [T, H] rows with BF16 or FP32 weight
- a fused residual add (H6144)
- per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a
shared w[256]
This PR follows #8577, #8626, #8655, #8704, #8810, #8827, #8848 and
#8849. It changes only the `deterministic=False` backward on sm_103a
(GB300). The change covers H1536, H2048 and H6144, with BF16 or FP32
weight, plain rows, in the windows listed below. FlashInfer main runs a
TwoSum single in all of these windows, the same program that stays above
T 131072. With this PR, the backward in each window runs one of two
program kinds:
- kb: a single kernel with blocked FP32 dw sums, the kernel kind that
#8655 ships on GB300
- det: the same programs as `deterministic=True`
The windows are:
- H1536, BF16 weight: kb on T 46593 to 53184 and on T 76353 to 131072;
det on T 53185 to 76352
- H1536, FP32 weight: kb on T 37889 to 40448 and on T 76353 to 131072;
det on T 40449 to 76352
- H2048, both weight dtypes: det on T 61953 to 68160 (the programs that
#8810 routes on T 58497 to 61952); kb on T 68161 to 131072
- H6144, BF16 weight: det on T 49153 to 57344; kb on T 57345 to 131072
- H6144, FP32 weight: det on T 49153 to 83392; kb on T 83393 to 131072
In every ratio below, a value under 1 means the baseline is faster than
new. The ranges cover all rows of the table, except that the SOTA ranges
leave out the rows marked !.
- **Against main.** main is the TwoSum single this PR replaces. new is
faster on every bwd and 2xfwd+bwd row where main was timed (20 shapes,
run A). bwd is 1.0279 to 1.2870 and 2xfwd+bwd is 1.0128 to 1.1468. These
ratios are det-anchored (see Measurement).
- **Against det on the 20 det shapes.** new runs det's programs there,
and the table repeats det's time, so det/new is 1.0000.
- **Against det on the 27 kb shapes.** det is faster than new on 25 of
27 bwd rows (det/new 0.9426 to 0.9989). new is faster on the other 2:
H6144 with BF16 weight at T 98304 and 131072 (1.0060 and 1.0115). On
2xfwd+bwd, det/new is 0.9699 to 1.0058, and det is faster on 24 of 27
rows. On 23 bwd rows and 18 2xfwd+bwd rows (marked -), det is faster
than new in all 4 sessions, each time by more than the row's A/A floor
in that session. new is faster in the same sense on 3 rows (marked +):
bwd of H6144 with BF16 weight at T 98304 and 131072, and 2xfwd+bwd at T
131072.
- **Against the strongest in-budget non-Cake SOTA.** On the 43 shapes
that have an in-budget SOTA, bwd is 0.9896 to 1.5171 and 2xfwd+bwd is
0.9964 to 1.3069. SOTA is faster than new on 1 bwd row (QA, H1536 T
131072 with FP32 weight) and on 1 2xfwd+bwd row (TC, H2048 T 114688 with
BF16 weight).
- On 4 shapes, no non-Cake candidate meets the dw budget (marked !):
H6144 T 131072 with BF16 weight, and H6144 T 114688, 122880 and 131072
with FP32 weight.
- On these shapes, the fastest candidate that has these outputs is QA.
Its worst dw element is 1.543 to 2.976 times the budget. It is faster
than new: bwd 0.9210 to 0.9630 and 2xfwd+bwd 0.9666 to 0.9829.
- **Forward.** This PR does not change the forward, and the fwd values
are as measured. main/new is 0.9961 to 1.0018. det/new on kb shapes is
0.9956 to 1.0030. SOTA/new is 0.9893 to 1.1306 and is below 1 on 7 rows:
every H6144 row with BF16 weight, and H2048 T 131072 with BF16 weight.
**Measurement**
- **Latency:** the median of the CUPTI span of one CUDA-graph replay
with L2 flushed (cold L2). Each run uses the 4 GPUs of one GB300 node,
one session per GPU. Each session times 201 replays per implementation
and op, split over 3 interleaved rounds, and the median pools the 4
sessions. Each implementation's own setup, reduction, cast and copy
kernels are inside the span. In each session, det, new and every
non-Cake candidate run in one process, and the implementation order
rotates. Over the table's rows, the largest |1 - A/A| in a session is
0.0066 to 0.0341 across the 12 sessions of the three runs. In the 4
sessions of main's process, it is 0.0080 to 0.0277.
- **Context:** on small rows, Cake gains somewhat more than the other
non-compile libraries from sharing the process with torch.compile. On
GB300, SOTA/new vs the non-compile candidates is up to about 1.10x
higher in the one-process measurement than in a process without
torch.compile; on large rows the difference is <= 1.02x. Every row here
is a large row (T >= 37889). Ratios vs det are unaffected.
- **Runs:** A, B and C ran on three different GB300 nodes. Each row
comes from one run (column run): from A where A measured the shape,
otherwise from B or C. A and one other run both measured 8 shapes. On
those shapes, the other run's det/new differs from A's by at most
0.0017, and its SOTA/new by at most 0.0084.
- **Builds:** det and new were timed on an earlier build of this branch.
On kb rows, that build runs the same kb program as this PR's head. It
also ran kb in the det windows, so on det rows new us is det's time (see
Key). Its det programs in these windows are the same as on FlashInfer
main at this PR's base, and this PR does not change them. main, and the
det it is anchored to, were timed on an earlier FlashInfer main build.
In these windows, that build's main and det programs are the same as on
FlashInfer main at this PR's base.
- **det and new:** det = deterministic Cake (`deterministic=True`),
which this PR does not change. new = this PR's non-deterministic Cake
(`deterministic=False`).
- **SOTA:** the fastest non-Cake candidate in the run that has these
outputs and meets the accuracy budget. The budget is dx atol = rtol =
1e-2 and dw per element atol 1e-4 + rtol 1e-4 vs FP64, taking the worst
of 5 seeds. new meets the dw budget in every changed window. On the 8
rows marked !, no candidate meets the budget (see Key).
- Candidates: QuACK 0.6.5 (heuristic config and per-shape autotuned),
torch.compile (torch 2.14.1+cu130; max-autotune-no-cudagraphs and
default mode) of the explicit FP32-statistics reference, Liger-Kernel
0.8.5, and Transformer Engine 2.20.2 (run as the FP32 problem).
- On BF16-weight rows, library APIs that return dw in BF16 do not have
these outputs, so they are never chosen as SOTA.
- **main:** this backend on FlashInfer main before this PR. Its
`deterministic=False` backward runs the TwoSum single in these windows.
- main was timed only in run A, in its own process together with det.
- main/new = (main / det of that process) / (new / det of new's
process).
- main us is main's latency in that process. main/new is det-anchored,
so it can differ slightly from main us / new us. n/a means main was not
timed at that shape.
- **Speedup:** baseline latency / new latency (> 1: new is faster; < 1:
the baseline is faster).
Key:
- **kb / det at the end of the workload:** the program new runs at that
T. kb = the single kernel with blocked FP32 dw sums. det = the programs
of `deterministic=True`. On det rows, new us is det's measurement from
the same run, so det/new is 1.0000 by construction.
- **op 2xfwd+bwd:** checkpointed step (forward, forward recompute,
backward).
- **+ / - after det/new:** in each of the 4 sessions of that run,
det/new differs from 1 by more than the row's A/A floor in that session,
and all 4 sessions agree that new (+) or det (-) is faster. Each session
also times a second copy of det and of new. A row's A/A floor in a
session is the larger |1 - A/A| of its det and new copies. Unmarked rows
do not meet that test.
- **! after the SOTA impl and after SOTA/new:** no non-Cake candidate
meets the dw budget. The table shows the fastest candidate that has
these outputs, with its worst dw element as a multiple of the budget in
parentheses, and computes SOTA/new against it.
- **SOTA names:**
- Q = QuACK, heuristic config
- QA = QuACK, per-shape autotuned
- TC = torch.compile, inductor mode max-autotune-no-cudagraphs
- TCd = torch.compile, inductor default mode
**GB300 (sm_103a): new non-deterministic Cake vs deterministic Cake, the
strongest non-Cake SOTA and main, every measured shape**
- GPU: GB300, sm_103a, 152 SMs, compute capability 10.3.
- Software: CUDA toolkit 13.4; torch 2.14.1+cu130 (CUDA runtime 13.0).
|GPU / arch|workload [T, H] x, w, form, program|op|det Cake us|SOTA us
(impl)|main us|new ND us|det/new|SOTA/new|main/new|run|
|---|---|---|---|---|---|---|---|---|---|---|
|GB300 / sm_103a|[46593, 1536] BF16, w BF16, plain,
kb|fwd|39.9360|41.9520 QA|40.0320|40.0000|0.9984|1.0488|1.0008|A|
|GB300 / sm_103a|[46593, 1536] BF16, w BF16, plain,
kb|bwd|75.2320|83.0730 QA|95.6480|76.9920|0.9771|1.0790|1.2413|A|
|GB300 / sm_103a|[46593, 1536] BF16, w BF16, plain,
kb|2xfwd+bwd|153.6650|164.1610
QA|176.9450|157.8570|0.9734|1.0399|1.1205|A|
|GB300 / sm_103a|[53248, 1536] BF16, w BF16, plain,
det|fwd|45.2800|47.6160 QA|n/a|45.2800|1.0000|1.0516|n/a|B|
|GB300 / sm_103a|[53248, 1536] BF16, w BF16, plain,
det|bwd|84.2240|95.3930 TC|n/a|84.2240|1.0000|1.1326|n/a|B|
|GB300 / sm_103a|[53248, 1536] BF16, w BF16, plain,
det|2xfwd+bwd|172.9930|194.2090 TC|n/a|172.9930|1.0000|1.1226|n/a|B|
|GB300 / sm_103a|[65536, 1536] BF16, w BF16, plain,
det|fwd|56.4480|58.3690 QA|56.3530|56.4480|1.0000|1.0340|0.9994|A|
|GB300 / sm_103a|[65536, 1536] BF16, w BF16, plain,
det|bwd|101.2160|108.1600 QA|128.9290|101.2160|1.0000|1.0686|1.2734|A|
|GB300 / sm_103a|[65536, 1536] BF16, w BF16, plain,
det|2xfwd+bwd|212.2890|221.8730
QA|243.3765|212.2890|1.0000|1.0451|1.1468|A|
|GB300 / sm_103a|[75776, 1536] BF16, w BF16, plain,
det|fwd|64.4480|67.0080 QA|n/a|64.4480|1.0000|1.0397|n/a|C|
|GB300 / sm_103a|[75776, 1536] BF16, w BF16, plain,
det|bwd|114.8170|121.3765 QA|n/a|114.8170|1.0000|1.0571|n/a|C|
|GB300 / sm_103a|[75776, 1536] BF16, w BF16, plain,
det|2xfwd+bwd|241.5050|251.5530 QA|n/a|241.5050|1.0000|1.0416|n/a|C|
|GB300 / sm_103a|[81920, 1536] BF16, w BF16, plain,
kb|fwd|69.9840|72.2560 QA|n/a|69.9840|1.0000|1.0325|n/a|B|
|GB300 / sm_103a|[81920, 1536] BF16, w BF16, plain,
kb|bwd|123.1040|134.9130 TC|n/a|125.4730|0.9811 -|1.0752|n/a|B|
|GB300 / sm_103a|[81920, 1536] BF16, w BF16, plain,
kb|2xfwd+bwd|260.6720|277.6645 TC|n/a|266.7195|0.9773 -|1.0410|n/a|B|
|GB300 / sm_103a|[98304, 1536] BF16, w BF16, plain,
kb|fwd|83.8090|86.4320 QA|n/a|83.9685|0.9981|1.0293|n/a|B|
|GB300 / sm_103a|[98304, 1536] BF16, w BF16, plain,
kb|bwd|146.3360|163.1680 TC|n/a|149.0560|0.9818 -|1.0947|n/a|B|
|GB300 / sm_103a|[98304, 1536] BF16, w BF16, plain,
kb|2xfwd+bwd|311.5210|333.9850 TC|n/a|317.6815|0.9806 -|1.0513|n/a|B|
|GB300 / sm_103a|[131072, 1536] BF16, w BF16, plain,
kb|fwd|112.3530|114.7210 QA|112.3520|112.2890|1.0006|1.0217|1.0003|A|
|GB300 / sm_103a|[131072, 1536] BF16, w BF16, plain,
kb|bwd|190.5610|200.9930 TC|241.4575|194.0810|0.9819 -|1.0356|1.2441|A|
|GB300 / sm_103a|[131072, 1536] BF16, w BF16, plain,
kb|2xfwd+bwd|412.7705|426.9460 TC|467.4585|420.2735|0.9821
-|1.0159|1.1123|A|
|GB300 / sm_103a|[61953, 2048] BF16, w BF16, plain,
det|fwd|70.8480|70.9920 Q|70.8480|70.8480|1.0000|1.0020|1.0018|A|
|GB300 / sm_103a|[61953, 2048] BF16, w BF16, plain,
det|bwd|122.1130|137.6000 TC|138.1130|122.1130|1.0000|1.1268|1.1319|A|
|GB300 / sm_103a|[61953, 2048] BF16, w BF16, plain,
det|2xfwd+bwd|263.7780|281.9050
QA|280.9290|263.7780|1.0000|1.0687|1.0658|A|
|GB300 / sm_103a|[64000, 2048] BF16, w BF16, plain,
det|fwd|73.0240|73.4080 Q|n/a|73.0240|1.0000|1.0053|n/a|C|
|GB300 / sm_103a|[64000, 2048] BF16, w BF16, plain,
det|bwd|126.5930|141.7280 TC|n/a|126.5930|1.0000|1.1196|n/a|C|
|GB300 / sm_103a|[64000, 2048] BF16, w BF16, plain,
det|2xfwd+bwd|272.5150|291.0400 TC|n/a|272.5150|1.0000|1.0680|n/a|C|
|GB300 / sm_103a|[66560, 2048] BF16, w BF16, plain,
det|fwd|75.9040|76.1930 Q|n/a|75.9040|1.0000|1.0038|n/a|C|
|GB300 / sm_103a|[66560, 2048] BF16, w BF16, plain,
det|bwd|130.5925|142.9770 TC|n/a|130.5925|1.0000|1.0948|n/a|C|
|GB300 / sm_103a|[66560, 2048] BF16, w BF16, plain,
det|2xfwd+bwd|281.9860|297.6650 TC|n/a|281.9860|1.0000|1.0556|n/a|C|
|GB300 / sm_103a|[69632, 2048] BF16, w BF16, plain,
kb|fwd|79.5190|79.7440 Q|n/a|79.4560|1.0008|1.0036|n/a|C|
|GB300 / sm_103a|[69632, 2048] BF16, w BF16, plain,
kb|bwd|135.9040|146.6570 TC|n/a|140.7040|0.9659 -|1.0423|n/a|C|
|GB300 / sm_103a|[69632, 2048] BF16, w BF16, plain,
kb|2xfwd+bwd|294.2420|307.4890 TC|n/a|300.0340|0.9807 -|1.0248|n/a|C|
|GB300 / sm_103a|[81920, 2048] BF16, w BF16, plain,
kb|fwd|93.4730|93.5360 Q|93.6005|93.4730|1.0000|1.0007|1.0010|A|
|GB300 / sm_103a|[81920, 2048] BF16, w BF16, plain,
kb|bwd|158.2885|166.4330 TC|177.2170|163.2330|0.9697 -|1.0196|1.0840|A|
|GB300 / sm_103a|[81920, 2048] BF16, w BF16, plain,
kb|2xfwd+bwd|345.4720|354.6250 TC|366.0505|351.6020|0.9826
-|1.0086|1.0405|A|
|GB300 / sm_103a|[98304, 2048] BF16, w BF16, plain,
kb|fwd|111.8400|111.8720 Q|n/a|111.8720|0.9997|1.0000|n/a|B|
|GB300 / sm_103a|[98304, 2048] BF16, w BF16, plain,
kb|bwd|186.3040|212.0655 TC|n/a|194.6240|0.9573 -|1.0896|n/a|B|
|GB300 / sm_103a|[98304, 2048] BF16, w BF16, plain,
kb|2xfwd+bwd|409.9200|436.1130 TC|n/a|419.7600|0.9766 -|1.0390|n/a|B|
|GB300 / sm_103a|[114688, 2048] BF16, w BF16, plain,
kb|fwd|130.7520|130.9115 Q|n/a|130.8480|0.9993|1.0005|n/a|B|
|GB300 / sm_103a|[114688, 2048] BF16, w BF16, plain,
kb|bwd|214.7850|225.1200 TC|n/a|223.8720|0.9594 -|1.0056|n/a|B|
|GB300 / sm_103a|[114688, 2048] BF16, w BF16, plain,
kb|2xfwd+bwd|475.8085|487.3610 TC|n/a|489.1210|0.9728 -|0.9964|n/a|B|
|GB300 / sm_103a|[131072, 2048] BF16, w BF16, plain,
kb|fwd|152.3850|151.4250 TC|153.0570|153.0570|0.9956|0.9893|0.9994|A|
|GB300 / sm_103a|[131072, 2048] BF16, w BF16, plain,
kb|bwd|242.6260|259.4095 TC|273.6330|252.8975|0.9594 -|1.0257|1.0820|A|
|GB300 / sm_103a|[131072, 2048] BF16, w BF16, plain,
kb|2xfwd+bwd|541.3790|558.6915 TC|577.4275|556.5780|0.9727
-|1.0038|1.0373|A|
|GB300 / sm_103a|[49153, 6144] BF16, w BF16, plain,
det|fwd|171.8410|170.7850 TC|171.8410|171.8410|1.0000|0.9939|1.0000|A|
|GB300 / sm_103a|[49153, 6144] BF16, w BF16, plain,
det|bwd|277.5530|329.5060 TC|294.3060|277.5530|1.0000|1.1872|1.0615|A|
|GB300 / sm_103a|[49153, 6144] BF16, w BF16, plain,
det|2xfwd+bwd|614.5470|667.9390
TC|632.1640|614.5470|1.0000|1.0869|1.0289|A|
|GB300 / sm_103a|[57344, 6144] BF16, w BF16, plain,
det|fwd|199.6480|199.1690 QA|n/a|199.6480|1.0000|0.9976|n/a|C|
|GB300 / sm_103a|[57344, 6144] BF16, w BF16, plain,
det|bwd|322.0815|336.8980 TC|n/a|322.0815|1.0000|1.0460|n/a|C|
|GB300 / sm_103a|[57344, 6144] BF16, w BF16, plain,
det|2xfwd+bwd|713.7625|747.2995 TC|n/a|713.7625|1.0000|1.0470|n/a|C|
|GB300 / sm_103a|[65536, 6144] BF16, w BF16, plain,
kb|fwd|227.9370|227.4260 Q|227.9370|227.7770|1.0007|0.9985|1.0010|A|
|GB300 / sm_103a|[65536, 6144] BF16, w BF16, plain,
kb|bwd|366.1775|389.7300 TC|384.8975|366.5950|0.9989|1.0631|1.0501|A|
|GB300 / sm_103a|[65536, 6144] BF16, w BF16, plain,
kb|2xfwd+bwd|815.2520|859.1410
TC|834.3090|814.8520|1.0005|1.0544|1.0241|A|
|GB300 / sm_103a|[73728, 6144] BF16, w BF16, plain,
kb|fwd|255.3610|254.4010 QA|n/a|255.3450|1.0001|0.9963|n/a|C|
|GB300 / sm_103a|[73728, 6144] BF16, w BF16, plain,
kb|bwd|409.2650|569.8910 TCd|n/a|410.9620|0.9959 -|1.3867|n/a|C|
|GB300 / sm_103a|[73728, 6144] BF16, w BF16, plain,
kb|2xfwd+bwd|914.2270|1112.1000 TCd|n/a|914.8355|0.9993|1.2156|n/a|C|
|GB300 / sm_103a|[98304, 6144] BF16, w BF16, plain,
kb|fwd|339.8420|338.7540 Q|339.8100|339.7940|1.0001|0.9969|1.0001|A|
|GB300 / sm_103a|[98304, 6144] BF16, w BF16, plain,
kb|bwd|542.4825|602.4350 TC|564.9945|539.2350|1.0060 +|1.1172|1.0477|A|
|GB300 / sm_103a|[98304, 6144] BF16, w BF16, plain,
kb|2xfwd+bwd|1214.8545|1298.5990
TC|1239.5595|1212.8055|1.0017|1.0707|1.0211|A|
|GB300 / sm_103a|[131072, 6144] BF16, w BF16, plain,
kb|fwd|451.8905|449.7940 Q|451.7470|451.7785|1.0002|0.9956|1.0000|A|
|GB300 / sm_103a|[131072, 6144] BF16, w BF16, plain,
kb|bwd|717.2845|682.9315 QA! (1.693)|746.1810|709.1390|1.0115
+|0.9630!|1.0514|A|
|GB300 / sm_103a|[131072, 6144] BF16, w BF16, plain,
kb|2xfwd+bwd|1616.0885|1579.4000 QA! (1.693)|1645.1775|1606.8245|1.0058
+|0.9829!|1.0237|A|
|GB300 / sm_103a|[37889, 1536] BF16, w FP32, plain,
kb|fwd|32.9920|37.4400 QA|33.2480|33.1200|0.9961|1.1304|0.9990|A|
|GB300 / sm_103a|[37889, 1536] BF16, w FP32, plain,
kb|bwd|60.6560|69.4725 QA|80.7685|64.3530|0.9426 -|1.0796|1.2561|A|
|GB300 / sm_103a|[37889, 1536] BF16, w FP32, plain,
kb|2xfwd+bwd|124.8650|141.7610 QA|145.0570|128.7335|0.9699
-|1.1012|1.1259|A|
|GB300 / sm_103a|[40449, 1536] BF16, w FP32, plain,
det|fwd|35.0410|39.6160 QA|n/a|35.0410|1.0000|1.1306|n/a|B|
|GB300 / sm_103a|[40449, 1536] BF16, w FP32, plain,
det|bwd|65.4410|72.3520 QA|n/a|65.4410|1.0000|1.1056|n/a|B|
|GB300 / sm_103a|[40449, 1536] BF16, w FP32, plain,
det|2xfwd+bwd|134.1445|149.0560 QA|n/a|134.1445|1.0000|1.1112|n/a|B|
|GB300 / sm_103a|[46593, 1536] BF16, w FP32, plain,
det|fwd|40.3840|43.9675 TC|n/a|40.3840|1.0000|1.0887|n/a|B|
|GB300 / sm_103a|[46593, 1536] BF16, w FP32, plain,
det|bwd|73.9210|80.5120 QA|n/a|73.9210|1.0000|1.0892|n/a|B|
|GB300 / sm_103a|[46593, 1536] BF16, w FP32, plain,
det|2xfwd+bwd|153.1200|167.9360 QA|n/a|153.1200|1.0000|1.0968|n/a|B|
|GB300 / sm_103a|[53248, 1536] BF16, w FP32, plain,
det|fwd|46.1440|50.7200 TC|n/a|46.1440|1.0000|1.0992|n/a|B|
|GB300 / sm_103a|[53248, 1536] BF16, w FP32, plain,
det|bwd|83.3600|126.4640 Q|n/a|83.3600|1.0000|1.5171|n/a|B|
|GB300 / sm_103a|[53248, 1536] BF16, w FP32, plain,
det|2xfwd+bwd|173.6000|226.8805 Q|n/a|173.6000|1.0000|1.3069|n/a|B|
|GB300 / sm_103a|[65536, 1536] BF16, w FP32, plain,
det|fwd|57.4730|60.8000 TC|57.2480|57.4730|1.0000|1.0579|0.9961|A|
|GB300 / sm_103a|[65536, 1536] BF16, w FP32, plain,
det|bwd|100.4490|106.0645 QA|129.2810|100.4490|1.0000|1.0559|1.2870|A|
|GB300 / sm_103a|[65536, 1536] BF16, w FP32, plain,
det|2xfwd+bwd|213.6015|228.5450
QA|241.9855|213.6015|1.0000|1.0700|1.1325|A|
|GB300 / sm_103a|[75776, 1536] BF16, w FP32, plain,
det|fwd|65.8240|69.5680 TC|n/a|65.8240|1.0000|1.0569|n/a|C|
|GB300 / sm_103a|[75776, 1536] BF16, w FP32, plain,
det|bwd|114.1760|119.2965 QA|n/a|114.1760|1.0000|1.0448|n/a|C|
|GB300 / sm_103a|[75776, 1536] BF16, w FP32, plain,
det|2xfwd+bwd|243.1055|259.3610 QA|n/a|243.1055|1.0000|1.0669|n/a|C|
|GB300 / sm_103a|[81920, 1536] BF16, w FP32, plain,
kb|fwd|71.1680|74.2080 TC|n/a|71.1360|1.0004|1.0432|n/a|B|
|GB300 / sm_103a|[81920, 1536] BF16, w FP32, plain,
kb|bwd|122.0800|127.1040 QA|n/a|125.6320|0.9717 -|1.0117|n/a|B|
|GB300 / sm_103a|[81920, 1536] BF16, w FP32, plain,
kb|2xfwd+bwd|261.3440|277.8565 QA|n/a|264.6080|0.9877|1.0501|n/a|B|
|GB300 / sm_103a|[98304, 1536] BF16, w FP32, plain,
kb|fwd|85.4085|88.1600 TC|n/a|85.1525|1.0030|1.0353|n/a|B|
|GB300 / sm_103a|[98304, 1536] BF16, w FP32, plain,
kb|bwd|145.2160|158.1120 TC|n/a|148.7365|0.9763 -|1.0630|n/a|B|
|GB300 / sm_103a|[98304, 1536] BF16, w FP32, plain,
kb|2xfwd+bwd|314.0160|331.1045 TC|n/a|316.2880|0.9928|1.0468|n/a|B|
|GB300 / sm_103a|[114688, 1536] BF16, w FP32, plain,
kb|fwd|99.5850|102.5600 TC|n/a|99.4560|1.0013|1.0312|n/a|B|
|GB300 / sm_103a|[114688, 1536] BF16, w FP32, plain,
kb|bwd|167.2650|176.5460 TC|n/a|171.5200|0.9752 -|1.0293|n/a|B|
|GB300 / sm_103a|[114688, 1536] BF16, w FP32, plain,
kb|2xfwd+bwd|363.6810|375.3610 TC|n/a|367.9040|0.9885|1.0203|n/a|B|
|GB300 / sm_103a|[131072, 1536] BF16, w FP32, plain,
kb|fwd|114.4320|116.8330 TC|114.2730|114.2405|1.0017|1.0227|1.0006|A|
|GB300 / sm_103a|[131072, 1536] BF16, w FP32, plain,
kb|bwd|189.6970|191.7450 QA|242.6735|193.7610|0.9790 -|0.9896|1.2526|A|
|GB300 / sm_103a|[131072, 1536] BF16, w FP32, plain,
kb|2xfwd+bwd|415.6345|428.9300
QA|468.8345|419.6170|0.9905|1.0222|1.1179|A|
|GB300 / sm_103a|[61953, 2048] BF16, w FP32, plain,
det|fwd|72.4800|73.9200 QA|72.5120|72.4800|1.0000|1.0199|0.9985|A|
|GB300 / sm_103a|[61953, 2048] BF16, w FP32, plain,
det|bwd|122.0165|132.7370 QA|145.7280|122.0165|1.0000|1.0879|1.1943|A|
|GB300 / sm_103a|[61953, 2048] BF16, w FP32, plain,
det|2xfwd+bwd|264.0020|275.4900
QA|288.5460|264.0020|1.0000|1.0435|1.0926|A|
|GB300 / sm_103a|[64000, 2048] BF16, w FP32, plain,
det|fwd|74.5600|75.8415 QA|n/a|74.5600|1.0000|1.0172|n/a|C|
|GB300 / sm_103a|[64000, 2048] BF16, w FP32, plain,
det|bwd|126.4640|136.2085 QA|n/a|126.4640|1.0000|1.0771|n/a|C|
|GB300 / sm_103a|[64000, 2048] BF16, w FP32, plain,
det|2xfwd+bwd|272.6090|283.1050 QA|n/a|272.6090|1.0000|1.0385|n/a|C|
|GB300 / sm_103a|[66560, 2048] BF16, w FP32, plain,
det|fwd|77.1840|78.6405 QA|n/a|77.1840|1.0000|1.0189|n/a|C|
|GB300 / sm_103a|[66560, 2048] BF16, w FP32, plain,
det|bwd|130.3680|141.2810 QA|n/a|130.3680|1.0000|1.0837|n/a|C|
|GB300 / sm_103a|[66560, 2048] BF16, w FP32, plain,
det|2xfwd+bwd|282.7195|294.4170 QA|n/a|282.7195|1.0000|1.0414|n/a|C|
|GB300 / sm_103a|[69632, 2048] BF16, w FP32, plain,
kb|fwd|81.0250|82.3370 QA|n/a|81.0560|0.9996|1.0158|n/a|C|
|GB300 / sm_103a|[69632, 2048] BF16, w FP32, plain,
kb|bwd|136.1920|151.1370 TC|n/a|140.1920|0.9715 -|1.0781|n/a|C|
|GB300 / sm_103a|[69632, 2048] BF16, w FP32, plain,
kb|2xfwd+bwd|295.6490|314.0810 TC|n/a|299.9050|0.9858 -|1.0473|n/a|C|
|GB300 / sm_103a|[81920, 2048] BF16, w FP32, plain,
kb|fwd|94.9760|96.2240 TC|95.0730|94.9755|1.0000|1.0131|1.0000|A|
|GB300 / sm_103a|[81920, 2048] BF16, w FP32, plain,
kb|bwd|158.1130|167.1690 TC|187.4255|162.2730|0.9744 -|1.0302|1.1555|A|
|GB300 / sm_103a|[81920, 2048] BF16, w FP32, plain,
kb|2xfwd+bwd|346.1140|356.9940 TC|376.2580|350.0660|0.9887
-|1.0198|1.0743|A|
|GB300 / sm_103a|[98304, 2048] BF16, w FP32, plain,
kb|fwd|113.2800|114.5600 TC|n/a|113.3440|0.9994|1.0107|n/a|B|
|GB300 / sm_103a|[98304, 2048] BF16, w FP32, plain,
kb|bwd|186.3680|202.3365 TC|n/a|194.1445|0.9599 -|1.0422|n/a|B|
|GB300 / sm_103a|[98304, 2048] BF16, w FP32, plain,
kb|2xfwd+bwd|410.3690|428.6400 TC|n/a|417.8245|0.9822 -|1.0259|n/a|B|
|GB300 / sm_103a|[114688, 2048] BF16, w FP32, plain,
kb|fwd|131.6480|132.8955 TC|n/a|131.5200|1.0010|1.0105|n/a|B|
|GB300 / sm_103a|[114688, 2048] BF16, w FP32, plain,
kb|bwd|215.3930|224.3205 TC|n/a|224.1915|0.9608 -|1.0006|n/a|B|
|GB300 / sm_103a|[114688, 2048] BF16, w FP32, plain,
kb|2xfwd+bwd|475.8400|485.4090 TC|n/a|485.2325|0.9806 -|1.0004|n/a|B|
|GB300 / sm_103a|[131072, 2048] BF16, w FP32, plain,
kb|fwd|151.2010|151.8090 TC|151.3290|151.3930|0.9987|1.0027|0.9987|A|
|GB300 / sm_103a|[131072, 2048] BF16, w FP32, plain,
kb|bwd|243.0410|285.8900 TCd|291.0740|252.8010|0.9614 -|1.1309|1.1509|A|
|GB300 / sm_103a|[131072, 2048] BF16, w FP32, plain,
kb|2xfwd+bwd|541.8270|604.2270 TCd|593.1225|553.6670|0.9786
-|1.0913|1.0708|A|
|GB300 / sm_103a|[49153, 6144] BF16, w FP32, plain,
det|fwd|171.7770|172.0970 TC|171.7450|171.7770|1.0000|1.0019|0.9998|A|
|GB300 / sm_103a|[49153, 6144] BF16, w FP32, plain,
det|bwd|278.0810|337.7620 TC|299.3620|278.0810|1.0000|1.2146|1.0780|A|
|GB300 / sm_103a|[49153, 6144] BF16, w FP32, plain,
det|2xfwd+bwd|612.3550|679.0275
TC|637.8920|612.3550|1.0000|1.1089|1.0420|A|
|GB300 / sm_103a|[57344, 6144] BF16, w FP32, plain,
det|fwd|198.8160|199.5840 QA|n/a|198.8160|1.0000|1.0039|n/a|B|
|GB300 / sm_103a|[57344, 6144] BF16, w FP32, plain,
det|bwd|321.1045|343.6170 TC|n/a|321.1045|1.0000|1.0701|n/a|B|
|GB300 / sm_103a|[57344, 6144] BF16, w FP32, plain,
det|2xfwd+bwd|710.6895|751.2810 TC|n/a|710.6895|1.0000|1.0571|n/a|B|
|GB300 / sm_103a|[65536, 6144] BF16, w FP32, plain,
det|fwd|227.1050|228.0010 QA|227.1370|227.1050|1.0000|1.0039|0.9999|A|
|GB300 / sm_103a|[65536, 6144] BF16, w FP32, plain,
det|bwd|366.9785|400.6095 TC|391.1380|366.9785|1.0000|1.0916|1.0660|A|
|GB300 / sm_103a|[65536, 6144] BF16, w FP32, plain,
det|2xfwd+bwd|812.2605|868.2600
TC|839.9080|812.2605|1.0000|1.0689|1.0342|A|
|GB300 / sm_103a|[81920, 6144] BF16, w FP32, plain,
det|fwd|281.3755|282.4960 QA|n/a|281.3755|1.0000|1.0040|n/a|B|
|GB300 / sm_103a|[81920, 6144] BF16, w FP32, plain,
det|bwd|452.8330|505.3130 TCd|n/a|452.8330|1.0000|1.1159|n/a|B|
|GB300 / sm_103a|[81920, 6144] BF16, w FP32, plain,
det|2xfwd+bwd|1008.2255|1085.5705 TCd|n/a|1008.2255|1.0000|1.0767|n/a|B|
|GB300 / sm_103a|[98304, 6144] BF16, w FP32, plain,
kb|fwd|338.4665|339.4095 QA|337.9380|337.9860|1.0014|1.0042|1.0000|A|
|GB300 / sm_103a|[98304, 6144] BF16, w FP32, plain,
kb|bwd|542.3710|620.4510 TC|574.5630|559.4425|0.9695 -|1.1091|1.0279|A|
|GB300 / sm_103a|[98304, 6144] BF16, w FP32, plain,
kb|2xfwd+bwd|1211.4785|1318.5985 TC|1245.5755|1229.9585|0.9850
-|1.0721|1.0128|A|
|GB300 / sm_103a|[114688, 6144] BF16, w FP32, plain,
kb|fwd|391.7115|392.9920 QA|n/a|391.9370|0.9994|1.0027|n/a|B|
|GB300 / sm_103a|[114688, 6144] BF16, w FP32, plain,
kb|bwd|626.9765|598.4330 QA! (2.976)|n/a|649.7455|0.9650
-|0.9210!|n/a|B|
|GB300 / sm_103a|[114688, 6144] BF16, w FP32, plain,
kb|2xfwd+bwd|1404.7860|1380.2270 QA! (2.976)|n/a|1427.8740|0.9838
-|0.9666!|n/a|B|
|GB300 / sm_103a|[122880, 6144] BF16, w FP32, plain,
kb|fwd|419.6010|420.7040 QA|n/a|419.6480|0.9999|1.0025|n/a|B|
|GB300 / sm_103a|[122880, 6144] BF16, w FP32, plain,
kb|bwd|675.8250|639.7455 QA! (2.731)|n/a|693.7920|0.9741
-|0.9221!|n/a|B|
|GB300 / sm_103a|[122880, 6144] BF16, w FP32, plain,
kb|2xfwd+bwd|1509.6185|1477.0905 QA! (2.731)|n/a|1527.4905|0.9883
-|0.9670!|n/a|B|
|GB300 / sm_103a|[131072, 6144] BF16, w FP32, plain,
kb|fwd|449.2815|450.4665 QA|449.0900|449.1540|1.0003|1.0029|0.9999|A|
|GB300 / sm_103a|[131072, 6144] BF16, w FP32, plain,
kb|bwd|715.4435|682.6280 QA! (1.543)|759.1400|727.8110|0.9830
-|0.9379!|1.0440|A|
|GB300 / sm_103a|[131072, 6144] BF16, w FP32, plain,
kb|2xfwd+bwd|1605.1605|1579.2400 QA! (1.543)|1652.2655|1619.4800|0.9912
-|0.9752!|1.0197|A|
Cake ships both modes of `cake_rmsnorm_train`, with all the forms above,
on sm_100a, sm_103a and sm_107a. This PR does not change the forward,
`deterministic=True`, the sm_100a or sm_107a routes, or any GB300 window
other than the ones above.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
<!-- coderabbit-ai-generated:v1 -->
## Summary by CodeRabbit
* **New Features**
* Expanded RMSNorm training support for SM103a GPUs across additional
workload sizes and weight precisions.
* Added prebuilt CUDA 12.9 kernel options for supported configurations.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
Co-authored-by: Yingyi Huang <averyh@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
flashinfer-bot
added
op: misc
yzh119
left a comment
flashinfer-bot
left a comment
Related issue: #8066
flashinfer.cake_rmsnorm_trainis RMSNorm for training. The forward returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and dw (FP32). The two modes,deterministic=Trueanddeterministic=False, use the same forward and differ only in the backward. The family ships for sm_100a (B200), sm_103a (GB300) and sm_107a (R200). The tables cover these forms in both modes:w32): H5120 and H6144, plus H2048 and H4096 in the B200 small-T rowsres, H6144) with BF16 or FP32 weightThis follow-up to #8208 and #8253 changes only the
deterministic=Falsebackward. It makes two changes:resH6144 route and their FP32-weight twins) now launch one CTA per row chunk plus one per column task, capped at the slot count, as the deterministic column-task programs already do; the grid changes only at small T, where a few CTAs used to run every column task in turn.Speedup = baseline latency / new latency. A value > 1 means the new kernel is faster. Geomeans over the 78-shape list, vs deterministic / vs SOTA:
The B200 vs-SOTA geomeans leave out the two B200
!rows.Added rows, vs deterministic / vs SOTA (small T = T1, T4, T16 and T64):
m): fwd 0.999x / 0.936x, bwd 1.028x / 1.007x, 2xfwd+bwd 1.025x / 1.008xmrows re-timed after fix(cake_rmsnorm_det): fixed-order dw for the smallest token counts of the SM100/SM103 backward #8512): fwd 1.003x / 0.888x, bwd 0.958x / 1.049x, 2xfwd+bwd 0.982x / 1.221xNotable rows:
m). There, B200 small-T bwd reads 0.94x to 1.36x vs det (geomean 1.028x; 19 of 36 rows < 1, 29 of 36 within 5 % of 1.00x); the lowest is H2048 T4 w32 bwd at 3.97 vs 3.71 us (new vs det, 0.94x). On the 32 shapes that fix(cake_rmsnorm_det): fixed-order dw for the smallest token counts of the SM100/SM103 backward #8512 changes, bwd reads 0.94x to 1.18x vs det (geomean 0.999x). fix(cake_rmsnorm_det): fixed-order dw for the smallest token counts of the SM100/SM103 backward #8512 also gives the GB300 deterministic backward a new single program at T1 to T16 (15 of the 20 shapes), so these rows were also re-timed on that main (m). GB300 small-T bwd now reads 0.74x to 1.02x vs det (geomean 0.958x; 11 of 20 rows < 1, 16 of 20 within 5 % of 1.00x). In the 3 re-timed bwd rows below 0.95x, the new mode runs the two programs that det ran before fix(cake_rmsnorm_det): fixed-order dw for the smallest token counts of the SM100/SM103 backward #8512: H6144 T4 bwd 5.86 vs 4.42 us (0.75x), H6144 T16 bwd 6.02 vs 5.02 us (0.84x) and H6144 T16 res w32 bwd 6.78 vs 4.99 us (0.74x), new vs det. R200 small-T rows are at parity (bwd 0.983x to 1.019x).resandres w32. Example, B200 H2048 T1 bwd: det 3.81 us, SOTA 2.14 us (tcd), new 3.87 us. On R200, SOTA is faster on all 8 small-T bwd rows (0.518x to 0.936x vs SOTA).!).Baselines and measurement. Every row has two baselines:
deterministic=True). Its GPU programs are those of FlashInfer main at a1b4ed3. This PR does not touch the deterministic package. On themrows, det runs the GPU programs of FlashInfer main 4046a86, which contains fix(cake_rmsnorm_det): fixed-order dw for the smallest token counts of the SM100/SM103 backward #8512.new = the non-deterministic Cake kernel of this PR (
deterministic=False). The accuracy budget takes the worst of 5 seeds: y and dx atol = rtol = 1e-2; rstd rel. error <= 1e-4; dw per element atol 1e-4 + rtol 1e-4 vs FP64. New and det meet it on every row on all three GPUs.Latency = median of 804 CUPTI spans per implementation per row. One span is one CUDA-graph replay after an L2 flush (cold L2). It includes the implementation's own setup, reduction, cast and copy kernels. For each GPU type, 4 sessions ran in parallel on 4 GPUs of one node (2 in row order, 2 in reverse). Each session ran 3 interleaved passes x 67 replays. eps = 1e-6. Latencies print to 0.01 us and per-row ratios to 0.01x.
In each session the implementations run in a fixed order: det first, then the non-Cake candidates, then new. A later check on B200 and GB300 found that this order alone moves an implementation's latency by up to about 5 %: rows where both modes launch the same GPU programs read 0.95x to 1.05x (up to 1.08x on
~rows), not 1.00x. A ratio within about 5 % of 1.00x is therefore not a measured difference. All implementations of a row run in one process that also loads every non-Cake candidate, including torch.compile. A later check on B200 and GB300 found that a torch.compile function in the process makes the GPU programs of every implementation run faster: on small rows the Cake kernels and the other libraries (quack, te, torch, liger) run up to about 30 % shorter than in a process without torch.compile, and torch.compile itself runs within about 3 % of the same speed in both. So the table compares all implementations in the same state. On small rows the Cake kernels gain somewhat more than the other libraries: the ratio vs those libraries reads up to about 1.10x (GB300) and 1.14x (B200) higher than in a process without torch.compile, so it can overstate the new kernel's lead by up to that factor. On large rows this difference is at most 1.02x. Ratios vs det are not affected (both are Cake kernels in the same process).Key:
w32.resadds the BF16 output h_new.[T,24,256]has strides (12288, 512, 1), and[T,4,256]is contiguous. Both share w[256].resrows, QuACK is torch add + QuACK~: for det, SOTA or new, (max - min) / min of the 4 session medians is above 5.0 % (compared at 0.1 % resolution).*: a faster non-Cake candidate fails the dw budget. These candidates are listed below the table.!: no matched candidate meets the dw budget. The fastest matched one is shown, and the row is left out of the vs-SOTA summary.m: B200 and GB300. This row was re-timed on FlashInfer main 4046a86 after fix(cake_rmsnorm_det): fixed-order dw for the smallest token counts of the SM100/SM103 backward #8512 (same node type, software and protocol). All columns of the row come from that run.B200 (sm_100a)
B200, 148 SMs, sm_100a, compute capability 10.0, driver 580.82.07, CUDA 13.4 (nvcc cuda_13.4.r13.4/compiler.38855100_0), torch 2.14.1+cu130. SM clock median 1965 MHz (99 % of samples at max), memory clock 3996 MHz, 4 GPUs of one node, no foreign GPU processes. Rows: 78 shapes + 36 small-T shapes, x 3 ops = 342 rows. The small-T shapes are T1, T4, T16 and T64 of H2048, H2048 w32, H4096, H4096 w32, H5120, H6144, H6144 w32, H6144 res and H6144 res w32. New and det meet the budget on every row. No matched candidate meets it on the 2
!rows. 39 rows were re-timed in a second run on another B200 node with the same software, because tc did not run on them in the first run: H1536 T65536, plus T256/T1024/T4096 of H5120 w32, H6144 w32, H6144 res and H6144 res w32. All columns of those rows come from the second run.~: 117 rows!: 2 rows, [T,24,256] T65536 bwd and 2xfwd+bwd. The fastest matched candidate is shown (tc, dw err/tol 1.122).m: 110 rows: every small-T row, plus H2048 T256 bwd and 2xfwd+bwd. They come from a later run on another B200 node on FlashInfer main 4046a86, which contains fix(cake_rmsnorm_det): fixed-order dw for the smallest token counts of the SM100/SM103 backward #8512, with the same software and protocol (SM clock median 1965 MHz, 100 % of samples at max; memory clock 3996 MHz; no foreign GPU processes).The vs-SOTA columns leave out the
!rows.GB300 (sm_103a)
GB300, 152 SMs, sm_103a, compute capability 10.3, driver 580.159.03, CUDA 13.4 (same nvcc as B200), torch 2.14.1+cu130. SM clock 2070 MHz (100 % of samples at max), memory clock 3996 MHz, 4 GPUs of one node. Rows: 78 shapes + 20 small-T shapes + 9 window shapes, x 3 ops = 321 rows. The small-T shapes are T1, T4, T16 and T64 of H5120, H6144, H6144 w32, H6144 res and H6144 res w32. The window shapes are H5120 T321/T512/T896, plus T449/T768/T1152 of H6144 and H6144 w32. New, det and SOTA meet the budget on every row.
~: 86 rows!or*rowsm: 58 rows: every small-T row except H6144 T64 res bwd and 2xfwd+bwd, where the new mode on FlashInfer main runs the backward of perf(cake_rmsnorm_train): run the deterministic package's pairs in more backward windows #8626. They come from a later run on another GB300 node (same clocks; no foreign GPU processes).R200 (sm_107a)
R200, 212 SMs, sm_107a, compute capability 10.7, driver 620.62, CUDA 13.4, torch 2.14.1+cu130. SM clock 2400-2424 MHz (median 2412), memory clock 4752 MHz, 4 GPUs of one node. This is a 212-SM R200 node. 216-SM R200 boards are timed separately and never pooled with it, and that measurement is pending. Rows: 78 shapes + 8 small-T shapes, x 3 ops = 258 rows. The small-T shapes are T1, T4, T16 and T64 of H6144 res and H6144 res w32. New and det meet the budget on every row.
~: 59 rows*: 19 rows (listed below the table)*faster candidates outside the dw budget (us; dw err/tol): H4096 T65536 bwd qa 139.33 (1.246); H4096 T65536 step qa 300.19 (1.246); H5120 T32768 bwd qa 92.35, q 123.23 (1.478); H5120 T32768 step qa 195.39, q 239.23 (1.478); H5120 T65536 bwd qa 169.06 (1.355); H5120 T65536 step qa 367.87 (1.355); H6144 T65536 bwd qa 194.21, q 228.54 (2.290); H6144 T65536 step qa 432.74 (2.290); H6144 T65536 res bwd qa 252.51, q 284.64 (2.042); H5120 T32768 w32 bwd qa 90.98, q 121.38 (1.339); H5120 T32768 w32 step qa 204.58 (1.339); H5120 T65536 w32 bwd qa 165.82, q 227.26 (1.327); H5120 T65536 w32 step qa 383.01 (1.327); H6144 T65536 w32 bwd qa 194.75 (1.123); H6144 T65536 w32 step qa 446.31 (1.123); H6144 T65536 res w32 bwd qa 252.75, q 280.87 (1.751); H6144 T65536 res w32 step qa 868.56 (1.751); [T,24,256] T65536 bwd tc 238.02 (4.454); [T,24,256] T65536 step tc 488.99, tcd 639.19 (4.454/5.751).The family has no other variants: it ships these forms in both modes on sm_100a, sm_103a and sm_107a, and all three are measured above.
Summary by CodeRabbit