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_gdn_kcp): plain split views of packed Q/K/V get contiguous dq / dk / dv in the non-deterministic KCP backward - #8435
flashinfer.experimental.cake_gdn_train_kcp is the non-deterministic KCP GDN training package. It supports BF16, the contiguous, zigzag and zigzag-per-document KCP layouts, and runs with or without in-kernel Q/K L2 norm. It covers forward, backward and checkpointed calls, in eager and CUDA graph modes. This PR changes only the host side. When q / k / v are plain split views of one packed QKV tensor (row.split + view), the backward now gives three contiguous dq / dk / dv to torch's split backward. The split-view backward then costs one contiguous concatenation instead of the slower path on main. Only q / k / v that come from one packed_qkv_views_kcp (or packed_qkv_views_det_kcp) split node keep the single packed dQKV row. There are no device code changes, and outputs and gradients are bitwise unchanged. Addresses #8367.
All rows use the PR head after it merged main, so the forward runs the fused forward output walk.
Compared with the same package on FlashInfer main (CP1, norm off), the plain-split-view backward gets faster:
B200: eager ckpt goes from 5.930 ms to 4.992 ms (1.188x).
GB300: eager bwd goes from 4.019 ms to 3.142 ms (1.279x).
R200 (208 SMs): eager bwd goes from 3.037 ms to 2.264 ms (1.341x).
Packed and contiguous forms are at parity (0.998-1.006x across all rows).
Against the deterministic Cake package in the same process (speedup = baseline ms / this PR ms), this PR is faster with plain split views:
eager bwd: 1.255-1.277x on B200, 1.255-1.280x on GB300 and 1.313-1.338x on R200.
eager ckpt: 1.222-1.245x on B200, 1.228-1.254x on GB300 and 1.261-1.286x on R200.
eager fwd: 1.170-1.192x on B200, 1.192-1.219x on GB300 and 1.172-1.198x on R200.
With packed views, bwd is at parity (1.000-1.003x), ckpt is 1.051-1.065x and fwd is 1.167-1.219x. Against the previous version, this PR's fwd and packed-view times are at parity. So the fwd and packed-view ckpt gains over deterministic Cake do not come from this PR.
Against the strongest qualified non-Cake implementation (both Q/K/V forms, norm off and on), this PR is slower on every claimed row:
GB300: versus cuDNN FROST, the speedup is 0.378-0.387x (eager fwd), 0.663-0.739x (eager bwd) and 0.764-0.821x (eager ckpt).
R200: versus cuDNN FROST, the speedup is 0.434-0.438x (eager fwd), 0.889-0.995x (eager bwd) and 0.928-0.993x (eager ckpt). FlashQLA does not support sm_107a.
B200: cuDNN FROST is faster than this PR in all 24 rows but fails the accuracy check (dbeta max-abs 0.135 of max|ref|, limit 0.02). The strongest qualified arm is FlashQLA 0.1.3. It runs in a separate process, so every B200 non-Cake ratio is shown but not claimed.
B200 (sm_100a, 148 SMs)
NVIDIA B200, sm_100a, 148 SMs. PyTorch 2.14.1 (cu130, CUDA 13.0 runtime), CUDA 13.4 toolkit, BF16. Workload: (H_qk, HV) = (16, 48), D = 128. CP1 = one rank, 32,768 tokens in five documents [12000, 9000, 6000, 3768, 2000], contiguous layout. CP4 zigzag = 4 ranks, documents [49152, 32768, 28672, 20480], 32,768 tokens per rank; each sample takes the slowest rank. Q/K/V forms: split = row.split + view (plain views), packed = packed_qkv_views_kcp, contig = contiguous copies. Modes: fwd = forward only; bwd = autograd backward after an eager forward. ckpt = torch.utils.checkpoint(use_reentrant=False), 2 forwards + 1 backward. graph = the package's graph="replay" mode with a caller-held workspace. Timing: 2 runs x 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls. Run 2 reverses the arm order. The shown value is the median of group medians. The previous version is a byte copy of the package from FlashInfer main before this PR. It runs in the same process as this PR, with the same build and libraries. Speedup = previous ms / this PR ms. Power-capped = GPU software power-cap share >= 5 % of an arm's timed-block samples, for any arm of a mode and norm setting. When the short runs hit the cap, the rows of that mode come from a sustained protocol instead. The sustained protocol uses 3 palindromic block pairs (block = 5 s idle + 20 warm-up calls + 300 timed calls). On B200, this applies to the CP1 norm-off eager fwd and bwd rows, all CP1 graph rows and the CP4 graph bwd norm-on rows. In the sustained runs, the CP1 norm-off eager fwd, eager bwd and graph bwd rows still hit the cap. They are marked "power-capped, not claimed". The CP1 norm-off eager fwd rows also fail the stability check (drift = last / first group - 1; every arm of the mode and norm setting must stay within 2 %; the previous version's split arm drifted 2.26 %). They are also marked "timing unstable". All other B200 rows are claimed.
CP1
GPU / architecture
workload / shape / dtype / mode
operation (Q/K/V form, norm)
previous version ms
this PR ms
speedup vs previous
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
split, norm off
4.326
3.389
1.276 (power-capped, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
packed, norm off
3.186
3.187
1.000 (power-capped, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
contig, norm off
3.388
3.388
1.000 (power-capped, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
split, norm on
4.611
3.680
1.253
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
packed, norm on
3.474
3.474
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
contig, norm on
3.678
3.679
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
split, norm off
5.930
4.992
1.188
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
packed, norm off
4.796
4.795
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
contig, norm off
6.210
6.210
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
split, norm on
6.388
5.446
1.173
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
packed, norm on
5.240
5.248
0.998 (this PR slower)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
contig, norm on
6.625
6.627
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
split, norm off
4.321
3.385
1.277 (power-capped, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
packed, norm off
3.180
3.178
1.001 (power-capped, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
split, norm on
4.639
3.672
1.263
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
packed, norm on
3.467
3.464
1.001
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
split, norm off
6.967
6.024
1.157
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
packed, norm off
5.829
5.831
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
split, norm on
7.516
6.549
1.148
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
packed, norm on
6.352
6.353
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
split, norm off
1.444
1.436
1.006 (power-capped, timing unstable, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
packed, norm off
1.431
1.432
0.999 (this PR slower; power-capped, timing unstable, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
contig, norm off
1.207
1.207
1.000 (power-capped, timing unstable, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
split, norm on
1.619
1.619
1.000
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
packed, norm on
1.608
1.611
0.998 (this PR slower)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
contig, norm on
1.332
1.331
1.001
CP4 zigzag
GPU / architecture
workload / shape / dtype / mode
operation (Q/K/V form, norm)
previous version ms
this PR ms
speedup vs previous
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
split, norm off
4.842
3.884
1.247
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
packed, norm off
3.676
3.679
0.999 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
contig, norm off
3.881
3.883
0.999 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
split, norm on
5.143
4.182
1.230
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
packed, norm on
3.976
3.975
1.000
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
contig, norm on
4.179
4.179
1.000
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
split, norm off
7.769
6.811
1.141
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
packed, norm off
6.608
6.614
0.999 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
contig, norm off
8.018
8.021
1.000
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
split, norm on
8.241
7.287
1.131
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
packed, norm on
7.088
7.091
1.000
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
contig, norm on
8.440
8.441
1.000
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
split, norm off
4.853
3.885
1.249
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
packed, norm off
3.679
3.676
1.001
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
split, norm on
5.151
4.188
1.230
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
packed, norm on
3.977
3.979
0.999 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
split, norm off
9.032
8.066
1.120
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
packed, norm off
7.865
7.858
1.001
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
split, norm on
9.557
8.611
1.110
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
packed, norm on
8.410
8.414
1.000
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
split, norm off
2.138
2.141
0.999 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
packed, norm off
2.126
2.131
0.998 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
contig, norm off
1.892
1.894
0.999 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
split, norm on
2.318
2.316
1.001
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
packed, norm on
2.306
2.308
0.999 (this PR slower)
B200 / sm_100a (148 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
contig, norm on
2.015
2.015
1.000
Two-baseline comparison (CP1)
These are separate runs from the previous-version tables above, so this PR's absolute times differ slightly between the two sets. Each run uses one process for deterministic Cake, this PR, FLA and cuDNN FROST. Run 2 reverses the arm order of run 1. Each run has 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls per arm. The shown value is the median of group medians. The norm-off eager fwd and bwd rows come from two sustained runs (3 block pairs; block = 5 s idle + 20 warm-up calls + 300 timed calls). No deterministic Cake, this PR or FLA cell is power-capped (every in-process arm < 5 %); the separate FlashQLA process has no power-cap record. Deterministic Cake is cake_gdn_det_train_kcp (chunk_gdn_det_kcp) on the same FlashInfer tree, unchanged by this PR. Its packed views use its own helper, packed_qkv_views_det_kcp. Its outputs and gradients are bitwise equal to this PR's. Strongest non-Cake is the fastest arm per row that passes the accuracy check. The check needs rel-RMS <= 1e-2 and max-abs <= 2e-2 of max|ref| on every tensor, with finite values. Candidates are FLA 0.6.0 Triton chunk_gated_delta_rule, cuDNN frontend 1.32.1 gated_delta_net and FlashQLA 0.1.3. The cuDNN plan is gdn_frost (cuDNN FROST): ck0 = the backward recomputes the states; ck64 = the forward stores BF16 states every 64 tokens. In "FlashQLA + FLA l2norm outside", FLA's Triton l2norm normalizes q / k outside the kernel (timed), and the in-kernel flag is off. On B200, cuDNN FROST fails the check on dbeta (max-abs 0.135 of max|ref|), although it is faster than this PR in all 24 rows. FlashQLA is therefore the strongest qualified arm. FlashQLA cannot share a process with the Cake package, because its TileLang 0.1.15 needs apache-tvm-ffi < 0.1.13. Its values come from a separate process that used the short protocol. Every B200 non-Cake ratio is therefore marked "cross-run, not claimed".
GPU / architecture
workload / shape / dtype / mode
operation
run
deterministic Cake ms
non-Cake ms (arm)
this PR ms
speedup vs deterministic Cake
speedup vs non-Cake
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager fwd
1
1.723
1.675 (FlashQLA)
1.445
1.192
1.159 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager fwd
2
1.721
1.675 (FlashQLA)
1.447
1.189
1.158 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager fwd
1
1.701
1.675 (FlashQLA)
1.436
1.185
1.166 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager fwd
2
1.705
1.675 (FlashQLA)
1.441
1.183
1.162 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager bwd
1
4.317
3.156 (FlashQLA)
3.385
1.275
0.932 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager bwd
2
4.325
3.150 (FlashQLA)
3.387
1.277
0.930 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager bwd
1
3.185
3.156 (FlashQLA)
3.181
1.001
0.992 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager bwd
2
3.189
3.150 (FlashQLA)
3.185
1.001
0.989 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager ckpt (2 fwd + bwd)
1
6.208
6.557 (FlashQLA)
4.996
1.243
1.312 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager ckpt (2 fwd + bwd)
2
6.194
6.549 (FlashQLA)
4.977
1.245
1.316 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager ckpt (2 fwd + bwd)
1
5.049
6.557 (FlashQLA)
4.800
1.052
1.366 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager ckpt (2 fwd + bwd)
2
5.039
6.549 (FlashQLA)
4.787
1.053
1.368 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager fwd
1
1.909
1.693 (FlashQLA + FLA l2norm outside)
1.632
1.170
1.037 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager fwd
2
1.902
1.690 (FlashQLA + FLA l2norm outside)
1.621
1.173
1.043 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager fwd
1
1.891
1.693 (FlashQLA + FLA l2norm outside)
1.621
1.167
1.044 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager fwd
2
1.883
1.690 (FlashQLA + FLA l2norm outside)
1.614
1.167
1.047 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager bwd
1
4.613
3.270 (FlashQLA + FLA l2norm outside)
3.676
1.255
0.890 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager bwd
2
4.612
3.263 (FlashQLA + FLA l2norm outside)
3.675
1.255
0.888 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager bwd
1
3.477
3.270 (FlashQLA + FLA l2norm outside)
3.471
1.002
0.942 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager bwd
2
3.478
3.263 (FlashQLA + FLA l2norm outside)
3.469
1.003
0.941 (REGRESSION: non-Cake faster; cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager ckpt (2 fwd + bwd)
1
6.675
6.781 (FlashQLA + FLA l2norm outside)
5.462
1.222
1.241 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager ckpt (2 fwd + bwd)
2
6.654
6.763 (FlashQLA + FLA l2norm outside)
5.438
1.224
1.244 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager ckpt (2 fwd + bwd)
1
5.529
6.781 (FlashQLA + FLA l2norm outside)
5.261
1.051
1.289 (cross-run, not claimed)
B200 / sm_100a (148 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager ckpt (2 fwd + bwd)
2
5.510
6.763 (FlashQLA + FLA l2norm outside)
5.239
1.052
1.291 (cross-run, not claimed)
GB300 (sm_103a, 152 SMs)
NVIDIA GB300, sm_103a, 152 SMs. PyTorch 2.14.1 (cu130, CUDA 13.0 runtime), CUDA 13.4 toolkit, BF16. Workload: (H_qk, HV) = (16, 48), D = 128. CP1 = one rank, 32,768 tokens in five documents [12000, 9000, 6000, 3768, 2000], contiguous layout. CP4 zigzag = 4 ranks, documents [49152, 32768, 28672, 20480], 32,768 tokens per rank; each sample takes the slowest rank. Q/K/V forms: split = row.split + view (plain views), packed = packed_qkv_views_kcp, contig = contiguous copies. Modes: fwd = forward only; bwd = autograd backward after an eager forward. ckpt = torch.utils.checkpoint(use_reentrant=False), 2 forwards + 1 backward. graph = the package's graph="replay" mode with a caller-held workspace. Timing: 2 runs x 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls. Run 2 reverses the arm order. The shown value is the median of group medians. The previous version is a byte copy of the package from FlashInfer main before this PR. It runs in the same process as this PR, with the same build and libraries. Speedup = previous ms / this PR ms. No GB300 row is power-capped (GPU software power-cap share >= 5 %). In the CP4 eager ckpt norm-on rows, the time level shifted within run 2: the drift (last / first group - 1) was 2.2-8.5 % there, and run 1 stayed within 0.3 %. These rows pass the paired check instead: every paired group ratio stays within ±1 % of its median. They are claimed and marked "level shift in run 2" with the paired ratio. The CP4 eager fwd norm-off rows come from a re-measure with two further runs.
CP1
GPU / architecture
workload / shape / dtype / mode
operation (Q/K/V form, norm)
previous version ms
this PR ms
speedup vs previous
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
split, norm off
4.019
3.142
1.279
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
packed, norm off
2.946
2.947
1.000
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
contig, norm off
3.139
3.140
1.000
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
split, norm on
4.316
3.440
1.255
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
packed, norm on
3.243
3.244
1.000
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
contig, norm on
3.434
3.436
0.999 (this PR slower)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
split, norm off
5.444
4.562
1.193
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
packed, norm off
4.384
4.387
0.999 (this PR slower)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
contig, norm off
5.770
5.772
1.000
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
split, norm on
5.916
5.039
1.174
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
packed, norm on
4.851
4.857
0.999 (this PR slower)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
contig, norm on
6.176
6.177
1.000
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
split, norm off
4.024
3.144
1.280
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
packed, norm off
2.951
2.947
1.001
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
split, norm on
4.331
3.441
1.259
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
packed, norm on
3.247
3.241
1.002
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
split, norm off
6.416
5.532
1.160
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
packed, norm off
5.353
5.349
1.001
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
split, norm on
6.923
6.040
1.146
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
packed, norm on
5.860
5.857
1.001
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
split, norm off
1.299
1.302
0.998 (this PR slower)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
packed, norm off
1.297
1.298
0.999 (this PR slower)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
contig, norm off
1.120
1.120
1.000
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
split, norm on
1.467
1.469
0.999 (this PR slower)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
packed, norm on
1.463
1.465
0.999 (this PR slower)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
contig, norm on
1.227
1.227
1.000
CP4 zigzag
GPU / architecture
workload / shape / dtype / mode
operation (Q/K/V form, norm)
previous version ms
this PR ms
speedup vs previous
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
split, norm off
4.479
3.580
1.251
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
packed, norm off
3.386
3.384
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
contig, norm off
3.579
3.578
1.000
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
split, norm on
4.779
3.872
1.234
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
packed, norm on
3.681
3.677
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager bwd
contig, norm on
3.874
3.872
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
split, norm off
7.162
6.264
1.143
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
packed, norm off
6.087
6.083
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
contig, norm off
7.428
7.426
1.000
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
split, norm on
7.625
6.731
1.133 (level shift in run 2; paired 1.133 [1.122, 1.136])
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
packed, norm on
6.555
6.555
1.000 (level shift in run 2; paired 1.000 [0.991, 1.002])
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager ckpt
contig, norm on
7.838
7.830
1.001 (level shift in run 2; paired 1.001 [0.999, 1.001])
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
split, norm off
4.485
3.583
1.252
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
packed, norm off
3.388
3.386
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
split, norm on
4.755
3.873
1.228
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph bwd
packed, norm on
3.679
3.677
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
split, norm off
8.323
7.421
1.122
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
packed, norm off
7.242
7.232
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
split, norm on
8.890
8.009
1.110
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, graph ckpt
packed, norm on
7.832
7.831
1.000
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
split, norm off
1.968
1.967
1.001
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
packed, norm off
1.959
1.963
0.998 (this PR slower)
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
contig, norm off
1.745
1.747
0.999 (this PR slower)
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
split, norm on
2.139
2.131
1.004
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
packed, norm on
2.126
2.127
1.000
GB300 / sm_103a (152 SMs)
CP4 zigzag, (16, 48), D 128, BF16, eager fwd
contig, norm on
1.853
1.851
1.001
Two-baseline comparison (CP1)
These are separate runs from the previous-version tables above, so this PR's absolute times differ slightly between the two sets. Each run uses one process for deterministic Cake, this PR, FLA and cuDNN FROST. Run 2 reverses the arm order of run 1. Each run has 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls per arm. The shown value is the median of group medians. No cell is power-capped or unstable. Deterministic Cake is cake_gdn_det_train_kcp (chunk_gdn_det_kcp) on the same FlashInfer tree, unchanged by this PR. Its packed views use its own helper, packed_qkv_views_det_kcp. Its outputs and gradients are bitwise equal to this PR's. Strongest non-Cake is the fastest arm per row that passes the accuracy check. The check needs rel-RMS <= 1e-2 and max-abs <= 2e-2 of max|ref| on every tensor, with finite values. Candidates are FLA 0.6.0 Triton chunk_gated_delta_rule, cuDNN frontend 1.32.1 gated_delta_net and FlashQLA 0.1.3. The cuDNN plan is gdn_frost (cuDNN FROST): ck0 = the backward recomputes the states; ck64 = the forward stores BF16 states every 64 tokens. On GB300, cuDNN FROST passes the check and is the strongest arm on every row: ck0 for fwd, ck64 for bwd and ckpt. FlashQLA ran in a separate process and was not faster than the shown arm in any row.
GPU / architecture
workload / shape / dtype / mode
operation
run
deterministic Cake ms
non-Cake ms (arm)
this PR ms
speedup vs deterministic Cake
speedup vs non-Cake
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager fwd
1
1.609
0.511 (cuDNN FROST ck0)
1.323
1.216
0.386 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager fwd
2
1.589
0.493 (cuDNN FROST ck0)
1.304
1.219
0.378 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager fwd
1
1.608
0.511 (cuDNN FROST ck0)
1.321
1.217
0.387 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager fwd
2
1.588
0.493 (cuDNN FROST ck0)
1.303
1.219
0.378 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager bwd
1
4.021
2.101 (cuDNN FROST ck64)
3.141
1.280
0.669 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager bwd
2
4.018
2.085 (cuDNN FROST ck64)
3.143
1.278
0.663 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager bwd
1
2.949
2.101 (cuDNN FROST ck64)
2.948
1.000
0.713 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager bwd
2
2.950
2.085 (cuDNN FROST ck64)
2.950
1.000
0.707 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager ckpt (2 fwd + bwd)
1
5.769
3.634 (cuDNN FROST ck64)
4.601
1.254
0.790 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager ckpt (2 fwd + bwd)
2
5.725
3.489 (cuDNN FROST ck64)
4.568
1.253
0.764 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager ckpt (2 fwd + bwd)
1
4.714
3.634 (cuDNN FROST ck64)
4.426
1.065
0.821 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager ckpt (2 fwd + bwd)
2
4.678
3.489 (cuDNN FROST ck64)
4.398
1.064
0.793 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager fwd
1
1.784
0.579 (cuDNN FROST ck0)
1.497
1.192
0.387 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager fwd
2
1.749
0.562 (cuDNN FROST ck0)
1.467
1.192
0.383 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager fwd
1
1.780
0.579 (cuDNN FROST ck0)
1.497
1.189
0.387 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager fwd
2
1.748
0.562 (cuDNN FROST ck0)
1.465
1.193
0.384 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager bwd
1
4.314
2.398 (cuDNN FROST ck64)
3.437
1.255
0.698 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager bwd
2
4.318
2.399 (cuDNN FROST ck64)
3.439
1.256
0.698 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager bwd
1
3.243
2.398 (cuDNN FROST ck64)
3.243
1.000
0.739 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager bwd
2
3.247
2.399 (cuDNN FROST ck64)
3.247
1.000
0.739 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager ckpt (2 fwd + bwd)
1
6.232
3.982 (cuDNN FROST ck64)
5.074
1.228
0.785 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager ckpt (2 fwd + bwd)
2
6.193
3.890 (cuDNN FROST ck64)
5.032
1.231
0.773 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager ckpt (2 fwd + bwd)
1
5.182
3.982 (cuDNN FROST ck64)
4.897
1.058
0.813 (REGRESSION: non-Cake faster)
GB300 / sm_103a (152 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager ckpt (2 fwd + bwd)
2
5.142
3.890 (cuDNN FROST ck64)
4.860
1.058
0.800 (REGRESSION: non-Cake faster)
R200 (sm_107a, 208 SMs)
NVIDIA R200 board, sm_107a, 208 SMs; these results cover the 208-SM board only. PyTorch 2.14.1 (cu130, CUDA 13.0 runtime), CUDA 13.4 toolkit, BF16. Workload: (H_qk, HV) = (16, 48), D = 128. CP1 = one rank, 32,768 tokens in five documents [12000, 9000, 6000, 3768, 2000], contiguous layout. Only one GPU was visible, so the 4-rank CP4 zigzag rows were not run on this board. Q/K/V forms: split = row.split + view (plain views), packed = packed_qkv_views_kcp, contig = contiguous copies. Modes: fwd = forward only; bwd = autograd backward after an eager forward. ckpt = torch.utils.checkpoint(use_reentrant=False), 2 forwards + 1 backward. graph = the package's graph="replay" mode with a caller-held workspace. Timing: 2 runs x 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls. Run 2 reverses the arm order. The shown value is the median of group medians. The previous version is a byte copy of the package from FlashInfer main before this PR. It runs in the same process as this PR, with the same build and libraries. Speedup = previous ms / this PR ms. No R200 row is power-capped (GPU software power-cap share >= 5 %).
CP1
GPU / architecture
workload / shape / dtype / mode
operation (Q/K/V form, norm)
previous version ms
this PR ms
speedup vs previous
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
split, norm off
3.037
2.264
1.341
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
packed, norm off
2.091
2.090
1.000
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
contig, norm off
2.267
2.266
1.000
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
split, norm on
3.232
2.460
1.314
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
packed, norm on
2.290
2.289
1.000
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager bwd
contig, norm on
2.464
2.462
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
split, norm off
4.082
3.309
1.234
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
packed, norm off
3.143
3.144
1.000
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
contig, norm off
5.054
5.049
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
split, norm on
4.414
3.641
1.212
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
packed, norm on
3.476
3.474
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager ckpt
contig, norm on
5.351
5.346
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
split, norm off
3.038
2.266
1.341
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
packed, norm off
2.090
2.088
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
split, norm on
3.233
2.462
1.313
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph bwd
packed, norm on
2.288
2.285
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
split, norm off
4.748
3.976
1.194
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
packed, norm off
3.802
3.799
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
split, norm on
5.128
4.353
1.178
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, graph ckpt
packed, norm on
4.184
4.181
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
split, norm off
0.947
0.945
1.002
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
packed, norm off
0.944
0.944
1.000
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
contig, norm off
0.805
0.803
1.002
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
split, norm on
1.075
1.075
1.000
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
packed, norm on
1.073
1.072
1.001
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, eager fwd
contig, norm on
0.900
0.895
1.006
Two-baseline comparison (CP1)
These are separate runs on a different R200 machine from the previous-version tables, so this PR's absolute times differ slightly between the two sets. Each run uses one process for deterministic Cake, this PR, FLA and cuDNN FROST. Run 2 reverses the arm order of run 1. Each run has 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls per arm. The shown value is the median of group medians. The norm-on rows come from a re-measure: two further runs with the same design. No cell is power-capped or unstable. Deterministic Cake is cake_gdn_det_train_kcp (chunk_gdn_det_kcp) on the same FlashInfer tree, unchanged by this PR. Its packed views use its own helper, packed_qkv_views_det_kcp. Its outputs and gradients are bitwise equal to this PR's. Strongest non-Cake is the fastest arm per row that passes the accuracy check. The check needs rel-RMS <= 1e-2 and max-abs <= 2e-2 of max|ref| on every tensor, with finite values. Candidates are FLA 0.6.0 Triton chunk_gated_delta_rule, cuDNN frontend 1.32.1 gated_delta_net and FlashQLA 0.1.3. The cuDNN plan is gdn_frost (cuDNN FROST): ck0 = the backward recomputes the states; ck64 = the forward stores BF16 states every 64 tokens. On R200, cuDNN FROST passes the check and is the strongest arm on every row: ck0 for fwd, ck64 for bwd and ckpt. FlashQLA does not support sm_107a; it accepts only sm90, sm100, sm103, sm120 and sm121. So R200 has no FlashQLA value.
GPU / architecture
workload / shape / dtype / mode
operation
run
deterministic Cake ms
non-Cake ms (arm)
this PR ms
speedup vs deterministic Cake
speedup vs non-Cake
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager fwd
1
1.134
0.415 (cuDNN FROST ck0)
0.948
1.196
0.438 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager fwd
2
1.136
0.414 (cuDNN FROST ck0)
0.948
1.198
0.437 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager fwd
1
1.131
0.415 (cuDNN FROST ck0)
0.948
1.193
0.438 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager fwd
2
1.131
0.414 (cuDNN FROST ck0)
0.946
1.196
0.438 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager bwd
1
3.013
2.004 (cuDNN FROST ck64)
2.254
1.337
0.889 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager bwd
2
3.017
2.004 (cuDNN FROST ck64)
2.255
1.338
0.889 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager bwd
1
2.085
2.004 (cuDNN FROST ck64)
2.084
1.000
0.962 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager bwd
2
2.087
2.004 (cuDNN FROST ck64)
2.086
1.000
0.961 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager ckpt (2 fwd + bwd)
1
4.253
3.077 (cuDNN FROST ck64)
3.313
1.284
0.929 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm off
eager ckpt (2 fwd + bwd)
2
4.256
3.073 (cuDNN FROST ck64)
3.310
1.286
0.928 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager ckpt (2 fwd + bwd)
1
3.330
3.077 (cuDNN FROST ck64)
3.148
1.058
0.977 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm off
eager ckpt (2 fwd + bwd)
2
3.329
3.073 (cuDNN FROST ck64)
3.149
1.057
0.976 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager fwd
1
1.252
0.463 (cuDNN FROST ck0)
1.068
1.172
0.434 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager fwd
2
1.256
0.465 (cuDNN FROST ck0)
1.069
1.175
0.435 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager fwd
1
1.247
0.463 (cuDNN FROST ck0)
1.066
1.170
0.434 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager fwd
2
1.250
0.465 (cuDNN FROST ck0)
1.066
1.173
0.436 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager bwd
1
3.206
2.261 (cuDNN FROST ck64)
2.442
1.313
0.926 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager bwd
2
3.208
2.263 (cuDNN FROST ck64)
2.443
1.313
0.926 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager bwd
1
2.279
2.261 (cuDNN FROST ck64)
2.274
1.002
0.994 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager bwd
2
2.281
2.263 (cuDNN FROST ck64)
2.274
1.003
0.995 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager ckpt (2 fwd + bwd)
1
4.560
3.425 (cuDNN FROST ck64)
3.617
1.261
0.947 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, plain split views, norm on
eager ckpt (2 fwd + bwd)
2
4.564
3.430 (cuDNN FROST ck64)
3.618
1.261
0.948 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager ckpt (2 fwd + bwd)
1
3.638
3.425 (cuDNN FROST ck64)
3.454
1.053
0.992 (REGRESSION: non-Cake faster)
R200 / sm_107a (208 SMs)
CP1, (16, 48), D 128, BF16, packed views, norm on
eager ckpt (2 fwd + bwd)
2
3.642
3.430 (cuDNN FROST ck64)
3.454
1.054
0.993 (REGRESSION: non-Cake faster)
This PR does not change the other Cake variants (the deterministic KCP package cake_gdn_det_train_kcp and the single-rank packages); architecture coverage stays sm_100a, sm_103a and sm_107a.
Summary by CodeRabbit
New Features
Backward gradients can now share a packed layout when Q/K/V inputs use a compatible packed view. Ordinary split or view inputs continue to produce separate gradients.
Documentation
Updated performance measurements for packed versus ordinary split views, including results for additional hardware and revised benchmark details.
Clarified when packed gradient handling applies and how gradient layouts differ between packed views and ordinary split/view inputs.
…ation in the non-deterministic KCP backward; use packed_qkv_views_kcp (flashinfer-ai#8367)
The non-deterministic mirror of the deterministic package's flashinfer-ai#8367 change.
Plain split / view head views of one packed Q/K/V row run the same programs
as packed_qkv_views_kcp views, but torch's split backward concatenates
dq / dk / dv into a new row: about 1.1 ms (33-36 %) more backward per call at
(16, 48) and 32,768 tokens on B200 and GB300. The docstrings of
packed_qkv_views_kcp and chunk_gdn_kcp (package and the facade's
deterministic=False route) and the package README say so, with the measured
cost. The autograd path is unchanged.
Tests: plain split views, packed views and contiguous operands with the
in-kernel norm off and on, direct and checkpointed: the hooks and
retain_grad on the views (and the split outputs) fire once with the
gradient; bitwise on the fake backend, within the deterministic route's
budgets on the device programs (CP1 rows and the CP <= 4 / CP8 dist rows);
packed views refuse in-place changes.
…f the split-view cost; check retain_grad on the split outputs (flashinfer-ai#8367)
README "Packed Q/K/V": the backward skips the concatenation only when qkv is
dense over [..., T, (2 H_qk + HV) * 128] (token stride exactly
(2 H_qk + HV) * 128) and the call keeps layout="tma"; the protocol caption
says a group is 2 blocks of 5 warm-up + 20 timed calls per view form; B200
rows for the norm off and on, measured on an exclusive node, replace the
shared-node row. The input bullet and the docstrings of
packed_qkv_views_kcp and chunk_gdn_kcp give the same conditions and the
range about 1.1-1.2 ms (33-37 %).
The facade's chunk_gdn_kcp docstring says the views of either package's
helper give the full backward speed, states the cost for the
non-deterministic route and cites only this package's README.
Tests: the split outputs' retain_grad gradient is asserted in cp1, l2n and
dist; the cp1 contiguous arm no longer checks .grad of q / k / v, which are
the leaves the row gradient is read from.
…dq / dk / dv in the non-deterministic KCP backward (flashinfer-ai#8367)
The non-deterministic KCP backward chose the packed dQKV gradient layout from
the operands' pointers and strides alone, so plain split / view views of a
packed Q/K/V row received three column views of one packed row as well. Only
the split node of packed_qkv_views_kcp hands such a row on as is; torch's split
backward concatenates the three gradients, and a concatenation of three
strided column views runs at a fraction of the bandwidth of one of three
contiguous tensors.
chunk_gdn_kcp now reads the gradients' consumer in the original forward: the
packed layout is kept only when q / k / v are the three outputs, in order, of
packed_qkv_views_kcp's split node (or of the deterministic package's
packed_qkv_views_det_kcp node, looked up without importing it); every other
form (plain split / view / slice views, leaf views) gets three contiguous
gradients. The decision lives on the autograd node, so a non-reentrant
checkpoint keeps the original forward's choice and a reentrant re-run makes its
own; it is part of the backward's CUDA-graph key. The explicit
chunk_gdn_bwd_kcp keeps the pointer rule. The programs, their operands and
the gradient values are unchanged; only where dq / dk / dv are written moves.
Tests: the layout per view form and checkpoint form (direct, non-reentrant,
reentrant, recompute handle) with bitwise values on the fake backend, the
explicit backward's rule, the graph key, and a CUDA-graph replay that
alternates packed and plain views.
…for plain views (flashinfer-ai#8367)
Plain split / view views of a packed Q/K/V row now get three contiguous
gradients in the non-deterministic KCP backward: the single-GPU and fake-backend
split-view tests expect grad_layout "separate" for them ("packed" only for
packed_qkv_views_kcp views), and the multi-rank rows check it per rank. Hooks,
retain_grad, PACKED_SPLIT_STATS and the values are unchanged.
…ntiguous concatenation; B200, GB300 and R200 rows (flashinfer-ai#8367)
The non-deterministic KCP backward writes one packed dQKV row only for the
outputs of packed_qkv_views_kcp (or packed_qkv_views_det_kcp). Plain
split / view views get three contiguous gradients, which torch's split
backward concatenates. Update the README bullet, the Packed Q/K/V section
and its table, the chunk_gdn_kcp docstring and the facade docstring to the
measured cost: about 0.2 ms (6-7 %) more backward than packed views on
B200 and GB300 and 0.17 ms (7-8 %) on R200 (208 SMs) at (16, 48) and
32,768 tokens, with the in-kernel norm off or on.
No actionable comments were generated in the recent review.
Walkthrough
The GDN KCP backward path now selects packed or separate Q/K/V gradients based on the recognized input split. The change carries this selection through autograd, backward allocation, and graph-cache keys. Tests and documentation describe the layouts and updated timing results.
Changes: Packed QKV gradients
Layer / File(s)
Summary
Detect and propagate packed-gradient eligibility flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py
chunk_gdn_kcp recognizes Q/K/V from the package’s packed split or the deterministic package’s split, then passes the packed-gradient setting through autograd to backward.
Allocate gradients and key backward graphs flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py
Backward allocation creates views into one packed row when the setting and storage are compatible; otherwise it allocates separate contiguous gradients. The graph-cache key includes the gradient-layout setting.
flashinfer-ai/flashinfer#8367 — The issue describes the plain-split gradient concatenation overhead addressed by this change.
flashinfer-ai/flashinfer#8063 — This change covers packed-gradient views in the package, but not the issue’s full requested non-deterministic implementation and qualification scope.
Possibly related PRs
flashinfer-ai/flashinfer#8393 — Documents the same packed-Q/K/V backward behavior for the non-deterministic KCP route.
No concrete merge-blocking issue is established for the gradient-layout change. Merge after normal checks pass.
Pre-merge checks: 4 · 1
Failed checks: 1 warning
Check name
Status
Explanation
Resolution
Docstring Coverage
Docstring coverage is 72.73% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…
Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed checks: 4 passed
Check name
Status
Explanation
Title check
The title names the main change: in the non-deterministic KCP backward, plain split views of packed Q/K/V now receive contiguous dq, dk and dv gradients. The raw summary, objectives and PR description…
Description check
The description covers the purpose, the scope of the change, the workload and benchmark methodology, and the performance results across B200, GB300 and R200. It links issue #8367. It explains the pack…
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.
Full details: Title check
Explanation
The title names the main change: in the non-deterministic KCP backward, plain split views of packed Q/K/V now receive contiguous dq, dk and dv gradients. The raw summary, objectives and PR description all describe this host-side change in the backward gradient layout. The title uses the conventional perf(...) prefix and stays under one sentence. It is long, at 119 characters, but it is specific and relevant.
Full details: Description check
Explanation
The description covers the purpose, the scope of the change, the workload and benchmark methodology, and the performance results across B200, GB300 and R200. It links issue #8367. It explains the packed-versus-split gradient behavior and states that outputs and gradients are bitwise unchanged. It does not fill in the template's test checklist, the experimental-track checklist, or the experimental-tests fence. Those sections are the template's own gating items and are left as placeholders. The description also includes a very large benchmark section that makes it hard to scan. Its core content is complete, so the check passes, but the template gaps should be closed before merge.
Full details: Docstring Coverage
Explanation
Docstring coverage is 72.73% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 22 functions across 6 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 1 remain after this review.
…he non-deterministic KCP package (#8521)
`flashinfer.experimental.cake_gdn_train_kcp` is the non-deterministic
KCP GDN training package (BF16; contiguous, zigzag and
zigzag-per-document layouts; optional in-kernel Q/K L2 norm; forward,
backward and checkpointed calls, eager and CUDA graph). This PR changes
only the host rule that picks the per-CTA value width (128, or 64 with
two CTAs per segment and head) for the forward and backward state walks:
the width now follows the plan's chain lengths, except that on sm_100a a
backward walk with a segment of at least 400 chunks keeps the previous
rule when the two-CTA grid would finish later. Kernels are unchanged;
outputs and gradients are bitwise equal to the previous version (eager
and CUDA graph).
**Conventions**
- Ratio = baseline ms / this PR ms. Above 1 means this PR is faster.
Below 1 (marked `*` in the tables) means this PR is slower.
- Workload label = shape letter (shapes are listed in the table legend),
then CPn = context parallel over n GPUs, then the layout: C =
contiguous, Z = zigzag, D = per-document zigzag. Example: `G CP4 Z`.
Pre-norm = q/k L2-normalised before the call; other rows normalise in
the kernel. A cell is one workload x operation (fwd, fwd+bwd, ckpt = 2 x
fwd + bwd).
- Table A = this PR vs previous, with an A/A arm (a second copy of
previous). Table B = det(main), the strongest non-Cake arm and this PR.
A multi-node CP16 sub-table (B200 2 x 8, GB300 4 x 4 GPUs) is separate
from both.
- A single-rank proxy (`proxy <workload> r<rank>`) runs one rank's share
of a CP > 1 workload on one GPU at CP1. It excludes communication and is
not a full step. Its documents are listed in [this
comment](https://github.com/flashinfer-ai/flashinfer/pull/8521#issuecomment-6612373113).
- Unmarked cells: timed at an earlier commit of this PR with identical
launch forms; `[this PR]`: timed at this head; `[merge]`: timed after
the merge with main.
- Per group, c = this PR ms / previous ms and a = A/A ms / previous ms.
A run is win if every c is below 1 and below every a, REGRESSION if
every c is above 1 and above every a, and parity otherwise.
- A workload regresses if two or more of its published runs are
REGRESSION and its fwd+bwd and ckpt cells point the same way.
- Win counts = win cells / changed workloads with a published cell. A
changed workload is one whose launch forms differ from the previous
version. A win cell has every used run classed win. Used runs are all
published runs of the cell, including earlier runs not shown in the
table, except in the B200 F CP2 D fwd+bwd and ckpt cells, whose launch
forms the last commit changes; there only the runs at this head count.
Medians and ranges are over the used runs of the win cells.
**Baselines**
- Previous = the non-deterministic package
`flashinfer.experimental.cake_gdn_train_kcp` of FlashInfer main.
det(main) = the deterministic KCP package
`flashinfer.experimental.cake_gdn_det_train_kcp` of FlashInfer main.
Both come from main bf0605d519ecf585a73372037baedda4a159fb45 for rows
timed at the rules commit, and from main
358c1756342f0fbbb0bb4192bbd0643153d7513f for rows re-timed after the
merge. The deterministic package is code-identical between the two
(docstrings and README only). The non-deterministic package differs only
in code these rows do not reach.
- Non-Cake arms: FLA = flash-linear-attention / fla-core 0.6.0 from PyPI
(upstream tag v0.6.0 = 413ad2c655ece57b94bc9ddc51fdbd43cf09d984). It
runs as FLA-CP (its own context parallelism; layouts C and Z only) or
FLA-HS (head split: all-to-all around FLA's single-GPU kernels). The
cuDNN FROST GDN plan runs in the same head split. FROST-ck0 recomputes
the states in the backward; FROST-ck64 stores them every 64 tokens. The
FROST plan is CuTe DSL code (nvidia-cutlass-dsl 4.8.1) shipped in
nvidia-cudnn-frontend 1.32.1 from PyPI (upstream tag v1.32.1 =
73f69cf0d704bbc384da4ecd1358f36579bf1d80). The environment has
nvidia-cudnn-cu13 9.24.0.43.
- The GDN training implementations with forward and backward that run on
these GPUs in the same process and accept the same packed
variable-length inputs are FLA and cuDNN FROST; FlashQLA needs its own
environment and is covered below.
- Per row, each arm's outputs and input gradients are checked against
this PR on that row's inputs: finite, rel-RMS <= 1e-2, max-abs <= 2e-2 x
max|ref|. Table B shows the fastest arm that passes.
**Setup (all tables)**
B200 (sm_100a, 148 SMs), GB300 (sm_103a, 152 SMs) and R200 (sm_107a, 208
SMs); CUDA 13.4 development container (CUDA 13.4.2), PyTorch
2.14.1+cu130, NCCL 2.30.7. BF16 q/k/v, head dim 128, FP32 g and beta;
every rank draws the same seeded inputs. Every timed arm runs eager;
CUDA graph replay was not timed. The arms of a run share one process and
run palindromically in groups (two blocks per arm per group; block =
warm-up calls + timed calls). Paired runs reverse the arm order. L2 is
cold: a 256 MiB buffer fill, a device synchronize and a barrier run
before every iteration. Timing uses CUDA events on every rank, and an
iteration's time is the slowest rank's. The statistic is the median over
all timed iterations of an arm.
| groups x blocks x timed (warm-up): count | B200 | GB300 | R200 |
|---|---|---|---|
| table A runs | 5x2x8 (5): 409; 5x2x12 (10): 70; sustained: 78 | 5x2x12
(10): 487; 5x2x8 (5): 18; 7x2x24 (10): 11 | 5x2x12 (10): 371 |
| table B rows | 5x2x8 (5): 85; sustained: 2 | 5x2x12 (10): 79; 5x2x8
(5): 59 | 5x2x12 (10): 13 CP1 + 149 proxy rows |
| CP16 runs | 5x2x12 (10) | 5x2x12 (10) | none |
Sustained (B200 only) = 3 groups, block = 5 s idle + 20 warm-up + 300
timed. The other protocols have no idle time.
Not published: runs with a group drift above 2 % (table A also accepts a
run when >= 5 per-group paired ratios lie within 1 % of their median;
table B uses the largest drift over all arms of the run), and short
table A runs with >= 5 % of timed samples under the software power cap
on any compared arm (sustained runs show that share when it is >= 5 %).
A/A noise (a = A/A ms / previous ms per published table A run, p5-p95,
full workloads / proxies): B200 0.997-1.003 / 0.996-1.003, GB300
0.998-1.002 / 0.994-1.004, R200 0.997-1.008 / 0.997-1.003.
**B200 (sm_100a, 148 SMs)**
- **vs previous (table A):** fwd 14/65 win cells (median 1.086x, range
1.000-1.114x), fwd+bwd 49/64 (1.013x, 1.000-1.058x), ckpt 44/64 (1.013x,
1.001-1.057x). Single-rank proxies: fwd 0/8, fwd+bwd 2/8 (1.014x,
1.010-1.017x), ckpt 2/8 (1.013x, 1.010-1.015x). Largest: A CP1 C
pre-norm fwd 1.114x in both runs, H CP1 C pre-norm fwd 1.111x/1.109x.
- **Runs below 1 (table A):** 0.980-0.999x (proxies 0.996-0.998x). 7
cells have one REGRESSION run (lowest 0.996x); each also has a parity or
win run. No workload regresses. All other runs below 1 are parity.
- **vs det(main) (table B):** faster on all 87 published rows: fwd
1.105-1.271x, fwd+bwd 1.047-1.269x, ckpt 1.046-1.358x.
- **vs non-Cake (table B):** fwd 0.361-1.952x, fwd+bwd 0.555-1.584x,
ckpt 0.694-1.924x. Below 1 on all 13 CP1 (single GPU) rows: FROST-ck0 in
4 fwd rows, FROST-ck64 in 5 fwd+bwd and 4 ckpt rows. None of the 74 CP2
rows is below 1.
**GB300 (sm_103a, 152 SMs)**
- **vs previous (table A):** fwd 10/54 win cells (median 1.083x, range
1.013-1.116x), fwd+bwd 44/54 (1.015x, 1.003-1.062x), ckpt 44/54 (1.014x,
1.002-1.062x). Single-rank proxies: fwd 1/16 (1.031x, 1.030-1.032x),
fwd+bwd 6/17 (1.014x, 1.008-1.035x), ckpt 5/17 (1.012x, 1.008-1.031x).
Largest: A CP1 C pre-norm fwd 1.116x and H CP1 C pre-norm fwd 1.114x,
each in both runs.
- **Runs below 1 (table A):** 0.992-0.999x (proxies 0.988-0.999x), all
parity. No run is classed REGRESSION.
- **vs det(main) (table B):** faster on all 138 published rows: fwd
1.091-1.366x, fwd+bwd 1.023-1.500x, ckpt 1.025-1.274x.
- **vs non-Cake (table B):** fwd 0.335-1.933x, fwd+bwd 0.549-1.747x,
ckpt 0.672-2.074x. Below 1 on all 14 CP1 (single GPU) rows: FROST-ck0 in
4 fwd rows, FROST-ck64 in 6 fwd+bwd and 4 ckpt rows. None of the 124
CP2/CP4 rows is below 1.
**R200 (sm_107a, 208 SMs)**
- **vs previous (table A):** R200 runs used one GPU. Table A has 6 CP1
workloads and 56 single-rank proxies, so most of the evidence, and all
CP > 1 evidence, comes from proxies, which exclude communication. CP1:
2/3 win cells per op (M CP1 C and M CP1 C pre-norm): fwd median 1.099x
(range 1.090-1.109x), fwd+bwd 1.054x (1.052-1.061x), ckpt 1.051x
(1.049-1.054x). Proxies: 14/56 win cells per op: fwd 1.067x
(1.027-1.105x), fwd+bwd 1.038x (1.019-1.064x), ckpt 1.037x
(1.010-1.060x). Largest: M CP1 C pre-norm fwd 1.105x/1.109x, proxy G CP4
Z r3 pre-norm fwd 1.104x/1.105x.
- **Runs below 1 (table A):** 0.994-0.999x (proxies 0.995-0.999x). 6
cells (1 CP1, 5 proxies) have one REGRESSION run (lowest 0.995x); each
also has a parity or win run. Proxy E CP4 C r3 has one REGRESSION run in
fwd and one in fwd+bwd, but both of its ckpt runs are parity. No
workload regresses. All other runs below 1 are parity.
- **vs det(main) (table B, timed at this head):** faster on all 13
published CP1 rows: fwd 1.162-1.311x, fwd+bwd 1.048-1.115x, ckpt
1.047-1.156x. Faster on all 149 published single-rank proxy rows too:
fwd 1.171-1.354x, fwd+bwd 1.053-1.271x, ckpt 1.050-1.157x.
- **vs non-Cake (table B, timed at this head):** this PR is slower on
all 13 published CP1 rows; cuDNN FROST is faster on every one. fwd
0.321-0.403x (FROST-ck0), fwd+bwd 0.717-0.764x (FROST-ck64), ckpt
0.834-0.925x (FROST-ck0 on 2 rows, FROST-ck64 on 3). Single-rank proxies
(one GPU, no communication; rows in [this
comment](https://github.com/flashinfer-ai/flashinfer/pull/8521#issuecomment-6614291049)):
below 1 on 140 of 149 published rows: fwd 51 of 51 (0.316-0.797x,
FROST-ck0), fwd+bwd 49 of 49 (0.736-0.852x, FROST-ck64), ckpt 40 of 49
(0.865-1.069x; FROST-ck0 on 7 of the 40, FROST-ck64 on 33).
**cuDNN FROST on B200.** On the inputs of the earlier comparison (#8435,
B200 CP1 shape M), both FROST modes fail the accuracy check (dbeta = 0
at one element where an FP64 reference gives 0.986; max-abs error 0.1349
of max|ref|, limit 0.02). The failure is input-dependent: with two other
seeds of the same recipe and with this benchmark's inputs, FROST's
largest max-abs error vs FP64 is 0.0082-0.0094 of max|ref|, and FROST
passes the per-row check on every B200 row where table B shows it.
**FlashQLA 0.1.3 is not an arm:** it needs its own environment and
cannot share a process with the FROST arm. On CP1 shape M (B200 and
GB300) it was slower than both FROST modes and faster than FLA in all 12
measured cells. On those B200 inputs FROST fails the accuracy check, so
there FlashQLA could be the strongest passing arm; no table B row has
that case. Other shapes and R200 were not measured.
## Performance tables
Previous = the non-deterministic KCP package of FlashInfer main; this PR
= this PR's head; both run in one process with det(main) and an A/A copy
of previous. Unmarked cells: timed at an earlier commit of this PR with
identical launch forms; [this PR]: timed at this head; [merge]: timed
after the merge with main.
Table A cells: previous ms -> this PR ms of run 1 (absolutes are run 1),
then the ratios previous / this PR of run 1 / run 2 and the class per
run (A/A envelope of per-group ratios: win, parity = inside the A/A
spread, REGRESSION). Table B: ratio = baseline ms / this PR ms from the
absolutes in the row. Shown: the latest step's published runs plus any
earlier REGRESSION run at the shipped launch forms (the B200 F CP2 D
backward runs before the last commit measured replaced launch forms and
are not shown); no workload regresses (a workload regresses if two or
more of its published runs are REGRESSION and its fwd+bwd and ckpt cells
point the same way).
Legend: '*' = ratio below 1.000 (this PR slower than that baseline);
'sust.' = sustained protocol, else short; 'cap x %' = timed samples
under the software power cap (>= 5 %); 'n/p cap' / 'n/p unstable' = not
published (power-capped short run / drift above 2 %), listed at the end.
Shapes (head dim 128, BF16 q / k / v; documents in tokens; heads = (q/k
heads, v heads)); layouts C = contiguous, Z = zigzag, D = per-document
zigzag; CPn = context parallel over n GPUs; pre-norm = q and k
L2-normalised before the call (other rows: in the kernel):
|shape|documents|tokens|heads|
|---|---|---|---|
|A|4: 98304/65536/57344/40960|262144|(16, 48)|
|B|3: 22672/131070/108245|261987|(16, 48)|
|C|2: 128147/131027|259174|(16, 48)|
|D|1: 254976|254976|(16, 48)|
|E|4: 28874/38814/60180/121988|249856|(16, 48)|
|F|6: 126565/32714/33696/15081/22979/10782|241817|(16, 48)|
|G|12:
3598/4930/6262/7593/8925/10256/11588/12920/14251/15583/16914/18252|131072|(16,
48)|
|H|4: 49152/32768/28672/20480|131072|(16, 48)|
|I|4: 24576/16384/14336/10240|65536|(16, 48)|
|J|3: 15945/21845/27746|65536|(16, 48)|
|K|4: 25583/5103/5103/5103|40892|(16, 48)|
|L|4: 20463/10223/5103/2543|38332|(16, 48)|
|M|5: 12000/9000/6000/3768/2000|32768|(16, 48)|
|N|3: 2049/3001/3142|8192|(2, 6)|
Table B arm = the fastest non-Cake arm that passes the accuracy check:
FLA-CP = flash-linear-attention's own context parallelism; FLA-HS = head
split (all-to-all) around FLA's single-GPU kernels; FROST-ck0 / -ck64 =
the same head split around cuDNN frontend's GDN (states recomputed /
stored every 64 tokens).
## B200 (sm_100a, 148 SMs)
### A. This PR vs previous version (same process)
|workload|fwd: previous ms -> this PR ms (ratio)|fwd+bwd|ckpt (2 x fwd +
bwd)|
|---|---|---|---|
|A CP1 C|11.075 -> 10.099 (1.097/1.097 win)|37.818 -> 36.785
(1.028/1.052 win)|38.927 -> 37.041 (1.051 win)|
|A CP1 C pre-norm|9.989 -> 8.966 (1.114/1.114 win)|n/p cap|n/p cap|
|A CP2 C|7.338 -> 7.339 (1.000/0.999* parity)|sust. 23.354 -> 23.358
(1.000 parity, cap 23.6 %)|sust. 24.851 -> 24.911 (0.998* parity, cap
23.6 %)|
|A CP2 D|7.840 -> 7.836 (1.001/0.999* parity)|sust. 23.932 -> 23.657
(1.012 win, cap 9.4 %)|sust. 26.007 -> 25.707 (1.012 win, cap 5.6 %)|
|A CP2 Z|7.838 -> 7.837 (1.000/0.998* parity)|sust. 24.400 -> 24.047
(1.015 win)|sust. 25.756 -> 25.406 (1.014 win)|
|A CP4 C|3.933 -> 3.917 (1.004/0.999* parity)|sust. 11.840 -> 11.843
(1.000 parity, cap 39.2 %)|sust. 12.700 -> 12.697 (1.000 parity, cap
26.9 %)|
|A CP4 D|4.043 -> 4.037 (1.001/0.998* parity)|sust. 11.724 -> 11.554
(1.015 win, cap 10.5 %)|sust. 13.251 -> 13.069 (1.014 win, cap 6.8 %)|
|A CP4 Z|4.054 -> 4.045 (1.002/1.000 parity)|sust. 12.143 -> 12.101
(1.003 win, cap 20.2 %)|sust. 13.107 -> 13.032 (1.006 win, cap 13.0 %)|
|A CP8 C|2.265 -> 2.264 (1.000/0.999* parity)|sust. 6.423 -> 6.422
(1.000 parity)|7.051 -> 7.049 (1.000/1.001 parity)|
|A CP8 D|2.309 -> 2.305 (1.002/0.999* parity)|sust. 6.433 -> 6.326
(1.017 win)|7.427 -> 7.317 (1.015/1.013 win)|
|A CP8 Z|2.538 -> 2.536 (1.001/0.999* parity)|6.660 -> 6.655
(1.001/1.000 parity)|7.669 -> 7.664 (1.001/1.000 parity)|
|B CP2 C|7.616 -> 7.620 (0.999*/1.000 parity)|sust. 23.822 -> 23.850
(0.999* parity, cap 21.2 %)|sust. 25.402 -> 25.387 (1.001 parity, cap
19.9 %); earlier sust. 25.354 -> 25.407 (0.998* REGRESSION, cap 22.4 %)|
|B CP2 D|7.694 -> 7.695 (1.000/0.999* parity)|sust. 23.822 -> 23.690
(1.006 win, cap 9.6 %)|sust. 25.718 -> 25.517 (1.008 win, cap 8.9 %)|
|B CP2 Z|7.542 -> 7.540 (1.000/0.999* parity)|sust. 23.343 -> 23.252
(1.004 win, cap 11.2 %)|sust. 24.982 -> 24.871 (1.004 win, cap 8.4 %)|
|B CP4 C|3.958 -> 3.955 (1.001 parity)|sust. 11.911 -> 11.916 (1.000
parity, cap 37.1 %)|sust. 12.808 -> 12.810 (1.000 parity, cap 24.2 %)|
|B CP4 D|4.469 -> 4.469 (1.000/0.999* parity)|sust. 13.003 -> 12.877
(1.010 win, cap 5.0 %)|sust. 14.210 -> 14.076 (1.010 win)|
|B CP4 Z|4.136 -> 4.132 (1.001/0.998* parity)|sust. 12.427 -> 12.421
(1.000 win)|sust. 13.348 -> 13.344 (1.000 parity)|
|C CP2 C|7.955 -> 7.962 (0.999*/0.999* parity)|sust. 23.910 -> 23.937
(0.999* parity, cap 21.6 %)|sust. 25.623 -> 25.551 (1.003 win, cap 20.4
%)|
|C CP2 D|sust. 7.372 -> 7.369 (1.000 win)|sust. 22.784 -> 22.710 (1.003
win, cap 16.8 %)|sust. 24.631 -> 24.445 (1.008 win, cap 13.7 %)|
|C CP2 Z|sust. 7.376 -> 7.373 (1.000 win)|sust. 22.925 -> 22.800 (1.005
win, cap 15.6 %)|sust. 24.739 -> 24.568 (1.007 win, cap 13.7 %)|
|C CP4 C|3.954 -> 3.950 (1.001/0.999* parity)|sust. 11.906 -> 11.755
(1.013 win)|sust. 12.937 -> 12.787 (1.012 win)|
|C CP4 D|4.545 -> 4.538 (1.002 parity)|sust. 13.139 -> 13.100 (1.003
win)|14.277 -> 14.229 (1.003/0.998* parity)|
|C CP4 Z|3.947 -> 3.944 (1.001/0.999* parity)|sust. 11.866 -> 11.703
(1.014 win)|sust. 12.910 -> 12.748 (1.013 win)|
|D CP2 C|7.803 -> 7.811 (0.999*/1.001 parity)|sust. 23.466 -> 23.480
(0.999* parity, cap 17.8 %); earlier sust. 23.435 -> 23.467 (0.999*
REGRESSION, cap 21.5 %)|sust. 25.118 -> 25.028 (1.004 win, cap 16.2 %)|
|D CP2 D|sust. 7.268 -> 7.272 (0.999* parity)|sust. 22.532 -> 22.406
(1.006 win, cap 14.5 %)|sust. 24.354 -> 24.172 (1.008 win, cap 11.7 %)|
|D CP2 Z|sust. 7.232 -> 7.232 (1.000 parity)|sust. 22.411 -> 22.352
(1.003 win, cap 15.3 %)|sust. 24.213 -> 24.042 (1.007 win, cap 12.4 %)|
|E CP2 C|7.487 -> 7.485 (1.000/1.000 parity)|[this PR] sust. 22.916 ->
22.661 (1.011 win)|[this PR] sust. 24.426 -> 24.174 (1.010 win)|
|E CP2 D|7.401 -> 7.398 (1.000/0.999* parity)|sust. 22.725 -> 22.453
(1.012 win, cap 9.6 %)|sust. 24.601 -> 24.288 (1.013 win)|
|E CP2 Z|7.517 -> 7.513 (1.001/0.999* parity/REGRESSION)|sust. 23.181 ->
22.888 (1.013 win)|sust. 24.594 -> 24.301 (1.012 win)|
|E CP4 C|3.824 -> 3.830 (0.998*/1.000 parity)|sust. 11.543 -> 11.501
(1.004 win, cap 29.0 %)|sust. 12.536 -> 12.447 (1.007 win, cap 13.3 %)|
|E CP4 D|3.928 -> 3.928 (1.000/0.998* parity)|sust. 11.417 -> 11.286
(1.012 win)|sust. 12.927 -> 12.799 (1.010 win)|
|E CP4 Z|3.869 -> 3.872 (0.999*/0.998* parity)|sust. 11.604 -> 11.441
(1.014 win)|sust. 12.604 -> 12.439 (1.013 win)|
|E CP8 C|2.235 -> 2.235 (1.000/0.999* parity)|sust. 6.287 -> 6.284
(1.000 win)|6.901 -> 6.906 (0.999*/0.999* parity)|
|E CP8 D|2.277 -> 2.278 (1.000 parity)|6.338 -> 6.279 (1.009 win)|7.332
-> 7.266 (1.009/1.008 win)|
|E CP8 Z|2.448 -> 2.445 (1.001/0.998* parity)|6.384 -> 6.381
(1.000/0.999* parity)|7.377 -> 7.374 (1.000/0.999* win/parity)|
|F CP2 C|7.429 -> 7.439 (0.999*/1.000 parity)|sust. 22.083 -> 21.973
(1.005 win, cap 20.4 %)|sust. 23.749 -> 23.675 (1.003 win, cap 18.7 %)|
|F CP2 D|7.018 -> 7.013 (1.001/0.998* parity)|[this PR] sust. 20.958 ->
20.954 (1.000/1.000 parity)|[this PR] sust. 23.438 -> 23.436
(1.000/1.000 parity)|
|F CP2 Z|6.957 -> 6.949 (1.001/0.997* parity)|sust. 21.712 -> 21.495
(1.010 win, cap 10.2 %)|sust. 23.256 -> 23.043 (1.009 win, cap 9.8 %)|
|F CP4 C|4.051 -> 4.049 (1.000/0.998* parity)|sust. 11.838 -> 11.836
(1.000 win)|12.851 -> 12.845 (1.000/1.001 parity)|
|F CP4 D|3.883 -> 3.884 (1.000/1.000 parity)|sust. 11.243 -> 11.238
(1.000 win)|sust. 12.789 -> 12.782 (1.001 win)|
|F CP4 Z|3.832 -> 3.835 (0.999*/0.998* parity)|sust. 11.478 -> 11.418
(1.005 win)|sust. 12.358 -> 12.295 (1.005 win)|
|F CP8 C|2.251 -> 2.252 (1.000/1.000 parity)|6.122 -> 6.116
(1.001/0.999* parity)|6.743 -> 6.743 (1.000/1.000 parity)|
|F CP8 D|2.329 -> 2.331 (0.999*/0.999* parity)|6.349 -> 6.337
(1.002/1.002 win/parity)|7.435 -> 7.418 (1.002/1.001 win/parity)|
|F CP8 Z|2.369 -> 2.373 (0.998*/0.998* parity)|6.235 -> 6.236
(1.000/0.999* parity)|7.207 -> 7.204 (1.000/0.999* parity)|
|G CP1 C|4.855 -> 4.865 (0.998*/1.000 parity)|18.107 -> 17.890
(1.012/1.011 win)|18.244 -> 18.034 (1.012/1.012 win)|
|G CP2 C|3.553 -> 3.550 (1.001/1.000 parity)|10.673 -> 10.681
(0.999*/1.000 parity)|11.397 -> 11.386 (1.001/1.000 parity)|
|G CP2 D|4.190 -> 4.192 (1.000/0.998* parity/REGRESSION)|12.005 ->
11.870 (1.011/1.009 win)|13.805 -> 13.650 (1.011/1.009 win)|
|G CP2 Z|3.164 -> 3.160 (1.001/1.000 parity)|10.404 -> 10.322
(1.008/1.010 win)|11.043 -> 10.962 (1.007/1.008 win)|
|G CP4 C|2.339 -> 2.301 (1.017/1.011 win)|6.341 -> 6.342 (1.000
parity)|7.020 -> 7.012 (1.001/0.999* parity)|
|G CP4 D|2.699 -> 2.695 (1.001/1.000 parity)|7.092 -> 7.007 (1.012/1.010
win)|8.479 -> 8.392 (1.010/1.008 win)|
|G CP4 Z|2.213 -> 2.045 (1.082/1.079 win)|6.332 -> 5.990 (1.057/1.054
win)|6.838 -> 6.590 (1.038/1.031 win)|
|G CP8 C|1.745 -> 1.744 (1.001/0.999* parity)|4.155 -> 4.155
(1.000/0.999* parity)|4.971 -> 4.978 (0.999*/1.000 parity)|
|G CP8 D|1.901 -> 1.898 (1.002/0.997* parity)|4.589 -> 4.544
(1.010/1.007 win/parity)|5.766 -> 5.716 (1.009/1.006 win/parity)|
|G CP8 Z|1.442 -> 1.413 (1.021/1.013 win)|3.629 -> 3.535 (1.027/1.022
win)|4.282 -> 4.220 (1.015/1.010 parity)|
|H CP1 C|5.642 -> 5.158 (1.094/1.095 win)|19.062 -> 18.503 (1.030/1.055
win)|19.208 -> 18.638 (1.031/1.053 win)|
|H CP1 C pre-norm|5.080 -> 4.574 (1.111/1.109 win)|17.801 -> 16.858
(1.056/1.056 win)|17.961 -> 17.003 (1.056/1.055 win)|
|H CP2 D|3.991 -> 3.993 (0.999*/0.998* parity)|11.646 -> 11.462 (1.016
win)|13.129 -> 12.941 (1.015/1.015 win)|
|H CP4 C|2.583 -> 2.586 (0.999*/0.999* parity)|6.638 -> 6.639
(1.000/1.001 parity)|7.702 -> 7.694 (1.001/1.000 parity)|
|H CP4 D|2.283 -> 2.276 (1.003/0.999* parity)|sust. 6.369 -> 6.272
(1.015 win, cap 7.3 %)|7.344 -> 7.229 (1.016/1.013 win)|
|H CP4 Z|2.243 -> 2.237 (1.003/0.999* parity)|6.307 -> 6.301
(1.001/1.000 win/parity)|7.098 -> 7.086 (1.002/1.000 parity)|
|H CP8 C|1.494 -> 1.491 (1.002/0.998* parity)|3.764 -> 3.760
(1.001/1.000 parity)|4.393 -> 4.388 (1.001/0.998* parity)|
|H CP8 D|1.444 -> 1.440 (1.003/0.997* parity)|3.692 -> 3.635
(1.016/1.013 win)|4.399 -> 4.349 (1.011/1.009 parity)|
|H CP8 Z|1.444 -> 1.442 (1.001/1.000 parity)|3.606 -> 3.601
(1.001/0.999* parity)|4.240 -> 4.226 (1.003/0.996* parity)|
|I CP1 C|2.924 -> 2.683 (1.090/1.091 win)|9.743 -> 9.450 (1.031/1.054
win)|9.880 -> 9.592 (1.030/1.052 win)|
|I CP1 C pre-norm|2.625 -> 2.376 (1.105/1.103 win)|9.070 -> 8.575
(1.058/1.058 win)|9.219 -> 8.718 (1.057/1.056 win)|
|I CP2 D|2.238 -> 2.234 (1.002/0.999* win/parity)|6.292 -> 6.184
(1.017/1.017 win)|7.217 -> 7.109 (1.015/1.014 win)|
|I CP4 C|1.477 -> 1.474 (1.002/0.999* parity)|3.577 -> 3.572
(1.001/1.000 parity)|4.199 -> 4.192 (1.002/0.999* parity)|
|I CP4 D|1.411 -> 1.408 (1.002/0.994* parity)|3.626 -> 3.573
(1.015/1.012 win)|4.303 -> 4.267 (1.008/1.010 parity/win)|
|I CP4 Z|1.293 -> 1.290 (1.002/0.998* parity)|3.399 -> 3.392
(1.002/0.999* parity)|4.008 -> 4.004 (1.001/0.999* parity)|
|I CP8 C|0.958 -> 0.956 (1.002/0.999* parity)|2.571 -> 2.555
(1.006/0.997* parity)|3.427 -> 3.437 (0.997*/0.996* parity)|
|I CP8 D|1.015 -> 1.016 (0.999*/0.994* parity)|2.598 -> 2.583
(1.006/1.004 parity)|3.407 -> 3.397 (1.003/1.003 parity/win)|
|I CP8 Z|[this PR] 0.912 -> 0.909 (1.003/0.998* parity)|[this PR] 2.498
-> 2.517 (0.992*/0.993* parity)|[this PR] 3.372 -> 3.376 (0.999*/0.999*
parity); earlier 3.340 -> 3.354 (0.996* REGRESSION)|
|J CP1 C|2.714 -> 2.726 (0.996*/1.000 parity)|9.514 -> 9.522
(0.999*/0.999* parity)|9.658 -> 9.663 (0.999*/1.000 parity)|
|K CP1 C|2.695 -> 2.696 (1.000/1.000 parity)|8.042 -> 7.968 (1.009/1.009
win)|8.675 -> 8.612 (1.007/1.008 win)|
|L CP1 C|2.484 -> 2.483 (1.000/0.999* parity)|7.370 -> 7.345
(1.003/1.003 win/parity)|7.939 -> 7.915 (1.003/1.003 win)|
|M CP1 C (control)|1.518 -> 1.416 (1.072/1.076 win)|[this PR] 5.137 ->
4.921 (1.044/1.043 win)|[this PR] 5.268 -> 5.044 (1.044/1.046 win)|
|M CP1 C pre-norm|1.354 -> 1.246 (1.087/1.085 win)|4.599 -> 4.383
(1.049/1.040 win)|4.767 -> 4.547 (1.048/1.042 win)|
|M CP2 C|1.284 -> 1.284 (1.000/1.002 parity)|3.292 -> 3.283 (1.003/1.003
win)|3.796 -> 3.741 (1.015/1.014 win)|
|M CP2 D|1.407 -> 1.402 (1.004/0.995* parity)|3.614 -> 3.583
(1.009/1.006 win)|4.302 -> 4.265 (1.009/1.004 win)|
|M CP2 Z|1.289 -> 1.242 (1.038/1.034 win)|3.365 -> 3.293 (1.022/1.018
win)|3.935 -> 3.908 (1.007/1.008 win/parity)|
|M CP4 C|0.924 -> 0.924 (1.000/0.999* parity)|2.495 -> 2.498
(0.999*/0.997* parity)|3.387 -> 3.381 (1.002/0.997* parity/REGRESSION)|
|M CP4 D|1.049 -> 1.043 (1.006/0.999* parity)|2.620 -> 2.608
(1.005/1.002 parity)|3.418 -> 3.414 (1.001/0.999* parity)|
|M CP4 Z|0.902 -> 0.896 (1.007/0.996* parity)|2.434 -> 2.421
(1.005/0.980* parity)|3.310 -> 3.278 (1.010/0.995* parity)|
|M CP8 C|0.759 -> 0.761 (0.997*/0.996* parity)|2.327 -> 2.331
(0.998*/0.999* parity)|3.429 -> 3.426 (1.001 parity)|
|M CP8 D|0.886 -> 0.885 (1.001/0.999* parity)|2.253 -> 2.252 (1.000
parity)|3.226 -> 3.238 (0.996*/0.996* parity)|
|M CP8 Z|0.688 -> 0.690 (0.997* parity)|2.202 -> 2.219 (0.992*/0.996*
parity)|3.219 -> 3.232 (0.996*/0.998* parity)|
|N CP1 C|0.307 -> 0.305 (1.007/0.997* parity)|1.466 -> 1.468
(0.999*/0.995* parity)|2.144 -> 2.148 (0.998*/0.996* parity)|
|N CP2 C|0.602 -> 0.602 (1.000 parity)|1.979 -> 1.987 (0.996*/0.998*
parity)|2.908 -> 2.908 (1.000/0.994* parity)|
|N CP2 D|0.659 -> 0.657 (1.003/0.995* parity)|2.121 -> 2.126
(0.998*/0.999* parity)|3.135 -> 3.134 (1.000/0.999* parity)|
|N CP2 Z|0.678 -> 0.667 (1.016/1.000 parity)|2.134 -> 2.140
(0.997*/0.994* parity)|3.150 -> 3.149 (1.000/0.993* parity)|
|N CP4 C|0.614 -> 0.607 (1.012/1.000 parity)|2.093 -> 2.100
(0.997*/0.987* parity)|3.040 -> 3.033 (1.002/0.997* parity)|
|N CP4 D|0.660 -> 0.656 (1.006 parity)|2.138 -> 2.158 (0.991*/0.995*
parity)|3.182 -> 3.165 (1.005/0.992* parity)|
|N CP4 Z|0.666 -> 0.662 (1.006/0.992* parity)|2.151 -> 2.143
(1.004/0.994* parity)|3.176 -> 3.153 (1.007/0.994* parity)|
|N CP8 C|0.625 -> 0.628 (0.995*/0.990* parity)|2.184 -> 2.195
(0.995*/0.998* parity)|3.142 -> 3.153 (0.997*/0.998* parity)|
|N CP8 D|0.700 -> 0.698 (1.003/0.993* parity)|2.252 -> 2.257
(0.998*/0.994* parity)|3.283 -> 3.292 (0.997*/0.998* parity)|
|N CP8 Z|0.700 -> 0.702 (0.997*/0.988* parity)|2.256 -> 2.265
(0.996*/0.994* REGRESSION/parity)|[this PR] 3.244 -> 3.265
(0.994*/0.995* parity)|
Single-rank proxies (single-rank proxy, not a full step): one GPU runs
the documents of rank r of the named CP configuration (documents in
[this
comment](https://github.com/flashinfer-ai/flashinfer/pull/8521#issuecomment-6612373113)).
|proxy of|fwd|fwd+bwd|ckpt|
|---|---|---|---|
|A CP16 D r0|0.782 -> 0.780 (1.003/1.001 parity)|2.619 -> 2.576
(1.017/1.014 win)|2.900 -> 2.857 (1.015/1.014 win)|
|E CP16 D r0|0.768 -> 0.765 (1.004/1.000 parity)|2.574 -> 2.542
(1.013/1.010 win)|2.853 -> 2.819 (1.012/1.010 win)|
|F CP16 D r0|0.748 -> 0.747 (1.001/1.003 parity)|2.527 -> 2.520
(1.003/1.000 parity)|2.807 -> 2.802 (1.002/1.001 parity)|
|G CP16 D r1|0.508 -> 0.508 (1.000/1.000 parity)|1.860 -> 1.846
(1.008/1.008 parity)|2.298 -> 2.286 (1.005/1.003 parity)|
|G CP16 Z r2|0.540 -> 0.539 (1.002/1.000 parity)|1.822 -> 1.817
(1.003/1.002 parity)|2.245 -> 2.240 (1.002/1.002 parity)|
|H CP16 D r0|0.496 -> 0.495 (1.002/0.996* parity)|1.792 -> 1.772
(1.011/1.010 win/parity)|2.250 -> 2.228 (1.010/1.008 win/parity)|
|I CP16 D r0|0.352 -> 0.351 (1.003 parity)|1.452 -> 1.455 (0.998*
parity)|2.122 -> 2.113 (1.004/1.000 parity)|
|M CP16 D r0|0.302 -> 0.301 (1.003/0.997* parity)|1.445 -> 1.439 (1.004
parity)|2.113 -> 2.094 (1.009/1.000 parity)|
### B. Two baselines (same process: det(main), strongest non-Cake, this
PR)
|workload|op|det(main) ms|strongest non-Cake ms (arm)|this PR
ms|det(main) / this PR|non-Cake / this PR|
|---|---|---|---|---|---|---|
|A CP1 C|fwd+bwd|40.660|20.715 (FROST-ck64)|37.341|1.089|0.555*|
|A CP1 C|ckpt|40.872|26.149 (FROST-ck64)|37.694|1.084|0.694*|
|A CP2 C|fwd|8.457|12.198 (FROST-ck0)|7.331|1.154|1.664|
|A CP2 C|fwd+bwd|24.793|31.765 (FROST-ck64)|23.317|1.063|1.362|
|A CP2 C|ckpt|26.315|40.848 (FROST-ck64)|24.921|1.056|1.639|
|A CP2 D|fwd|8.952|12.207 (FROST-ck0)|7.861|1.139|1.553|
|A CP2 D|fwd+bwd|25.595|31.769 (FROST-ck64)|24.116|1.061|1.317|
|A CP2 Z|fwd|8.831|12.206 (FROST-ck0)|7.899|1.118|1.545|
|A CP2 Z|fwd+bwd|25.685|31.760 (FROST-ck64)|24.398|1.053|1.302|
|A CP2 Z|ckpt|27.068|40.850 (FROST-ck64)|25.749|1.051|1.586|
|B CP2 C|fwd|8.698|12.111 (FROST-ck0)|7.621|1.141|1.589|
|B CP2 C|fwd+bwd|25.246|31.710 (FROST-ck64)|23.826|1.060|1.331|
|B CP2 C|ckpt|26.727|40.764 (FROST-ck64)|25.423|1.051|1.603|
|B CP2 D|fwd|8.836|12.116 (FROST-ck0)|7.693|1.149|1.575|
|B CP2 D|fwd+bwd|25.524|31.701 (FROST-ck64)|24.050|1.061|1.318|
|B CP2 D|ckpt|27.246|40.756 (FROST-ck64)|25.956|1.050|1.570|
|B CP2 Z|fwd|8.640|12.122 (FROST-ck0)|7.539|1.146|1.608|
|B CP2 Z|fwd+bwd|25.160|31.708 (FROST-ck64)|23.800|1.057|1.332|
|B CP2 Z|ckpt|26.712|40.761 (FROST-ck64)|25.448|1.050|1.602|
|C CP2 C|fwd|8.921|11.038 (FLA-CP)|7.954|1.122|1.388|
|C CP2 C|fwd+bwd|24.994|34.529 (FROST-ck64)|23.830|1.049|1.449|
|C CP2 C|ckpt|26.671|44.880 (FROST-ck64)|25.502|1.046|1.760|
|C CP2 D|fwd|8.489|13.346 (FROST-ck0)|7.383|1.150|1.808|
|C CP2 D|fwd+bwd|24.737|34.521 (FROST-ck64)|23.276|1.063|1.483|
|C CP2 D|ckpt|26.405|44.864 (FROST-ck64)|24.990|1.057|1.795|
|C CP2 Z|fwd|8.522|13.347 (FROST-ck0)|7.396|1.152|1.805|
|C CP2 Z|fwd+bwd|24.913|34.525 (FROST-ck64)|23.428|1.063|1.474|
|C CP2 Z|ckpt|26.605|44.872 (FROST-ck64)|25.154|1.058|1.784|
|D CP2 C|fwd|8.758|13.062 (FROST-ck0)|7.843|1.117|1.665|
|D CP2 C|fwd+bwd|24.673|33.779 (FROST-ck64)|23.506|1.050|1.437|
|D CP2 C|ckpt|26.182|43.908 (FROST-ck64)|25.012|1.047|1.755|
|D CP2 D|fwd|8.414|13.064 (FROST-ck0)|7.286|1.155|1.793|
|D CP2 D|fwd+bwd|24.469|33.769 (FROST-ck64)|23.002|1.064|1.468|
|D CP2 D|ckpt|26.207|43.895 (FROST-ck64)|24.649|1.063|1.781|
|D CP2 Z|fwd|8.368|13.064 (FROST-ck0)|7.238|1.156|1.805|
|D CP2 Z|fwd+bwd|24.350|33.790 (FROST-ck64)|22.958|1.061|1.472|
|D CP2 Z|ckpt|25.960|43.895 (FROST-ck64)|24.542|1.058|1.789|
|E CP2 C|fwd|8.429|11.731 (FROST-ck0)|7.499|1.124|1.564|
|E CP2 C|fwd+bwd|23.827|30.874 (FROST-ck64)|22.635|1.053|1.364|
|E CP2 C|ckpt|25.336|39.603 (FROST-ck64)|24.159|1.049|1.639|
|E CP2 D|fwd|8.506|11.741 (FROST-ck0)|7.453|1.141|1.575|
|E CP2 D|fwd+bwd|23.809|30.857 (FROST-ck64)|22.457|1.060|1.374|
|E CP2 D|ckpt|25.675|39.596 (FROST-ck64)|24.352|1.054|1.626|
|E CP2 Z|fwd|8.474|11.738 (FROST-ck0)|7.557|1.121|1.553|
|E CP2 Z|fwd+bwd|24.091|30.867 (FROST-ck64)|22.874|1.053|1.349|
|E CP2 Z|ckpt|25.501|39.606 (FROST-ck64)|24.303|1.049|1.630|
|F CP2 C|fwd|8.332|11.258 (FROST-ck0)|7.436|1.120|1.514|
|F CP2 C|fwd+bwd|23.127|28.713 (FROST-ck64)|21.930|1.055|1.309|
|F CP2 C|ckpt|24.740|37.145 (FROST-ck64)|23.596|1.048|1.574|
|F CP2 D|fwd|8.146|11.260 (FROST-ck0)|7.036|1.158|1.600|
|F CP2 D|fwd+bwd|22.079|28.990 (FROST-ck64)|[this PR]
20.978|1.052|1.382|
|F CP2 D|ckpt|24.573|37.524 (FROST-ck64)|[this PR] 23.474|1.047|1.599|
|F CP2 Z|fwd|8.027|11.259 (FROST-ck0)|6.967|1.152|1.616|
|F CP2 Z|fwd+bwd|23.334|28.707 (FROST-ck64)|21.913|1.065|1.310|
|F CP2 Z|ckpt|24.777|37.149 (FROST-ck64)|23.451|1.057|1.584|
|G CP1 C|fwd|5.956|1.777 (FROST-ck0)|4.884|1.219|0.364*|
|G CP1 C|fwd+bwd|19.417|10.420 (FROST-ck64)|17.961|1.081|0.580*|
|G CP2 C|fwd|4.031|6.171 (FROST-ck0)|3.565|1.131|1.731|
|G CP2 C|fwd+bwd|11.160|16.329 (FROST-ck64)|10.660|1.047|1.532|
|G CP2 C|ckpt|11.916|21.023 (FROST-ck64)|11.379|1.047|1.848|
|G CP2 D|fwd|4.819|6.170 (FROST-ck0)|4.235|1.138|1.457|
|G CP2 D|fwd+bwd|12.607|16.331 (FROST-ck64)|11.879|1.061|1.375|
|G CP2 D|ckpt|14.402|21.030 (FROST-ck64)|13.673|1.053|1.538|
|G CP2 Z|fwd|3.716|6.160 (FROST-ck0)|3.156|1.177|1.952|
|G CP2 Z|fwd+bwd|10.989|16.331 (FROST-ck64)|10.307|1.066|1.584|
|G CP2 Z|ckpt|11.565|21.026 (FROST-ck64)|10.930|1.058|1.924|
|H CP1 C|fwd|6.565|1.863 (FROST-ck0)|5.165|1.271|0.361*|
|H CP1 C|fwd+bwd|20.393|10.451 (FROST-ck64)|18.546|1.100|0.564*|
|I CP1 C|fwd|3.397|0.997 (FROST-ck0)|2.692|1.262|0.370*|
|I CP1 C|fwd+bwd|10.422|5.592 (FROST-ck64)|9.479|1.099|0.590*|
|I CP1 C|ckpt|10.578|6.855 (FROST-ck64)|9.632|1.098|0.712*|
|J CP1 C|fwd|3.202|1.115 (FROST-ck0)|2.734|1.171|0.408*|
|J CP1 C|fwd+bwd|9.977|5.951 (FROST-ck64)|9.524|1.048|0.625*|
|J CP1 C|ckpt|10.124|7.357 (FROST-ck64)|9.660|1.048|0.762*|
|M CP1 C|ckpt|5.553|4.221 (FROST-ck64)|5.071|1.095|0.832*|
|M CP2 C|fwd|1.599|1.889 (FROST-ck0)|1.292|1.238|1.462|
|M CP2 C|fwd+bwd|4.179|4.782 (FROST-ck64)|3.294|1.269|1.452|
|M CP2 C|ckpt|5.136|6.329 (FROST-ck64)|3.782|1.358|1.673|
|M CP2 D|fwd|1.572|1.890 (FROST-ck0)|1.422|1.105|1.329|
|M CP2 D|fwd+bwd|3.961|4.787 (FROST-ck64)|3.597|1.101|1.331|
|M CP2 D|ckpt|5.000|6.323 (FROST-ck64)|4.293|1.165|1.473|
|M CP2 Z|fwd|1.463|1.887 (FROST-ck0)|1.244|1.176|1.517|
|M CP2 Z|fwd+bwd|3.627|4.788 (FROST-ck64)|3.288|1.103|1.456|
|M CP2 Z|ckpt|4.505|6.327 (FROST-ck64)|3.935|1.145|1.608|
|N CP2 C|fwd+bwd|2.520|2.262 (FROST-ck0)|2.068|1.219|1.094|
|N CP2 C|ckpt|3.526|3.442 (FROST-ck0)|2.948|1.196|1.168|
|N CP2 Z|fwd+bwd|2.718|2.237 (FROST-ck64)|2.175|1.250|1.029|
Published rows: 87; det(main) faster in 0, non-Cake faster in 13; 15
unstable rows not published (listed at the end).
## GB300 (sm_103a, 152 SMs)
### A. This PR vs previous version (same process)
|workload|fwd: previous ms -> this PR ms (ratio)|fwd+bwd|ckpt (2 x fwd +
bwd)|
|---|---|---|---|
|A CP1 C|10.375 -> 9.404 (1.103/1.016 win/parity)|34.915 -> 34.059
(1.025/1.031 win)|35.063 -> 34.220 (1.025/1.031 win)|
|A CP1 C pre-norm|9.275 -> 8.308 (1.116/1.116 win)|31.560 -> 30.729
(1.027/1.062 win)|31.695 -> 30.860 (1.027/1.062 win)|
|A CP2 C|6.990 -> 6.993 (1.000/0.999* parity)|21.576 -> 21.366
(1.010/1.010 win)|23.254 -> 23.039 (1.009/1.009 win)|
|A CP2 D|7.241 -> 7.243 (1.000/1.000 parity)|22.148 -> 21.777
(1.017/1.016 win)|24.029 -> 23.674 (1.015/1.015 win)|
|A CP2 Z|7.587 -> 7.587 (1.000/1.001 parity)|23.376 -> 23.056
(1.014/1.014 win)|25.086 -> 24.814 (1.011/1.011 win)|
|A CP4 C|3.678 -> 3.673 (1.001/1.000 parity)|10.864 -> 10.864
(1.000/1.000 parity)|11.767 -> 11.771 (1.000/1.000 parity)|
|A CP4 D|3.771 -> 3.774 (0.999*/0.999* parity)|10.926 -> 10.698
(1.021/1.021 win)|12.345 -> 12.118 (1.019/1.018 win)|
|A CP4 Z|3.783 -> 3.790 (0.998*/0.998* parity)|11.310 -> 11.127
(1.016/1.016 win)|12.229 -> 12.041 (1.016/1.016 win)|
|B CP2 C|6.976 -> 6.969 (1.001/1.000 parity)|21.721 -> 21.447
(1.013/1.012 win)|23.401 -> 23.126 (1.012/1.012 win)|
|B CP2 D|7.114 -> 7.115 (1.000/0.999* parity)|22.096 -> 21.759
(1.015/1.015 win)|23.848 -> 23.514 (1.014/1.014 win)|
|B CP2 Z|7.009 -> 7.005 (1.001/0.998* parity)|21.746 -> 21.361
(1.018/1.018 win)|23.282 -> 22.897 (1.017/1.016 win)|
|B CP4 C|3.708 -> 3.712 (0.999*/0.999* parity)|10.933 -> 10.934
(1.000/1.000 parity)|11.874 -> 11.882 (0.999*/1.000 parity)|
|B CP4 D|4.172 -> 4.173 (1.000/1.000 parity)|12.125 -> 11.949
(1.015/1.016 win)|13.252 -> 13.080 (1.013/1.014 win)|
|B CP4 Z|3.859 -> 3.859 (1.000/0.999* parity)|11.509 -> 11.516
(0.999*/1.000 parity)|12.369 -> 12.382 (0.999*/1.000 parity)|
|C CP2 C|6.480 -> 6.471 (1.001/0.999* parity)|23.238 -> 23.243
(1.000/1.000 parity)|23.476 -> 23.481 (1.000/1.000 parity)|
|C CP2 D|6.873 -> 6.874 (1.000/0.999* parity)|21.227 -> 20.890
(1.016/1.015 win)|22.985 -> 22.661 (1.014/1.015 win)|
|C CP2 Z|7.924 -> 7.927 (1.000/1.000 parity)|23.839 -> 23.559
(1.012/1.011 win)|26.005 -> 25.719 (1.011/1.011 win)|
|C CP4 C|3.680 -> 3.679 (1.000/1.000 parity)|11.043 -> 10.881
(1.015/1.014 win)|11.997 -> 11.839 (1.013/1.013 win)|
|C CP4 D|4.102 -> 4.104 (1.000/0.999* parity)|12.085 -> 12.021
(1.005/1.006 win)|13.028 -> 12.956 (1.006/1.006 win)|
|C CP4 Z|3.663 -> 3.667 (0.999*/0.999* parity)|10.976 -> 10.815
(1.015/1.015 win)|11.939 -> 11.776 (1.014/1.013 win)|
|D CP2 C|6.779 -> 6.784 (0.999*/0.998* parity)|21.069 -> 20.862
(1.010/1.011 win)|22.797 -> 22.606 (1.008/1.009 win)|
|D CP2 D|7.190 -> 7.194 (0.999*/1.000 parity)|21.776 -> 21.491
(1.013/1.013 win)|23.620 -> 23.338 (1.012/1.013 win)|
|D CP2 Z|7.142 -> 7.146 (0.999*/1.000 parity)|21.657 -> 21.375
(1.013/1.013 win)|23.499 -> 23.213 (1.012/1.012 win)|
|E CP2 C|6.762 -> 6.757 (1.001/0.999* parity)|21.090 -> 20.750
(1.016/1.016 win)|22.597 -> 22.255 (1.015/1.015 win)|
|E CP2 D|6.898 -> 6.890 (1.001/1.000 parity)|21.006 -> 20.757
(1.012/1.012 win)|22.825 -> 22.544 (1.012/1.012 win)|
|E CP2 Z|7.370 -> 7.363 (1.001/0.998* parity)|22.570 -> 22.273
(1.013/1.012 win)|24.422 -> 24.127 (1.012/1.012 win)|
|E CP4 C|3.582 -> 3.576 (1.002/0.999* parity)|10.701 -> 10.553
(1.014/1.014 win)|11.634 -> 11.489 (1.013/1.012 win)|
|E CP4 D|3.667 -> 3.668 (1.000/0.998* parity)|10.603 -> 10.431
(1.016/1.014 win)|12.006 -> 11.826 (1.015/1.013 win)|
|E CP4 Z|3.629 -> 3.627 (1.001/1.001 parity)|10.789 -> 10.612
(1.017/1.016 win)|11.717 -> 11.538 (1.016/1.016 win)|
|F CP2 C|6.512 -> 6.520 (0.999*/0.999* parity)|19.886 -> 19.671
(1.011/1.011 win)|21.562 -> 21.352 (1.010/1.011 win)|
|F CP2 D|6.558 -> 6.556 (1.000/0.999* parity)|19.472 -> 19.390
(1.004/1.003 win)|21.776 -> 21.706 (1.003/1.002 win)|
|F CP2 Z|6.435 -> 6.425 (1.002/1.000 parity)|20.031 -> 19.807
(1.011/1.011 win)|21.484 -> 21.230 (1.012/1.012 win)|
|F CP4 C|3.787 -> 3.788 (1.000/0.999* parity)|11.015 -> 11.008
(1.001/0.999* parity)|11.941 -> 11.941 (1.000/1.000 parity)|
|F CP4 D|3.653 -> 3.649 (1.001/1.000 parity)|10.387 -> 10.375
(1.001/1.002 parity/win)|11.825 -> 11.800 (1.002/1.002 parity/win)|
|F CP4 Z|3.590 -> 3.590 (1.000/0.999* parity)|10.632 -> 10.582
(1.005/1.004 win)|11.481 -> 11.398 (1.007/1.007 win)|
|G CP1 C|4.549 -> 4.525 (1.005/1.002 parity)|16.773 -> 16.595
(1.011/1.010 win)|16.929 -> 16.745 (1.011/1.010 win)|
|G CP2 C|3.324 -> 3.325 (1.000/1.001 parity)|9.887 -> 9.888 (1.000/1.000
parity)|10.510 -> 10.511 (1.000/0.999* parity)|
|G CP2 D|3.933 -> 3.930 (1.001/1.000 parity)|11.110 -> 10.962
(1.014/1.014 win)|12.810 -> 12.668 (1.011/1.011 win)|
|G CP2 Z|3.002 -> 3.009 (0.998*/0.998* parity)|9.647 -> 9.553
(1.010/1.009 win)|10.220 -> 10.126 (1.009/1.008 win)|
|G CP4 C|2.241 -> 2.208 (1.015/1.013 win)|5.876 -> 5.874 (1.000/0.999*
parity)|6.525 -> 6.522 (1.000/1.000 parity)|
|G CP4 D|2.543 -> 2.544 (1.000/0.999* parity)|6.570 -> 6.496
(1.011/1.010 win)|7.877 -> 7.794 (1.011/1.009 win)|
|G CP4 Z|2.079 -> 1.951 (1.066/1.064 win)|5.818 -> 5.532 (1.052/1.050
win)|6.339 -> 6.113 (1.037/1.036 win)|
|H CP1 C|5.301 -> 4.821 (1.100/1.018 win)|17.597 -> 17.190 (1.024/1.031
win)|17.743 -> 17.322 (1.024/1.031 win)|
|H CP1 C pre-norm|4.757 -> 4.271 (1.114/1.114 win)|15.948 -> 15.517
(1.028/1.061 win)|16.084 -> 15.653 (1.028/1.061 win)|
|H CP2 C|3.742 -> 3.736 (1.002/0.999* parity)|11.007 -> 11.001
(1.001/0.999* parity)|11.907 -> 11.905 (1.000/0.999* parity)|
|H CP2 D|3.737 -> 3.740 (0.999*/0.998* parity)|10.842 -> 10.607
(1.022/1.022 win)|12.230 -> 11.997 (1.019/1.019 win)|
|H CP2 Z|3.813 -> 3.815 (0.999*/1.000 parity)|11.469 -> 11.468
(1.000/1.000 parity)|12.276 -> 12.273 (1.000/1.000 parity)|
|H CP4 D|2.180 -> 2.174 (1.003/1.001 parity)|5.927 -> 5.812 (1.020/1.019
win)|6.855 -> 6.739 (1.017/1.017 win)|
|I CP1 C|2.785 -> 2.549 (1.093 win)|8.974 -> 8.761 (1.024/1.031
win)|9.115 -> 8.908 (1.023/1.030 win)|
|I CP1 C pre-norm|2.500 -> 2.263 (1.105/1.105 win)|8.121 -> 7.909
(1.027/1.059 win)|8.258 -> 8.047 (1.026/1.058 win)|
|I CP2 C|2.166 -> 2.164 (1.001/1.000 parity)|5.862 -> 5.848 (1.002/1.001
parity)|6.462 -> 6.451 (1.002/1.001 parity)|
|I CP2 D|2.145 -> 2.142 (1.001/1.000 parity)|5.833 -> 5.720 (1.020/1.021
win)|6.715 -> 6.597 (1.018/1.017 win)|
|I CP2 Z|2.065 -> 2.058 (1.003/1.000 parity)|5.566 -> 5.561 (1.001/1.002
parity)|6.294 -> 6.283 (1.002/1.001 parity)|
|I CP4 D|1.372 -> 1.368 (1.003/0.999* parity)|3.417 -> 3.361
(1.017/1.015 win)|4.184 -> 4.122 (1.015/1.011 win)|
|J CP1 C|2.588 -> 2.597 (0.997*/0.999* parity)|8.827 -> 8.837
(0.999*/0.999* parity)|8.979 -> 8.989 (0.999*/0.999* parity)|
|M CP1 C (control)|1.484 -> 1.383 (1.073/1.059 win)|[merge] 4.735 ->
4.563 (1.038/1.037 win)|[merge] 4.864 -> 4.697 (1.036/1.035 win)|
|M CP1 C pre-norm|1.323 -> 1.220 (1.084/1.083 win)|4.245 -> 4.079
(1.041/1.050 win)|4.383 -> 4.212 (1.041/1.048 win)|
|M CP2 C|1.263 -> 1.259 (1.003/0.998* parity)|3.095 -> 3.099
(0.999*/0.999* parity)|3.863 -> 3.860 (1.001/0.999* parity)|
|M CP2 D|1.379 -> 1.375 (1.003/1.001 parity)|3.378 -> 3.348 (1.009/1.007
parity/win)|4.141 -> 4.117 (1.006/1.007 parity)|
|M CP2 Z|1.299 -> 1.219 (1.066/1.060 win)|3.190 -> 3.187 (1.001/0.998*
parity)|4.141 -> 4.137 (1.001/0.996* parity)|
|M CP4 C|1.069 -> 1.073 (0.996*/0.993* parity)|2.539 -> 2.544
(0.998*/1.001 parity)|3.870 -> 3.866 (1.001 parity)|
|M CP4 D|1.047 -> 1.044 (1.003 parity)|2.523 -> 2.507 (1.006/1.001
parity)|3.680 -> 3.688 (0.998*/0.999* parity)|
|M CP4 Z|1.028 -> 1.024 (1.004 parity)|2.403 -> 2.402 (1.000/0.995*
parity)|3.711 -> 3.710 (1.000/0.999* parity)|
|N CP1 C|0.411 -> 0.410 (1.002/1.000 parity)|1.201 -> 1.198 (1.003
parity)|1.971 -> 1.968 (1.002 parity)|
|N CP2 C|0.874 -> 0.876 (0.998* parity)|2.041 -> 2.030 (1.005/0.992*
parity)|3.247 -> 3.259 (0.996*/0.993* parity)|
|N CP2 D|n/p unstable|2.212 -> 2.212 (1.000/0.996* parity)|3.587 ->
3.591 (0.999*/0.996* parity)|
|N CP2 Z|0.994 -> 0.996 (0.998* parity)|2.264 -> 2.254 (1.004/0.996*
parity)|3.568 -> 3.573 (0.999* parity)|
|N CP4 C|0.917 -> 0.914 (1.003 parity)|2.126 -> 2.133 (0.997*/0.997*
parity)|3.397 -> 3.396 (1.000/0.997* parity)|
|N CP4 D|0.989 -> 0.992 (0.997* parity)|2.190 -> 2.190 (1.000/0.998*
parity)|3.545 -> 3.541 (1.001/0.997* parity)|
|N CP4 Z|0.986 -> 0.986 (1.000/0.999* parity)|2.209 -> 2.216
(0.997*/0.996* parity)|3.585 -> 3.581 (1.001/0.997* parity)|
Single-rank proxies (single-rank proxy, not a full step): one GPU runs
the documents of rank r of the named CP configuration (documents in
[this
comment](https://github.com/flashinfer-ai/flashinfer/pull/8521#issuecomment-6612373113)).
|proxy of|fwd|fwd+bwd|ckpt|
|---|---|---|---|
|A CP16 D = H CP8 D r0|0.817 -> 0.815 (1.002/1.000 parity)|2.451 ->
2.416 (1.014/1.016 win)|2.669 -> 2.644 (1.009/1.012 win)|
|A CP8 D r0|1.369 -> 1.362 (1.005/1.000 parity)|4.551 -> 4.479
(1.016/1.018 win)|4.691 -> 4.619 (1.016/1.017 win)|
|C CP8 C r0|1.365 -> 1.359 (1.004 parity)|4.383 -> 4.381 (1.000/1.001
parity)|4.523 -> 4.518 (1.001/1.001 parity)|
|E CP16 D r0|0.789 -> 0.789 (1.000 parity)|2.394 -> 2.373 (1.009/1.008
win)|2.634 -> 2.609 (1.010/1.008 win)|
|E CP8 D r0|1.311 -> 1.305 (1.005/1.004 parity)|4.339 -> 4.285
(1.013/1.012 win)|4.475 -> 4.424 (1.012/1.012 win)|
|F CP16 D r0|0.776 -> 0.774 (1.003 parity)|2.316 -> 2.315 (1.000/1.000
parity)|2.553 -> 2.558 (0.998*/1.000 parity)|
|F CP8 D r0|1.300 -> 1.294 (1.005/1.002 parity)|4.292 -> 4.290
(1.000/1.001 parity)|4.429 -> 4.430 (1.000/1.002 parity)|
|G CP16 D r1|0.573 -> 0.568 (1.009/0.998* parity)|1.625 -> 1.614 (1.007
parity)|2.157 -> 2.141 (1.007/0.999* parity)|
|G CP16 Z r2|0.591 -> 0.590 (1.002/1.002 parity)|1.603 -> 1.605
(0.999*/1.001 parity)|2.119 -> 2.120 (1.000/1.001 parity)|
|G CP8 C r3|0.973 -> 0.975 (0.998*/1.002 parity)|2.780 -> 2.780
(1.000/1.001 parity)|2.924 -> 2.926 (0.999*/1.001 parity)|
|G CP8 D r1|0.809 -> 0.807 (1.002/0.999* parity)|2.477 -> 2.455
(1.009/1.008 win)|2.675 -> 2.653 (1.008/1.008 parity)|
|G CP8 Z r1|0.922 -> 0.893 (1.032/1.030 win)|2.670 -> 2.579 (1.035/1.034
win)|2.804 -> 2.721 (1.031/1.029 win)|
|H CP16 D = I CP8 D r0|0.556 -> 0.555 (1.002 parity)|1.585 -> 1.566
(1.012/1.010 win/parity)|2.143 -> 2.143 (1.000/1.008 parity/win)|
|I CP16 D r0|0.434 -> 0.437 (0.993* parity)|1.250 -> 1.249 (1.001
parity)|2.008 -> 2.017 (0.996* parity)|
|M CP16 D r0|0.409 -> 0.411 (0.995*/0.995* parity)|1.222 -> 1.226
(0.997* parity)|1.995 -> 2.020 (0.988*/0.993* parity)|
|M CP8 D r0|0.438 -> 0.442 (0.991*/0.998* parity)|1.209 -> 1.219
(0.992*/0.997* parity)|1.977 -> 1.976 (1.001/0.993* parity)|
|M CP8 Z r2|n/p unstable|1.220 -> 1.230 (0.992* parity)|2.010 -> 2.007
(1.001/0.991* parity)|
### B. Two baselines (same process: det(main), strongest non-Cake, this
PR)
|workload|op|det(main) ms|strongest non-Cake ms (arm)|this PR
ms|det(main) / this PR|non-Cake / this PR|
|---|---|---|---|---|---|---|
|A CP1 C|fwd|12.119|3.152 (FROST-ck0)|9.397|1.290|0.335*|
|A CP1 C|fwd+bwd|37.707|18.705 (FROST-ck64)|34.081|1.106|0.549*|
|A CP1 C|ckpt|37.883|23.031 (FROST-ck64)|34.260|1.106|0.672*|
|A CP2 C|fwd|8.213|11.456 (FROST-ck0)|7.051|1.165|1.625|
|A CP2 C|fwd+bwd|22.801|29.422 (FROST-ck64)|21.440|1.063|1.372|
|A CP2 C|ckpt|24.507|38.027 (FROST-ck64)|23.124|1.060|1.644|
|A CP2 D|fwd|8.359|11.457 (FROST-ck0)|7.276|1.149|1.575|
|A CP2 D|fwd+bwd|23.245|29.424 (FROST-ck64)|21.816|1.066|1.349|
|A CP2 D|ckpt|25.158|38.018 (FROST-ck64)|23.707|1.061|1.604|
|A CP2 Z|fwd|8.728|11.463 (FROST-ck0)|7.627|1.144|1.503|
|A CP2 Z|fwd+bwd|24.412|29.437 (FROST-ck64)|23.111|1.056|1.274|
|A CP2 Z|ckpt|26.284|38.029 (FROST-ck64)|24.865|1.057|1.529|
|A CP4 C|fwd|4.282|7.183 (FROST-ck0)|3.717|1.152|1.932|
|A CP4 C|fwd+bwd|11.446|18.745 (FROST-ck64)|10.883|1.052|1.722|
|A CP4 C|ckpt|12.355|24.417 (FROST-ck64)|11.772|1.050|2.074|
|A CP4 D|fwd|4.352|7.182 (FROST-ck0)|3.825|1.138|1.878|
|A CP4 D|fwd+bwd|11.481|18.745 (FROST-ck64)|10.729|1.070|1.747|
|A CP4 D|ckpt|12.914|24.421 (FROST-ck64)|12.167|1.061|2.007|
|A CP4 Z|fwd|4.360|7.182 (FROST-ck0)|3.820|1.141|1.880|
|A CP4 Z|fwd+bwd|11.855|18.748 (FROST-ck64)|11.128|1.065|1.685|
|A CP4 Z|ckpt|12.786|24.413 (FROST-ck64)|12.062|1.060|2.024|
|B CP2 C|fwd|8.167|11.449 (FROST-ck0)|6.999|1.167|1.636|
|B CP2 C|fwd+bwd|22.921|29.311 (FROST-ck64)|21.478|1.067|1.365|
|B CP2 C|ckpt|24.619|37.968 (FROST-ck64)|23.205|1.061|1.636|
|B CP2 D|fwd|8.287|11.450 (FROST-ck0)|7.151|1.159|1.601|
|B CP2 D|fwd+bwd|23.277|29.309 (FROST-ck64)|21.802|1.068|1.344|
|B CP2 D|ckpt|25.032|37.973 (FROST-ck64)|23.568|1.062|1.611|
|B CP2 Z|fwd|8.155|11.458 (FROST-ck0)|7.060|1.155|1.623|
|B CP2 Z|fwd+bwd|22.895|29.303 (FROST-ck64)|21.407|1.070|1.369|
|B CP2 Z|ckpt|24.445|37.977 (FROST-ck64)|22.964|1.064|1.654|
|B CP4 C|fwd|4.337|7.114 (FROST-ck0)|3.755|1.155|1.895|
|B CP4 C|fwd+bwd|11.526|18.247 (FROST-ck64)|10.950|1.053|1.666|
|B CP4 C|ckpt|12.489|23.824 (FROST-ck64)|11.919|1.048|1.999|
|B CP4 D|fwd|4.751|7.116 (FROST-ck0)|4.221|1.126|1.686|
|B CP4 D|fwd+bwd|12.702|18.250 (FROST-ck64)|11.982|1.060|1.523|
|B CP4 D|ckpt|13.841|23.828 (FROST-ck64)|13.118|1.055|1.816|
|B CP4 Z|fwd|4.418|7.113 (FROST-ck0)|3.919|1.127|1.815|
|B CP4 Z|fwd+bwd|12.069|18.248 (FROST-ck64)|11.575|1.043|1.577|
|B CP4 Z|ckpt|12.947|23.822 (FROST-ck64)|12.432|1.041|1.916|
|C CP2 C|fwd|7.099|10.355 (FLA-CP)|6.509|1.091|1.591|
|C CP2 C|fwd+bwd|23.774|32.712 (FROST-ck64)|23.244|1.023|1.407|
|C CP2 C|ckpt|24.097|42.426 (FROST-ck0)|23.516|1.025|1.804|
|C CP2 D|fwd|7.994|12.574 (FROST-ck0)|6.899|1.159|1.823|
|C CP2 D|fwd+bwd|22.350|32.707 (FROST-ck64)|20.925|1.068|1.563|
|C CP2 D|ckpt|24.088|42.426 (FROST-ck0)|22.679|1.062|1.871|
|C CP2 Z|fwd|9.039|12.575 (FROST-ck0)|7.953|1.137|1.581|
|C CP2 Z|fwd+bwd|24.931|32.708 (FROST-ck64)|23.573|1.058|1.388|
|C CP2 Z|ckpt|27.102|42.427 (FROST-ck0)|25.740|1.053|1.648|
|C CP4 C|fwd|4.291|6.993 (FROST-ck0)|3.724|1.152|1.878|
|C CP4 C|fwd+bwd|11.624|18.116 (FROST-ck64)|10.900|1.066|1.662|
|C CP4 C|ckpt|12.621|23.607 (FROST-ck64)|11.886|1.062|1.986|
|C CP4 D|fwd|4.621|6.995 (FROST-ck0)|4.147|1.114|1.687|
|C CP4 D|fwd+bwd|12.607|18.115 (FROST-ck64)|12.070|1.044|1.501|
|C CP4 D|ckpt|13.553|23.608 (FROST-ck64)|13.011|1.042|1.814|
|C CP4 Z|fwd|4.297|6.996 (FROST-ck0)|3.730|1.152|1.876|
|C CP4 Z|fwd+bwd|11.588|18.114 (FROST-ck64)|10.850|1.068|1.669|
|C CP4 Z|ckpt|12.590|23.599 (FROST-ck64)|11.848|1.063|1.992|
|D CP2 C|fwd|7.994|12.281 (FROST-ck0)|6.811|1.174|1.803|
|D CP2 C|fwd+bwd|22.277|31.962 (FROST-ck64)|20.892|1.066|1.530|
|D CP2 C|ckpt|23.994|41.457 (FROST-ck0)|22.647|1.059|1.831|
|D CP2 D|fwd|8.262|12.279 (FROST-ck0)|7.213|1.145|1.702|
|D CP2 D|fwd+bwd|22.826|31.958 (FROST-ck64)|21.493|1.062|1.487|
|D CP2 D|ckpt|24.670|41.464 (FROST-ck0)|23.345|1.057|1.776|
|D CP2 Z|fwd|8.250|12.296 (FROST-ck0)|7.213|1.144|1.705|
|D CP2 Z|fwd+bwd|22.782|31.928 (FROST-ck64)|21.455|1.062|1.488|
|D CP2 Z|ckpt|24.621|41.522 (FROST-ck64)|23.294|1.057|1.783|
|E CP2 C|fwd|7.803|10.970 (FROST-ck0)|6.778|1.151|1.618|
|E CP2 C|fwd+bwd|22.133|28.226 (FROST-ck64)|20.787|1.065|1.358|
|E CP2 C|ckpt|23.659|36.486 (FROST-ck64)|22.310|1.060|1.635|
|E CP2 D|fwd|7.980|10.972 (FROST-ck0)|6.930|1.152|1.583|
|E CP2 D|fwd+bwd|22.103|28.222 (FROST-ck64)|20.818|1.062|1.356|
|E CP2 D|ckpt|23.899|36.498 (FROST-ck64)|22.600|1.057|1.615|
|E CP2 Z|fwd|8.443|10.970 (FROST-ck0)|7.376|1.145|1.487|
|E CP2 Z|fwd+bwd|23.626|28.223 (FROST-ck64)|22.279|1.060|1.267|
|E CP2 Z|ckpt|25.480|36.483 (FROST-ck64)|24.130|1.056|1.512|
|E CP4 C|fwd|4.157|6.744 (FROST-ck0)|3.619|1.149|1.863|
|E CP4 C|fwd+bwd|11.269|17.339 (FROST-ck64)|10.585|1.065|1.638|
|E CP4 C|ckpt|12.224|22.625 (FROST-ck64)|11.533|1.060|1.962|
|E CP4 D|fwd|4.267|6.738 (FROST-ck0)|3.716|1.148|1.813|
|E CP4 D|fwd+bwd|11.198|17.337 (FROST-ck64)|10.478|1.069|1.655|
|E CP4 D|ckpt|12.607|22.627 (FROST-ck64)|11.889|1.060|1.903|
|E CP4 Z|fwd|4.187|6.742 (FROST-ck0)|3.671|1.141|1.837|
|E CP4 Z|fwd+bwd|11.333|17.335 (FROST-ck64)|10.637|1.065|1.630|
|E CP4 Z|ckpt|12.280|22.628 (FROST-ck64)|11.573|1.061|1.955|
|F CP2 C|fwd|7.636|10.658 (FROST-ck0)|6.534|1.169|1.631|
|F CP2 C|fwd+bwd|21.006|27.060 (FROST-ck64)|19.705|1.066|1.373|
|F CP2 C|ckpt|22.681|35.094 (FROST-ck64)|21.379|1.061|1.642|
|F CP2 D|fwd|7.664|10.657 (FROST-ck0)|6.585|1.164|1.618|
|F CP2 D|fwd+bwd|20.538|27.069 (FROST-ck64)|19.394|1.059|1.396|
|F CP2 D|ckpt|22.869|35.107 (FROST-ck64)|21.721|1.053|1.616|
|F CP2 Z|fwd|7.539|10.664 (FROST-ck0)|6.462|1.167|1.650|
|F CP2 Z|fwd+bwd|21.137|27.067 (FROST-ck64)|19.833|1.066|1.365|
|F CP2 Z|ckpt|22.593|35.108 (FROST-ck64)|21.279|1.062|1.650|
|F CP4 C|fwd|4.291|5.857 (FROST-ck0)|3.825|1.122|1.531|
|F CP4 C|fwd+bwd|11.501|14.962 (FROST-ck64)|11.059|1.040|1.353|
|F CP4 C|ckpt|12.444|19.362 (FROST-ck64)|11.986|1.038|1.615|
|F CP4 D|fwd|4.243|5.855 (FROST-ck0)|3.674|1.155|1.594|
|F CP4 D|fwd+bwd|10.980|14.958 (FROST-ck64)|10.385|1.057|1.440|
|F CP4 D|ckpt|12.429|19.368 (FROST-ck64)|11.827|1.051|1.638|
|F CP4 Z|fwd|4.137|5.861 (FROST-ck0)|3.637|1.137|1.611|
|F CP4 Z|fwd+bwd|11.165|14.966 (FROST-ck64)|10.627|1.051|1.408|
|F CP4 Z|ckpt|12.031|19.360 (FROST-ck64)|11.441|1.052|1.692|
|G CP1 C|fwd|5.661|1.671 (FROST-ck0)|4.551|1.244|0.367*|
|G CP1 C|fwd+bwd|17.861|9.765 (FROST-ck64)|16.583|1.077|0.589*|
|G CP1 C|ckpt|18.013|11.984 (FROST-ck64)|16.740|1.076|0.716*|
|G CP2 C|fwd|3.788|5.845 (FROST-ck0)|3.359|1.128|1.740|
|G CP2 C|fwd+bwd|10.317|15.350 (FROST-ck64)|9.892|1.043|1.552|
|G CP2 C|ckpt|11.091|19.833 (FROST-ck64)|10.549|1.051|1.880|
|G CP2 D|fwd|4.566|5.842 (FROST-ck0)|3.972|1.150|1.471|
|G CP2 D|fwd+bwd|11.719|15.349 (FROST-ck64)|10.979|1.067|1.398|
|G CP2 D|ckpt|13.421|19.830 (FROST-ck64)|12.678|1.059|1.564|
|G CP2 Z|fwd|3.539|5.844 (FROST-ck0)|3.023|1.171|1.933|
|G CP2 Z|fwd+bwd|10.230|15.351 (FROST-ck64)|9.582|1.068|1.602|
|G CP2 Z|ckpt|10.800|19.833 (FROST-ck64)|10.176|1.061|1.949|
|G CP4 C|fwd|2.590|3.427 (FROST-ck0)|2.257|1.148|1.518|
|G CP4 C|fwd+bwd|6.247|8.459 (FROST-ck64)|5.917|1.056|1.430|
|G CP4 C|ckpt|7.534|11.082 (FROST-ck64)|6.553|1.150|1.691|
|G CP4 D|fwd|2.912|3.429 (FROST-ck0)|2.599|1.120|1.319|
|G CP4 D|fwd+bwd|6.916|8.460 (FROST-ck64)|6.533|1.059|1.295|
|G CP4 D|ckpt|8.465|11.085 (FROST-ck64)|7.856|1.078|1.411|
|G CP4 Z|fwd|2.392|3.424 (FROST-ck0)|1.992|1.201|1.719|
|G CP4 Z|fwd+bwd|6.127|8.459 (FROST-ck64)|5.583|1.097|1.515|
|G CP4 Z|ckpt|6.691|11.088 (FROST-ck64)|6.198|1.080|1.789|
|H CP1 C|fwd+bwd|19.030|9.946 (FROST-ck64)|17.215|1.105|0.578*|
|H CP1 C|ckpt|19.220|12.273 (FROST-ck64)|17.382|1.106|0.706*|
|I CP1 C|fwd|3.283|1.027 (FROST-ck0)|2.581|1.272|0.398*|
|I CP1 C|fwd+bwd|9.722|5.311 (FROST-ck64)|8.787|1.106|0.604*|
|J CP1 C|fwd|3.077|1.113 (FROST-ck0)|2.600|1.183|0.428*|
|J CP1 C|fwd+bwd|9.338|5.596 (FROST-ck64)|8.854|1.055|0.632*|
|J CP1 C|ckpt|9.489|6.943 (FROST-ck64)|9.004|1.054|0.771*|
|M CP1 C|fwd+bwd|5.049|3.063 (FROST-ck64)|4.565|1.106|0.671*|
|M CP2 C|fwd+bwd|4.674|4.508 (FROST-ck64)|3.117|1.500|1.446|
|M CP2 Z|fwd+bwd|3.928|4.540 (FROST-ck64)|3.284|1.196|1.382|
|N CP4 C|fwd|1.355|1.313 (FLA-CP)|0.992|1.366|1.324|
|N CP4 C|fwd+bwd|3.175|3.245 (FLA-CP)|2.343|1.355|1.385|
|N CP4 C|ckpt|4.812|5.165 (FLA-CP)|3.776|1.274|1.368|
|N CP4 Z|fwd+bwd|3.165|3.750 (FLA-CP)|2.366|1.338|1.585|
|N CP4 Z|ckpt|4.869|6.085 (FLA-CP)|3.889|1.252|1.565|
Published rows: 138; det(main) faster in 0, non-Cake faster in 14; 33
unstable rows not published (listed at the end).
## R200 (sm_107a, 208 SMs)
### A. This PR vs previous version (same process)
|workload|fwd: previous ms -> this PR ms (ratio)|fwd+bwd|ckpt (2 x fwd +
bwd)|
|---|---|---|---|
|A CP1 C|6.528 -> 6.541 (0.998*/0.997* parity)|23.815 -> 23.790
(1.001/1.000 parity)|23.865 -> 23.870 (1.000/1.000 parity)|
|G CP1 C|2.953 -> 2.951 (1.001/1.000 parity)|11.238 -> 11.235
(1.000/1.001 parity)|11.369 -> 11.360 (1.001/1.001 parity)|
|J CP1 C|1.803 -> 1.808 (0.997*/0.994* parity)|6.170 -> 6.183
(0.998*/0.999* parity/REGRESSION)|6.315 -> 6.325 (0.998*/0.999* parity)|
|M CP1 C|0.990 -> 0.908 (1.090/1.092 win)|3.276 -> 3.110 (1.053/1.052
win)|3.360 -> 3.199 (1.050/1.049 win)|
|M CP1 C pre-norm|0.884 -> 0.800 (1.105/1.109 win)|3.004 -> 2.848
(1.055/1.061 win)|3.133 -> 2.976 (1.053/1.054 win)|
|N CP1 C|0.191 -> 0.192 (0.995*/1.005 parity)|0.692 -> 0.696
(0.994*/0.997* parity)|1.264 -> 1.265 (0.999*/0.994* parity)|
Single-rank proxies (single-rank proxy, not a full step): one GPU runs
the documents of rank r of the named CP configuration (documents in
[this
comment](https://github.com/flashinfer-ai/flashinfer/pull/8521#issuecomment-6612373113)).
|proxy of|fwd|fwd+bwd|ckpt|
|---|---|---|---|
|A CP2 C r0|2.861 -> 2.860 (1.000/0.998* parity)|11.066 -> 11.055
(1.001/1.000 parity)|11.193 -> 11.180 (1.001/1.002 parity)|
|A CP2 D r0|2.921 -> 2.914 (1.002/0.999* parity)|11.205 -> 11.204
(1.000/1.001 parity)|11.346 -> 11.337 (1.001/1.001 parity)|
|A CP2 Z r1|2.849 -> 2.853 (0.999*/1.004 parity)|11.190 -> 11.190
(1.000/1.001 parity)|11.311 -> 11.313 (1.000/1.000 parity)|
|A CP4 Z = B CP4 Z r0|1.556 -> 1.556 (1.000/0.999* parity)|5.771 ->
5.778 (0.999*/1.002 parity)|5.853 -> 5.874 (0.996*/1.001 parity)|
|A CP8 C r0|0.841 -> 0.838 (1.004/1.004 parity)|3.025 -> 3.024
(1.000/1.002 parity)|3.151 -> 3.150 (1.000/1.001 parity)|
|B CP16 D r0|0.500 -> 0.485 (1.031/1.027 win)|1.651 -> 1.616
(1.022/1.019 win)|1.728 -> 1.695 (1.019/1.010 win)|
|B CP2 C r0|2.901 -> 2.895 (1.002/0.995* parity)|11.035 -> 11.017
(1.002/0.999* parity)|11.108 -> 11.094 (1.001/0.999* parity)|
|B CP2 D r0|3.040 -> 3.033 (1.002/1.008 parity)|11.498 -> 11.486
(1.001/1.001 parity)|11.617 -> 11.605 (1.001/1.001 parity)|
|B CP2 Z r1|2.972 -> 2.976 (0.999*/1.001 parity)|11.206 -> 11.208
(1.000/0.999* parity)|11.306 -> 11.320 (0.999*/0.999* parity)|
|B CP4 D r0|1.498 -> 1.497 (1.001/0.997* parity)|5.659 -> 5.663
(0.999*/0.999* parity)|5.784 -> 5.786 (1.000/0.998* parity)|
|B CP8 D r0|0.880 -> 0.844 (1.043/1.052 win)|3.136 -> 3.066 (1.023/1.026
win)|3.260 -> 3.189 (1.022/1.025 win)|
|C CP2 C r1|2.950 -> 2.954 (0.999*/1.001 parity)|11.018 -> 11.023
(1.000/1.002 parity)|11.138 -> 11.140 (1.000/1.001 parity)|
|C CP2 D r0|3.059 -> 3.053 (1.002/1.003 parity)|11.252 -> 11.243
(1.001/1.001 parity)|11.323 -> 11.346 (0.998*/1.001 parity)|
|C CP2 Z r0|2.779 -> 2.782 (0.999*/0.999* parity)|10.906 -> 10.899
(1.001/1.001 parity)|11.018 -> 11.020 (1.000/1.001 parity)|
|C CP4 C r0|1.530 -> 1.536 (0.996*/0.997* REGRESSION/parity)|5.650 ->
5.654 (0.999*/0.999* parity)|5.768 -> 5.772 (0.999*/0.999* parity)|
|C CP4 Z r0|1.535 -> 1.532 (1.002/0.995* parity/REGRESSION)|5.652 ->
5.648 (1.001/0.999* parity)|5.779 -> 5.775 (1.001/0.999* parity)|
|C CP8 C r0|0.825 -> 0.824 (1.001/1.002 parity)|2.958 -> 2.957
(1.000/1.002 parity)|3.077 -> 3.074 (1.001/1.003 parity)|
|D CP2 C r0|2.933 -> 2.932 (1.000/0.999* parity)|10.813 -> 10.822
(0.999*/1.001 parity)|10.958 -> 10.954 (1.000/1.001 parity)|
|E CP16 D r0|0.493 -> 0.492 (1.002/1.000 parity)|1.651 -> 1.651
(1.000/1.000 parity)|1.826 -> 1.821 (1.003/1.002 parity)|
|E CP2 C r1|2.836 -> 2.838 (0.999*/0.998* parity)|10.615 -> 10.611
(1.000/1.000 parity)|10.716 -> 10.716 (1.000/1.000 parity)|
|E CP2 D r0|3.066 -> 3.064 (1.001/1.001 parity)|11.292 -> 11.287
(1.000/1.000 parity)|11.408 -> 11.396 (1.001/1.000 parity)|
|E CP2 Z r0|2.921 -> 2.908 (1.004/1.006 parity)|10.746 -> 10.739
(1.001/1.001 parity)|10.831 -> 10.827 (1.000/1.000 parity)|
|E CP4 C r3|1.472 -> 1.479 (0.995*/0.999* REGRESSION/parity)|5.437 ->
5.451 (0.997*/0.999* REGRESSION/parity)|5.561 -> 5.572 (0.998*/0.999*
parity)|
|E CP4 D r0|1.590 -> 1.492 (1.066/1.064 win)|5.749 -> 5.535 (1.039/1.039
win)|5.879 -> 5.662 (1.038/1.038 win)|
|E CP4 Z r2|1.548 -> 1.479 (1.047/1.051 win)|5.629 -> 5.501 (1.023/1.025
win)|5.746 -> 5.622 (1.022/1.024 win)|
|E CP8 C r1|0.794 -> 0.793 (1.001/0.996* parity/REGRESSION)|2.861 ->
2.859 (1.001/0.999* parity)|2.987 -> 2.985 (1.001/0.997* parity)|
|E CP8 D r0|0.852 -> 0.840 (1.014/1.002 win/parity)|2.975 -> 2.969
(1.002/1.001 win/parity)|3.093 -> 3.085 (1.003/1.002 parity)|
|F CP16 D r0|0.497 -> 0.472 (1.053/1.052 win)|1.636 -> 1.583
(1.033/1.036 win)|1.818 -> 1.792 (1.015/1.016 win)|
|F CP2 C r0|2.801 -> 2.798 (1.001/1.001 parity)|10.385 -> 10.397
(0.999*/1.001 parity)|10.472 -> 10.480 (0.999*/1.001 parity)|
|F CP2 D r0|2.625 -> 2.626 (1.000/1.001 parity)|10.117 -> 10.114
(1.000/1.001 parity/win)|10.242 -> 10.229 (1.001/1.001 parity/win)|
|F CP2 Z r1|2.799 -> 2.784 (1.005/1.003 parity)|10.453 -> 10.450
(1.000/1.000 parity)|10.564 -> 10.565 (1.000/1.000 parity)|
|F CP4 C r0|1.452 -> 1.457 (0.997*/1.001 parity)|5.325 -> 5.320
(1.001/0.999* parity)|5.401 -> 5.395 (1.001/1.005 parity)|
|F CP4 D r2|1.520 -> 1.424 (1.067/1.073 win)|5.451 -> 5.239 (1.040/1.041
win)|5.574 -> 5.368 (1.038/1.040 win)|
|F CP4 Z r0|1.548 -> 1.448 (1.069/1.076 win)|5.581 -> 5.393 (1.035/1.037
win)|5.717 -> 5.522 (1.035/1.037 win)|
|F CP4 Z r0 pre-norm|1.361 -> 1.260 (1.080/1.079 win)|5.031 -> 4.843
(1.039/1.043 win)|5.191 -> 5.008 (1.037/1.041 win)|
|F CP8 D r0|0.859 -> 0.808 (1.063/1.065 win)|2.955 -> 2.839 (1.041/1.038
win)|3.136 -> 3.022 (1.038/1.036 win)|
|G CP16 D r1|0.314 -> 0.313 (1.003/1.003 parity)|0.965 -> 0.965
(1.000/1.000 parity)|1.187 -> 1.185 (1.002/1.003 parity)|
|G CP2 C r0|1.591 -> 1.594 (0.998*/1.001 parity)|5.911 -> 5.916
(0.999*/1.000 parity)|6.043 -> 6.044 (1.000/0.999* parity)|
|G CP2 D r1|1.487 -> 1.489 (0.999*/1.001 parity)|5.722 -> 5.722
(1.000/1.000 parity)|5.849 -> 5.850 (1.000/1.000 parity)|
|G CP2 Z r1|1.702 -> 1.703 (0.999*/1.001 parity)|6.219 -> 6.217
(1.000/1.000 parity)|6.303 -> 6.281 (1.004/1.001 parity)|
|G CP4 C r1|0.941 -> 0.939 (1.002/0.998* parity)|3.208 -> 3.193
(1.005/1.001 parity)|3.338 -> 3.333 (1.002/1.001 parity)|
|G CP4 D r1|0.813 -> 0.813 (1.000/1.006 parity)|3.052 -> 3.048
(1.001/1.002 parity)|3.176 -> 3.171 (1.002/1.002 parity)|
|G CP4 Z r3|1.042 -> 0.956 (1.090/1.091 win)|3.401 -> 3.218 (1.057/1.059
win)|3.491 -> 3.308 (1.055/1.057 win)|
|G CP4 Z r3 pre-norm|0.924 -> 0.837 (1.104/1.105 win)|3.100 -> 2.922
(1.061/1.064 win)|3.225 -> 3.049 (1.058/1.060 win)|
|G CP8 D r1|0.482 -> 0.482 (1.000/1.000 parity)|1.705 -> 1.705
(1.000/1.000 parity)|1.727 -> 1.728 (0.999*/1.003 parity)|
|G CP8 Z r3|0.559 -> 0.559 (1.000/1.002 parity)|1.828 -> 1.826
(1.001/1.000 parity)|1.945 -> 1.945 (1.000/1.002 parity)|
|H CP2 C r0|1.539 -> 1.541 (0.999*/0.997* parity)|5.800 -> 5.803
(0.999*/1.001 parity)|5.916 -> 5.919 (0.999*/1.000 parity)|
|H CP2 Z r0|1.688 -> 1.557 (1.084/1.083 win)|6.104 -> 5.879 (1.038/1.041
win)|6.224 -> 6.001 (1.037/1.040 win)|
|H CP2 Z r0 pre-norm|1.465 -> 1.352 (1.084/1.084 win)|5.449 -> 5.254
(1.037/1.044 win)|5.505 -> 5.305 (1.038/1.043 win)|
|I CP2 C r1|0.978 -> 0.976 (1.002/1.003 parity)|3.304 -> 3.296
(1.002/1.003 parity)|3.426 -> 3.417 (1.003/1.003 parity)|
|M CP16 D r0|0.178 -> 0.178 (1.000 parity)|0.514 -> 0.516 (0.996*/0.997*
parity)|1.263 -> 1.263 (1.000/0.995* parity)|
|M CP2 C r0|0.515 -> 0.514 (1.002/1.000 parity)|1.721 -> 1.718
(1.002/1.001 parity)|1.878 -> 1.875 (1.002/1.003 parity)|
|M CP2 D r0|0.486 -> 0.486 (1.000/1.002 parity)|1.684 -> 1.682
(1.001/1.001 parity)|1.870 -> 1.869 (1.001/0.999* parity)|
|M CP2 Z r0|0.510 -> 0.485 (1.052/1.049 win)|1.720 -> 1.673 (1.028/1.030
win)|1.879 -> 1.849 (1.016/1.014 win)|
|M CP4 D r0|0.310 -> 0.310 (1.000/1.000 parity)|1.063 -> 1.063
(1.000/1.004 parity)|1.419 -> 1.415 (1.003/1.001 parity)|
|M CP8 D r0|0.219 -> 0.219 (1.000/1.009 parity)|0.792 -> 0.790
(1.003/0.996* parity)|1.269 -> 1.267 (1.002/0.996* parity)|
### B. Two baselines (same process: det(main), strongest non-Cake, this
PR)
|workload|op|det(main) ms|strongest non-Cake ms (arm)|this PR
ms|det(main) / this PR|non-Cake / this PR|
|---|---|---|---|---|---|---|
|A CP1 C|fwd|8.011|2.143 (FROST-ck0)|[this PR] 6.674|1.200|0.321*|
|A CP1 C|fwd+bwd|25.594|17.449 (FROST-ck64)|[this PR]
24.320|1.052|0.717*|
|A CP1 C|ckpt|25.719|20.369 (FROST-ck0)|[this PR] 24.423|1.053|0.834*|
|G CP1 C|fwd|3.897|1.046 (FROST-ck0)|[this PR] 3.000|1.299|0.349*|
|G CP1 C|fwd+bwd|12.350|8.750 (FROST-ck64)|[this PR]
11.447|1.079|0.764*|
|G CP1 C|ckpt|12.466|10.283 (FROST-ck64)|[this PR] 11.568|1.078|0.889*|
|J CP1 C|fwd|2.138|0.612 (FROST-ck0)|[this PR] 1.840|1.162|0.333*|
|J CP1 C|fwd+bwd|6.602|4.640 (FROST-ck64)|[this PR] 6.301|1.048|0.736*|
|J CP1 C|ckpt|6.648|5.443 (FROST-ck64)|[this PR] 6.352|1.047|0.857*|
|M CP1 C|fwd|1.211|0.372 (FROST-ck0)|[this PR] 0.924|1.311|0.403*|
|M CP1 C|fwd+bwd|3.586|2.394 (FROST-ck64)|[this PR] 3.216|1.115|0.744*|
|M CP1 C|ckpt|3.705|3.003 (FROST-ck64)|[this PR] 3.347|1.107|0.897*|
|N CP1 C|ckpt|1.478|1.182 (FROST-ck0)|[this PR] 1.278|1.156|0.925*|
Published rows: 13; det(main) faster in 0, non-Cake faster in 13; 2
unstable rows not published (listed at the end).
R200 single-rank proxies: table B in [this
comment](https://github.com/flashinfer-ai/flashinfer/pull/8521#issuecomment-6614291049)
(same columns and rules).
## Multi-node CP16 (sub-table, published runs; not part of tables A / B)
|GPU (nodes x GPUs)|workload|operation|run|protocol|previous ms|this PR
ms|previous / this PR|det(main) ms|det(main) / this PR|class vs
previous|
|---|---|---|---|---|---|---|---|---|---|---|
|B200 (2 x 8)|A CP16
D|ckpt|run3|short|93.466|92.180|1.014|92.059|0.999*|parity|
|GB300 (4 x 4)|F CP16
D|fwd+bwd|run1|short|41.894|42.062|0.996*|42.240|1.004|parity|
|GB300 (4 x 4)|F CP16
D|fwd+bwd|run2|short|40.499|40.557|0.999*|40.730|1.004|parity|
|GB300 (4 x 4)|F CP16
D|ckpt|run2|short|63.689|63.608|1.001|63.815|1.003|parity|
## Cells not published or not measured
- B200, table A, power-capped (SW power-cap share >= 5 % in a short-call
run): A CP1 C pre-norm fwd+bwd, A CP1 C pre-norm ckpt.
- B200, table B, unstable: group drift above 2 % (largest over the
record's arms): A CP1 C fwd, A CP2 D ckpt, G CP1 C ckpt, H CP1 C ckpt, M
CP1 C fwd, M CP1 C fwd+bwd, N CP1 C fwd, N CP1 C fwd+bwd, N CP1 C ckpt,
N CP2 C fwd, N CP2 D fwd, N CP2 D fwd+bwd, N CP2 D ckpt, N CP2 Z fwd, N
CP2 Z ckpt.
- GB300, table A, unstable (group drift above 2 % or paired ratios not
within 1 %): N CP2 D fwd, proxy M CP8 Z r2 fwd.
- GB300, table B, unstable: group drift above 2 % (largest over the
record's arms): H CP1 C fwd, I CP1 C ckpt, M CP1 C fwd, M CP1 C ckpt, M
CP2 C fwd, M CP2 C ckpt, M CP2 D fwd, M CP2 D fwd+bwd, M CP2 D ckpt, M
CP2 Z fwd, M CP2 Z ckpt, M CP4 C fwd, M CP4 C fwd+bwd, M CP4 C ckpt, M
CP4 D fwd, M CP4 D fwd+bwd, M CP4 D ckpt, M CP4 Z fwd, M CP4 Z fwd+bwd,
M CP4 Z ckpt, N CP1 C fwd, N CP1 C fwd+bwd, N CP1 C ckpt, N CP2 C fwd, N
CP2 C fwd+bwd, N CP2 C ckpt, N CP2 D fwd, N CP2 D fwd+bwd, N CP2 D ckpt,
N CP2 Z fwd, N CP2 Z fwd+bwd, N CP2 Z ckpt, N CP4 Z fwd.
- R200, table B, unstable: group drift above 2 % (largest over the
record's arms): N CP1 C fwd, N CP1 C fwd+bwd.
- B200, table CP16, unstable (group drift above 2 % or paired ratios not
within 1 %): F CP16 D fwd+bwd (2 nodes, run1), F CP16 D fwd+bwd (2
nodes, run2), F CP16 D ckpt (2 nodes, run1), F CP16 D ckpt (2 nodes,
run2).
- GB300, table CP16, unstable (group drift above 2 % or paired ratios
not within 1 %): F CP16 D ckpt (4 nodes, run1).
- B200, 2 nodes: shape A CP16 D, second arm order not measured
(multi-node time limit).
- GB300, 4 nodes: shape A CP16 D not measured (multi-node time limit).
- R200: no multi-node window; the CP > 1 configurations are covered by
single-rank proxies only. R200 table A rows were timed at an earlier
commit of this PR with identical launch forms; R200 table B rows (CP1
and single-rank proxies) were timed at this head.
This PR does not change the other Cake GDN training variants (the
single-rank packages `cake_gdn_train` and `cake_gdn_det_train` and the
deterministic KCP package `cake_gdn_det_train_kcp`); architecture
coverage stays sm_100a, sm_103a and sm_107a.
---------
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: linear attention
flashinfer.experimental.cake_gdn_train_kcpis the non-deterministic KCP GDN training package. It supports BF16, the contiguous, zigzag and zigzag-per-document KCP layouts, and runs with or without in-kernel Q/K L2 norm. It covers forward, backward and checkpointed calls, in eager and CUDA graph modes. This PR changes only the host side. When q / k / v are plain split views of one packed QKV tensor (row.split+view), the backward now gives three contiguous dq / dk / dv to torch's split backward. The split-view backward then costs one contiguous concatenation instead of the slower path on main. Only q / k / v that come from onepacked_qkv_views_kcp(orpacked_qkv_views_det_kcp) split node keep the single packed dQKV row. There are no device code changes, and outputs and gradients are bitwise unchanged. Addresses #8367.All rows use the PR head after it merged main, so the forward runs the fused forward output walk.
Compared with the same package on FlashInfer main (CP1, norm off), the plain-split-view backward gets faster:
Packed and contiguous forms are at parity (0.998-1.006x across all rows).
Against the deterministic Cake package in the same process (speedup = baseline ms / this PR ms), this PR is faster with plain split views:
With packed views, bwd is at parity (1.000-1.003x), ckpt is 1.051-1.065x and fwd is 1.167-1.219x. Against the previous version, this PR's fwd and packed-view times are at parity. So the fwd and packed-view ckpt gains over deterministic Cake do not come from this PR.
Against the strongest qualified non-Cake implementation (both Q/K/V forms, norm off and on), this PR is slower on every claimed row:
B200 (sm_100a, 148 SMs)
NVIDIA B200, sm_100a, 148 SMs. PyTorch 2.14.1 (cu130, CUDA 13.0 runtime), CUDA 13.4 toolkit, BF16. Workload: (H_qk, HV) = (16, 48), D = 128. CP1 = one rank, 32,768 tokens in five documents [12000, 9000, 6000, 3768, 2000], contiguous layout. CP4 zigzag = 4 ranks, documents [49152, 32768, 28672, 20480], 32,768 tokens per rank; each sample takes the slowest rank. Q/K/V forms: split =
row.split+view(plain views), packed =packed_qkv_views_kcp, contig = contiguous copies. Modes: fwd = forward only; bwd = autograd backward after an eager forward. ckpt =torch.utils.checkpoint(use_reentrant=False), 2 forwards + 1 backward. graph = the package'sgraph="replay"mode with a caller-held workspace. Timing: 2 runs x 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls. Run 2 reverses the arm order. The shown value is the median of group medians. The previous version is a byte copy of the package from FlashInfer main before this PR. It runs in the same process as this PR, with the same build and libraries. Speedup = previous ms / this PR ms. Power-capped = GPU software power-cap share >= 5 % of an arm's timed-block samples, for any arm of a mode and norm setting. When the short runs hit the cap, the rows of that mode come from a sustained protocol instead. The sustained protocol uses 3 palindromic block pairs (block = 5 s idle + 20 warm-up calls + 300 timed calls). On B200, this applies to the CP1 norm-off eager fwd and bwd rows, all CP1 graph rows and the CP4 graph bwd norm-on rows. In the sustained runs, the CP1 norm-off eager fwd, eager bwd and graph bwd rows still hit the cap. They are marked "power-capped, not claimed". The CP1 norm-off eager fwd rows also fail the stability check (drift = last / first group - 1; every arm of the mode and norm setting must stay within 2 %; the previous version's split arm drifted 2.26 %). They are also marked "timing unstable". All other B200 rows are claimed.CP1
CP4 zigzag
Two-baseline comparison (CP1)
These are separate runs from the previous-version tables above, so this PR's absolute times differ slightly between the two sets. Each run uses one process for deterministic Cake, this PR, FLA and cuDNN FROST. Run 2 reverses the arm order of run 1. Each run has 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls per arm. The shown value is the median of group medians. The norm-off eager fwd and bwd rows come from two sustained runs (3 block pairs; block = 5 s idle + 20 warm-up calls + 300 timed calls). No deterministic Cake, this PR or FLA cell is power-capped (every in-process arm < 5 %); the separate FlashQLA process has no power-cap record. Deterministic Cake is
cake_gdn_det_train_kcp(chunk_gdn_det_kcp) on the same FlashInfer tree, unchanged by this PR. Its packed views use its own helper,packed_qkv_views_det_kcp. Its outputs and gradients are bitwise equal to this PR's. Strongest non-Cake is the fastest arm per row that passes the accuracy check. The check needs rel-RMS <= 1e-2 and max-abs <= 2e-2 of max|ref| on every tensor, with finite values. Candidates are FLA 0.6.0 Tritonchunk_gated_delta_rule, cuDNN frontend 1.32.1gated_delta_netand FlashQLA 0.1.3. The cuDNN plan is gdn_frost (cuDNN FROST): ck0 = the backward recomputes the states; ck64 = the forward stores BF16 states every 64 tokens. In "FlashQLA + FLA l2norm outside", FLA's Triton l2norm normalizes q / k outside the kernel (timed), and the in-kernel flag is off. On B200, cuDNN FROST fails the check on dbeta (max-abs 0.135 of max|ref|), although it is faster than this PR in all 24 rows. FlashQLA is therefore the strongest qualified arm. FlashQLA cannot share a process with the Cake package, because its TileLang 0.1.15 needs apache-tvm-ffi < 0.1.13. Its values come from a separate process that used the short protocol. Every B200 non-Cake ratio is therefore marked "cross-run, not claimed".GB300 (sm_103a, 152 SMs)
NVIDIA GB300, sm_103a, 152 SMs. PyTorch 2.14.1 (cu130, CUDA 13.0 runtime), CUDA 13.4 toolkit, BF16. Workload: (H_qk, HV) = (16, 48), D = 128. CP1 = one rank, 32,768 tokens in five documents [12000, 9000, 6000, 3768, 2000], contiguous layout. CP4 zigzag = 4 ranks, documents [49152, 32768, 28672, 20480], 32,768 tokens per rank; each sample takes the slowest rank. Q/K/V forms: split =
row.split+view(plain views), packed =packed_qkv_views_kcp, contig = contiguous copies. Modes: fwd = forward only; bwd = autograd backward after an eager forward. ckpt =torch.utils.checkpoint(use_reentrant=False), 2 forwards + 1 backward. graph = the package'sgraph="replay"mode with a caller-held workspace. Timing: 2 runs x 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls. Run 2 reverses the arm order. The shown value is the median of group medians. The previous version is a byte copy of the package from FlashInfer main before this PR. It runs in the same process as this PR, with the same build and libraries. Speedup = previous ms / this PR ms. No GB300 row is power-capped (GPU software power-cap share >= 5 %). In the CP4 eager ckpt norm-on rows, the time level shifted within run 2: the drift (last / first group - 1) was 2.2-8.5 % there, and run 1 stayed within 0.3 %. These rows pass the paired check instead: every paired group ratio stays within ±1 % of its median. They are claimed and marked "level shift in run 2" with the paired ratio. The CP4 eager fwd norm-off rows come from a re-measure with two further runs.CP1
CP4 zigzag
Two-baseline comparison (CP1)
These are separate runs from the previous-version tables above, so this PR's absolute times differ slightly between the two sets. Each run uses one process for deterministic Cake, this PR, FLA and cuDNN FROST. Run 2 reverses the arm order of run 1. Each run has 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls per arm. The shown value is the median of group medians. No cell is power-capped or unstable. Deterministic Cake is
cake_gdn_det_train_kcp(chunk_gdn_det_kcp) on the same FlashInfer tree, unchanged by this PR. Its packed views use its own helper,packed_qkv_views_det_kcp. Its outputs and gradients are bitwise equal to this PR's. Strongest non-Cake is the fastest arm per row that passes the accuracy check. The check needs rel-RMS <= 1e-2 and max-abs <= 2e-2 of max|ref| on every tensor, with finite values. Candidates are FLA 0.6.0 Tritonchunk_gated_delta_rule, cuDNN frontend 1.32.1gated_delta_netand FlashQLA 0.1.3. The cuDNN plan is gdn_frost (cuDNN FROST): ck0 = the backward recomputes the states; ck64 = the forward stores BF16 states every 64 tokens. On GB300, cuDNN FROST passes the check and is the strongest arm on every row: ck0 for fwd, ck64 for bwd and ckpt. FlashQLA ran in a separate process and was not faster than the shown arm in any row.R200 (sm_107a, 208 SMs)
NVIDIA R200 board, sm_107a, 208 SMs; these results cover the 208-SM board only. PyTorch 2.14.1 (cu130, CUDA 13.0 runtime), CUDA 13.4 toolkit, BF16. Workload: (H_qk, HV) = (16, 48), D = 128. CP1 = one rank, 32,768 tokens in five documents [12000, 9000, 6000, 3768, 2000], contiguous layout. Only one GPU was visible, so the 4-rank CP4 zigzag rows were not run on this board. Q/K/V forms: split =
row.split+view(plain views), packed =packed_qkv_views_kcp, contig = contiguous copies. Modes: fwd = forward only; bwd = autograd backward after an eager forward. ckpt =torch.utils.checkpoint(use_reentrant=False), 2 forwards + 1 backward. graph = the package'sgraph="replay"mode with a caller-held workspace. Timing: 2 runs x 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls. Run 2 reverses the arm order. The shown value is the median of group medians. The previous version is a byte copy of the package from FlashInfer main before this PR. It runs in the same process as this PR, with the same build and libraries. Speedup = previous ms / this PR ms. No R200 row is power-capped (GPU software power-cap share >= 5 %).CP1
Two-baseline comparison (CP1)
These are separate runs on a different R200 machine from the previous-version tables, so this PR's absolute times differ slightly between the two sets. Each run uses one process for deterministic Cake, this PR, FLA and cuDNN FROST. Run 2 reverses the arm order of run 1. Each run has 8 palindromic groups; each group has 2 blocks of 5 warm-up + 20 timed calls per arm. The shown value is the median of group medians. The norm-on rows come from a re-measure: two further runs with the same design. No cell is power-capped or unstable. Deterministic Cake is
cake_gdn_det_train_kcp(chunk_gdn_det_kcp) on the same FlashInfer tree, unchanged by this PR. Its packed views use its own helper,packed_qkv_views_det_kcp. Its outputs and gradients are bitwise equal to this PR's. Strongest non-Cake is the fastest arm per row that passes the accuracy check. The check needs rel-RMS <= 1e-2 and max-abs <= 2e-2 of max|ref| on every tensor, with finite values. Candidates are FLA 0.6.0 Tritonchunk_gated_delta_rule, cuDNN frontend 1.32.1gated_delta_netand FlashQLA 0.1.3. The cuDNN plan is gdn_frost (cuDNN FROST): ck0 = the backward recomputes the states; ck64 = the forward stores BF16 states every 64 tokens. On R200, cuDNN FROST passes the check and is the strongest arm on every row: ck0 for fwd, ck64 for bwd and ckpt. FlashQLA does not support sm_107a; it accepts only sm90, sm100, sm103, sm120 and sm121. So R200 has no FlashQLA value.This PR does not change the other Cake variants (the deterministic KCP package
cake_gdn_det_train_kcpand the single-rank packages); architecture coverage stays sm_100a, sm_103a and sm_107a.Summary by CodeRabbit