Skip to content

fix(cake_dense_projection_det_gemm): take the cluster rendezvous before the 2-CTA tensor-memory allocation, and add sm_107a row rules - #8303

Merged
yyihuang merged 1 commit into
flashinfer-ai:mainfrom
yyihuang:dpg-det-r14-rendezvous
Nov 2, 2026
Merged

yyihuang merged 1 commit into
flashinfer-ai:mainfrom
yyihuang:dpg-det-r14-rendezvous

Conversation

@yyihuang

@yyihuang yyihuang commented Nov 2, 2026 •

Copy link
Copy Markdown
Collaborator

The Cake dense_projection_det_gemm package in flashinfer.experimental is the deterministic dense projection GEMM. It has 127 programs built from 124 templates. Most of them are BF16-operand dense GEMMs that write BF16 or FP32 output. One is an FP32 router GEMM that runs through BF16x3 emulation. This PR fixes a hang in the programs that run as 2-CTA clusters and adds sm_107a row rules. It does not add or remove any program, template or registry record.

In 49 of the package's kernel sources, the programs run as 2-CTA clusters. Each CTA allocated tensor memory (tcgen05.alloc.cta_group::2) before the cluster barrier. A CTA could therefore allocate before its peer CTA had started, and the launch could wait forever. This happened on GB200, GB300 and Rubin in long-running processes. The cluster barrier now comes first. Machine code changes only in the program prologue, except for one of these programs on sm_103a, where one more instruction also moves. Outputs are bitwise identical to the previous version.

The PR also adds 34 sm_107a row rules for Rubin, measured on the 216-SM part. Of these, 17 apply at every token count T, 9 apply only at T = 16231 and 8 apply only at T = 16384. They route rows to programs that already ship. The total rule count goes from 307 to 341, and the sm_107a count goes from 90 to 124.

Speedup is the time of the previous Cake version (the package on main) divided by the time of this PR, on the same GPU with CUDA 13.4, over 294 shape rows. On GB200 the geometric mean is 0.9999, with rows from 0.992 to 1.015. On GB300 it is 0.9997, with rows from 0.993 to 1.006. No row on either GPU is below 0.99, and the fixed programs are within 0.8 % of the previous version. On Rubin the geometric mean is 1.025 over all 294 rows. It is 1.157 over the 51 rows that the new rules change and 0.9988 over the other 243 rows.

The largest Rubin gains are in FP32 weight gradients. 4096→512 goes from 74.3 to 33.0 µs at T = 16384 and from 73.9 to 32.9 µs at T = 16231, which is 2.25x in both cases. 5120→512 at T = 16231 goes from 74.2 to 40.5 µs (1.83x). 4096→4096 at T = 16384 goes from 235.3 to 157.6 µs (1.49x). 4096→1024 at T = 16231 goes from 74.9 to 53.0 µs (1.41x). Thin FP32 weight gradients from 4096 or 5120 to 32, 48 or 64 are 1.13x to 1.19x faster. Rows that use the new FP32 input-gradient rules are 1.05x to 1.09x faster.

The prologue fix has a small cost on Rubin. One row is slower in both measured passes: the 5120→5120 BF16 input gradient at T = 16231 goes from 225.4 to 229.0 µs (0.984). Eight more Rubin rows read 0.982 to 0.990 but are not confirmed: the slowdown shows in only one of the two passes or in only one of the two order-balanced halves.

The package ships BF16 and FP32 outputs in the kk, kn, nk and nn operand layouts, along with weight- and input-gradient forms that include the half-tile split forms. It also ships eight-warp CTA programs on sm_100a and sm_103a, and the FP32 router GEMM through BF16x3 emulation. It has programs for sm_100a (GB200), sm_103a (GB300) and sm_107a (Rubin): 256 build entries (83 / 86 / 87) and 127 registry records.

Times in µs on NVIDIA GB200 (sm_100a, 152 SMs), NVIDIA GB300 (sm_103a, 152 SMs) and Rubin (sm_107a, 216 SMs) with CUDA 13.4, cold L2 and order-balanced, with one row per shape at T = 16172, 16231 or 16384. Speedup = previous Cake version / this PR; † and ‡ are explained below the table.

shape (K→N) op out T GB200 previous µs GB200 this PR µs GB200 speedup GB300 previous µs GB300 this PR µs GB300 speedup Rubin previous µs Rubin this PR µs Rubin speedup
indexer_q (2048→4096) dgrad bf16 16172 173.8 173.8 1.000† 141.7 141.7 1.000† 79.6 79.7 0.999
indexer_q (2048→4096) dgrad bf16 16231 173.8 173.8 1.000† 141.8 141.8 1.000† 79.6 79.7 1.000
indexer_q (2048→4096) fwd bf16 16172 170.2 170.2 1.000† 138.8 138.8 1.000† 79.2 79.2 1.000
indexer_q (2048→4096) fwd bf16 16231 171.0 171.0 1.000† 139.6 139.6 1.000† 79.2 79.2 1.000
indexer_q (2048→4096) wgrad bf16 16172 179.2 179.2 1.000† 147.6 147.6 1.000† 79.1 79.1 1.000
indexer_q (2048→4096) wgrad bf16 16231 180.4 180.4 1.000† 148.3 148.3 1.000† 79.2 79.3 0.999
indexer_q (2048→4096) wgrad f32 16172 173.7 173.6 1.001 145.2 145.3 0.999 80.8 80.8 1.000
indexer_q (2048→4096) wgrad f32 16231 175.3 175.5 0.999 145.9 146.1 0.999 81.1 81.1 1.000
shared_down (2048→6144) dgrad bf16 16172 258.1 258.1 1.000† 216.4 216.4 1.000† 116.3 116.4 0.999
shared_down (2048→6144) dgrad bf16 16231 256.1 256.1 1.000† 215.8 215.8 1.000† 116.3 116.4 0.999
shared_down (2048→6144) dgrad f32 16172 261.6 262.5 0.997 221.5 221.5 1.000 118.1 118.1 1.000
shared_down (2048→6144) dgrad f32 16231 261.3 260.9 1.002 220.5 220.5 1.000 118.1 118.1 1.000
shared_down (2048→6144) fwd bf16 16172 249.3 250.4 0.996 212.0 212.2 0.999 111.2 111.3 1.000
shared_down (2048→6144) fwd bf16 16231 249.7 250.4 0.997 213.9 214.1 0.999 111.2 111.2 0.999
shared_down (2048→6144) wgrad bf16 16172 259.4 259.4 1.000† 218.4 218.4 1.000† 116.1 116.1 1.000
shared_down (2048→6144) wgrad bf16 16231 261.2 261.2 1.000† 218.9 218.9 1.000† 116.3 116.3 1.000
shared_down (2048→6144) wgrad f32 16172 262.1 262.1 1.000† 219.4 219.4 1.000† 116.9 117.1 0.999
shared_down (2048→6144) wgrad f32 16231 262.6 262.6 1.000† 220.2 220.2 1.000† 117.4 117.5 0.999
q_b (2048→16384) dgrad bf16 16172 696.1 695.2 1.001 601.1 601.9 0.999 316.3 317.0 0.998
q_b (2048→16384) dgrad bf16 16231 707.3 697.1 1.015 599.8 602.2 0.996 312.5 314.2 0.994
q_b (2048→16384) dgrad f32 16172 696.4 695.2 1.002 592.7 593.0 0.999 315.7 311.7 1.013
q_b (2048→16384) dgrad f32 16231 694.9 691.6 1.005 595.2 594.6 1.001 322.2 321.0 1.004
q_b (2048→16384) fwd bf16 16172 613.3 613.3 1.000† 549.8 549.8 1.000† 297.7 301.2 0.988
q_b (2048→16384) fwd bf16 16231 615.5 615.5 1.000† 548.5 548.5 1.000† 309.1 310.9 0.994
q_b (2048→16384) wgrad bf16 16172 673.1 677.9 0.993 586.6 588.6 0.997 304.2 298.0 1.021
q_b (2048→16384) wgrad bf16 16231 682.3 681.1 1.002 588.4 588.3 1.000 299.1 298.0 1.004
q_b (2048→16384) wgrad f32 16172 703.5 695.7 1.011 586.0 588.4 0.996 297.0 295.2 1.006
q_b (2048→16384) wgrad f32 16231 710.3 705.8 1.006 590.1 590.6 0.999 299.6 297.4 1.007
4096→32 dgrad bf16 16231 24.8 24.8 1.000† 23.2 23.2 1.000† 16.9 16.9 1.001
4096→32 dgrad bf16 16384 24.4 24.4 1.000† 23.3 23.3 1.000† 17.0 17.0 1.001
4096→32 dgrad f32 16231 44.8 44.8 1.000† 42.9 42.9 1.000† 30.2 28.3 1.067‡
4096→32 dgrad f32 16384 45.2 45.2 1.000† 43.4 43.4 1.000† 30.5 28.6 1.064‡
4096→32 fwd bf16 16231 25.9 25.9 1.000† 27.0 27.0 1.000† 19.7 19.5 1.012‡
4096→32 fwd bf16 16384 25.7 25.7 1.000† 26.6 26.6 1.000† 19.3 19.1 1.013‡
4096→32 wgrad bf16 16231 27.4 27.4 1.000† 27.4 27.4 1.000† 20.3 20.4 0.994
4096→32 wgrad bf16 16384 27.1 27.1 1.000† 27.6 27.6 1.000† 20.7 20.8 0.997
4096→32 wgrad f32 16231 27.4 27.4 1.000† 27.5 27.5 1.000† 24.1 20.6 1.173‡
4096→32 wgrad f32 16384 27.1 27.1 1.000† 27.7 27.7 1.000† 24.3 20.8 1.167‡
4096→48 dgrad bf16 16231 24.6 24.6 1.000† 22.9 22.9 1.000† 17.3 17.3 1.000
4096→48 dgrad bf16 16384 24.8 24.8 1.000† 23.3 23.3 1.000† 17.4 17.4 1.000
4096→48 dgrad f32 16231 45.2 45.2 1.000† 43.2 43.2 1.000† 30.4 28.5 1.068‡
4096→48 dgrad f32 16384 45.6 45.6 1.000† 43.6 43.6 1.000† 30.7 28.8 1.067‡
4096→48 fwd bf16 16231 26.1 26.1 1.000† 27.0 27.0 1.000† 19.9 19.6 1.014‡
4096→48 fwd bf16 16384 25.8 25.8 1.000† 26.8 26.8 1.000† 19.3 19.1 1.012‡
4096→48 wgrad bf16 16231 27.7 27.7 1.000† 27.9 27.9 1.000† 20.5 20.6 0.995
4096→48 wgrad bf16 16384 27.5 27.5 1.000† 28.1 28.1 1.000† 21.0 21.0 0.997
4096→48 wgrad f32 16231 27.7 27.7 1.000† 27.9 27.9 1.000† 24.5 20.8 1.179‡
4096→48 wgrad f32 16384 27.6 27.6 1.000† 28.0 28.0 1.000† 24.6 21.2 1.159‡
4096→64 dgrad bf16 16231 24.5 24.5 1.000† 23.6 23.6 1.000† 16.9 16.9 1.001
4096→64 dgrad bf16 16384 24.6 24.6 1.000† 23.7 23.7 1.000† 17.0 17.1 0.999
4096→64 dgrad f32 16231 45.3 45.3 1.000† 43.6 43.6 1.000† 30.4 28.5 1.068‡
4096→64 dgrad f32 16384 45.5 45.5 1.000† 43.6 43.6 1.000† 30.5 28.7 1.066‡
4096→64 fwd bf16 16231 26.2 26.2 1.000† 27.2 27.2 1.000† 20.0 19.6 1.019‡
4096→64 fwd bf16 16384 26.1 26.1 1.000† 27.0 27.0 1.000† 19.4 19.1 1.018‡
4096→64 wgrad bf16 16231 27.8 27.8 1.000† 28.0 28.0 1.000† 20.5 20.6 0.993
4096→64 wgrad bf16 16384 27.4 27.4 1.000† 27.8 27.8 1.000† 20.9 21.0 0.995
4096→64 wgrad f32 16231 27.7 27.7 1.000† 27.9 27.9 1.000† 24.7 20.6 1.195‡
4096→64 wgrad f32 16384 27.4 27.4 1.000† 28.0 28.0 1.000† 24.7 21.0 1.176‡
4096→128 dgrad bf16 16231 25.8 25.8 1.000† 24.4 24.4 1.000† 18.8 18.9 0.993
4096→128 dgrad bf16 16384 26.1 26.1 1.000† 24.5 24.5 1.000† 18.8 18.9 0.993
4096→128 dgrad f32 16231 47.6 47.6 1.000† 44.8 44.8 1.000† 32.0 29.4 1.090‡
4096→128 dgrad f32 16384 47.6 47.6 1.000† 44.9 44.9 1.000† 32.3 29.7 1.086‡
4096→128 fwd bf16 16231 27.2 27.2 1.000† 28.1 28.1 1.000† 20.4 19.8 1.027‡
4096→128 fwd bf16 16384 26.7 26.7 1.000† 27.5 27.5 1.000† 19.8 19.3 1.025‡
4096→128 wgrad bf16 16231 28.4 28.4 1.000† 28.6 28.6 1.000† 21.1 21.2 0.997
4096→128 wgrad bf16 16384 28.3 28.3 1.000† 28.6 28.6 1.000† 21.5 21.6 0.999
4096→128 wgrad f32 16231 28.5 28.5 1.000† 28.9 28.9 1.000† 21.0 21.2 0.991
4096→128 wgrad f32 16384 28.2 28.2 1.000† 28.5 28.5 1.000† 21.4 21.6 0.991
4096→512 dgrad bf16 16231 49.3 49.3 1.000† 41.2 41.2 1.000† 30.1 30.2 0.999
4096→512 dgrad bf16 16384 49.5 49.5 1.000† 41.0 41.0 1.000† 30.3 30.3 1.001
4096→512 dgrad f32 16231 68.7 68.9 0.998 59.8 59.9 0.998 39.6 39.7 0.997
4096→512 dgrad f32 16384 69.0 69.1 0.999 60.1 60.2 0.998 39.8 39.9 0.999
4096→512 fwd bf16 16231 53.2 53.1 1.000 47.7 47.8 0.999 33.9 33.9 1.000
4096→512 fwd bf16 16384 54.8 54.8 0.999 48.3 48.3 1.000 33.9 33.9 1.000
4096→512 wgrad bf16 16231 51.3 51.3 1.000† 43.4 43.4 1.000† 36.5 36.5 1.000
4096→512 wgrad bf16 16384 51.7 51.7 1.000† 43.3 43.3 1.000† 36.8 36.8 1.000
4096→512 wgrad f32 16231 51.3 51.3 1.000† 43.1 43.1 1.000† 73.9 32.9 2.246‡
4096→512 wgrad f32 16384 51.4 51.4 1.000† 43.4 43.4 1.000† 74.3 33.0 2.249‡
4096→1024 dgrad bf16 16231 93.3 93.3 1.000† 72.3 72.3 1.000† 45.4 45.5 0.999
4096→1024 dgrad bf16 16384 93.6 93.6 1.000† 72.4 72.4 1.000† 45.4 45.5 0.999
4096→1024 fwd bf16 16231 95.8 95.8 1.000† 75.2 75.2 1.000† 46.7 46.8 0.998
4096→1024 fwd bf16 16384 96.4 96.4 1.000† 75.1 75.1 1.000† 46.8 46.9 0.998
4096→1024 wgrad bf16 16231 96.3 96.3 1.000† 73.6 73.6 1.000† 53.5 53.5 1.000
4096→1024 wgrad bf16 16384 96.6 96.6 1.000† 74.5 74.5 1.000† 54.1 54.0 1.000
4096→1024 wgrad f32 16231 97.8 97.8 1.000† 75.8 75.8 1.000† 74.9 53.0 1.413‡
4096→1024 wgrad f32 16384 99.1 99.1 1.000† 76.7 76.7 1.000† 76.5 76.5 1.000
4096→4096 dgrad bf16 16231 316.2 316.2 1.000† 272.7 272.7 1.000† 150.1 150.2 1.000
4096→4096 dgrad bf16 16384 318.0 318.0 1.000† 270.3 270.3 1.000† 149.9 150.5 0.996
4096→4096 fwd bf16 16231 318.6 318.6 1.000† 270.9 270.9 1.000† 149.7 149.5 1.002
4096→4096 fwd bf16 16384 318.2 318.2 1.000† 271.1 271.1 1.000† 149.7 149.9 0.998
4096→4096 wgrad bf16 16231 329.3 329.3 1.000† 281.2 281.2 1.000† 176.2 154.6 1.140‡
4096→4096 wgrad bf16 16384 332.0 332.0 1.000† 283.6 283.6 1.000† 177.5 155.5 1.141‡
4096→4096 wgrad f32 16231 334.5 333.7 1.003 286.5 286.7 1.000 233.6 156.7 1.491‡
4096→4096 wgrad f32 16384 338.3 336.7 1.005 289.8 289.9 1.000 235.3 157.6 1.493‡
4096→12288 dgrad bf16 16231 928.6 928.6 1.000† 831.5 831.5 1.000† 498.0 497.3 1.001
4096→12288 dgrad bf16 16384 935.4 935.4 1.000† 838.7 838.7 1.000† 500.5 497.4 1.006
4096→12288 fwd bf16 16231 893.9 893.9 1.000† 809.1 809.1 1.000† 480.2 484.6 0.991
4096→12288 fwd bf16 16384 917.4 917.4 1.000† 812.1 812.1 1.000† 483.5 483.7 1.000
4096→12288 wgrad bf16 16231 921.2 921.2 1.000† 832.4 832.4 1.000† 498.1 499.6 0.997
4096→12288 wgrad bf16 16384 911.3 911.3 1.000† 842.6 842.6 1.000† 500.7 503.5 0.994
4096→12288 wgrad f32 16231 1056.5 1056.0 1.000 904.0 906.2 0.998 494.1 490.1 1.008
4096→12288 wgrad f32 16384 1054.9 1058.1 0.997 905.4 905.9 0.999 498.1 495.8 1.005
5120→32 dgrad bf16 16231 29.0 29.0 1.000† 27.5 27.5 1.000† 20.0 20.0 1.000
5120→32 dgrad bf16 16384 29.2 29.2 1.000† 27.6 27.6 1.000† 20.0 20.0 1.000
5120→32 dgrad f32 16231 53.6 53.6 1.000† 51.5 51.5 1.000† 36.4 34.2 1.063‡
5120→32 dgrad f32 16384 54.1 54.1 1.000† 52.0 52.0 1.000† 36.8 34.6 1.064‡
5120→32 fwd bf16 16231 30.9 30.9 1.000† 31.6 31.6 1.000† 23.2 23.3 0.997
5120→32 fwd bf16 16384 31.6 31.6 1.000† 32.7 32.7 1.000† 23.4 23.4 0.996
5120→32 wgrad bf16 16231 32.4 32.4 1.000† 34.1 34.1 1.000† 23.1 23.1 0.999
5120→32 wgrad bf16 16384 32.8 32.8 1.000† 34.3 34.3 1.000† 23.1 23.2 0.996
5120→32 wgrad f32 16231 32.2 32.2 1.000† 34.1 34.1 1.000† 26.0 22.7 1.144‡
5120→32 wgrad f32 16384 32.7 32.7 1.000† 34.2 34.2 1.000† 26.1 23.0 1.134‡
5120→48 dgrad bf16 16231 29.4 29.4 1.000† 27.7 27.7 1.000† 20.4 20.4 1.000
5120→48 dgrad bf16 16384 29.7 29.7 1.000† 27.9 27.9 1.000† 20.5 20.5 1.000
5120→48 dgrad f32 16231 54.2 54.2 1.000† 51.9 51.9 1.000† 37.0 34.5 1.070‡
5120→48 dgrad f32 16384 54.7 54.7 1.000† 52.3 52.3 1.000† 37.2 34.9 1.069‡
5120→48 fwd bf16 16231 30.9 30.9 1.000† 31.7 31.7 1.000† 23.3 23.3 0.997
5120→48 fwd bf16 16384 31.5 31.5 1.000† 32.6 32.6 1.000† 23.4 23.5 0.998
5120→48 wgrad bf16 16231 32.5 32.5 1.000† 34.5 34.5 1.000† 23.3 23.3 0.999
5120→48 wgrad bf16 16384 32.9 32.9 1.000† 34.7 34.7 1.000† 23.4 23.4 1.002
5120→48 wgrad f32 16231 32.5 32.5 1.000† 34.2 34.2 1.000† 26.3 23.1 1.139‡
5120→48 wgrad f32 16384 33.0 33.0 1.000† 34.8 34.8 1.000† 26.4 23.3 1.132‡
5120→64 dgrad bf16 16231 29.3 29.3 1.000† 27.9 27.9 1.000† 20.1 20.1 0.999
5120→64 dgrad bf16 16384 29.6 29.6 1.000† 28.0 28.0 1.000† 20.2 20.2 1.001
5120→64 dgrad f32 16231 54.7 54.7 1.000† 52.3 52.3 1.000† 36.8 34.3 1.073‡
5120→64 dgrad f32 16384 55.1 55.1 1.000† 52.7 52.7 1.000† 37.1 34.7 1.069‡
5120→64 fwd bf16 16231 31.4 31.4 1.000† 32.3 32.3 1.000† 23.4 23.4 0.997
5120→64 fwd bf16 16384 32.0 32.0 1.000† 33.1 33.1 1.000† 23.5 23.6 0.997
5120→64 wgrad bf16 16231 33.2 33.2 1.000† 34.5 34.5 1.000† 23.2 23.3 0.999
5120→64 wgrad bf16 16384 33.4 33.4 1.000† 34.9 34.9 1.000† 23.4 23.4 1.000
5120→64 wgrad f32 16231 33.1 33.1 1.000† 34.4 34.4 1.000† 26.4 22.9 1.155‡
5120→64 wgrad f32 16384 33.6 33.6 1.000† 34.7 34.7 1.000† 26.6 23.3 1.142‡
5120→128 dgrad bf16 16231 31.4 31.4 1.000† 29.2 29.2 1.000† 21.8 21.9 0.996
5120→128 dgrad bf16 16384 31.2 31.2 1.000† 29.2 29.2 1.000† 21.9 22.0 0.996
5120→128 dgrad f32 16231 57.1 57.1 1.000† 54.2 54.2 1.000† 39.1 35.9 1.090‡
5120→128 dgrad f32 16384 57.4 57.4 1.000† 54.3 54.3 1.000† 39.4 36.1 1.090‡
5120→128 fwd bf16 16231 32.7 32.7 1.000† 33.5 33.5 1.000† 23.8 23.2 1.024‡
5120→128 fwd bf16 16384 32.3 32.3 1.000† 33.2 33.2 1.000† 24.0 23.5 1.022‡
5120→128 wgrad bf16 16231 33.7 33.7 1.000† 35.6 35.6 1.000† 23.8 23.8 1.000
5120→128 wgrad bf16 16384 34.2 34.2 1.000† 36.1 36.1 1.000† 23.9 24.0 0.999
5120→128 wgrad f32 16231 33.6 33.6 1.000† 35.6 35.6 1.000† 23.8 23.9 0.996
5120→128 wgrad f32 16384 34.2 34.2 1.000† 35.9 35.9 1.000† 23.9 24.1 0.994
5120→512 dgrad bf16 16231 62.9 62.9 1.000† 49.7 49.7 1.000† 35.2 35.2 1.000
5120→512 dgrad bf16 16384 63.4 63.4 1.000† 50.1 50.1 1.000† 35.2 35.2 1.000
5120→512 dgrad f32 16231 86.3 86.4 0.999 73.9 73.9 1.000 47.3 47.5 0.997
5120→512 dgrad f32 16384 85.7 85.7 0.999 73.6 73.6 0.999 47.7 47.8 0.997
5120→512 fwd bf16 16231 66.5 66.5 1.001 57.7 57.8 0.999 39.8 39.8 1.000
5120→512 fwd bf16 16384 68.3 68.4 0.999 58.2 58.3 0.998 39.9 40.0 1.000
5120→512 wgrad bf16 16231 68.6 68.6 1.000† 53.2 53.2 1.000† 38.4 38.4 1.000
5120→512 wgrad bf16 16384 69.7 69.7 1.000† 54.0 54.0 1.000† 38.6 38.6 1.000
5120→512 wgrad f32 16231 70.0 70.0 1.000† 55.1 55.1 1.000† 74.2 40.5 1.831‡
5120→512 wgrad f32 16384 70.2 70.2 1.000† 55.0 55.0 1.000† 74.6 40.8 1.827‡
5120→1024 dgrad bf16 16231 113.5 113.5 1.000† 90.1 90.1 1.000† 53.5 53.5 0.999
5120→1024 dgrad bf16 16384 112.8 112.8 1.000† 89.9 89.9 1.000† 53.6 53.7 0.999
5120→1024 fwd bf16 16231 117.1 117.1 1.000† 93.4 93.4 1.000† 55.9 56.0 0.998
5120→1024 fwd bf16 16384 118.0 118.0 1.000† 93.6 93.6 1.000† 56.1 56.1 0.999
5120→1024 wgrad bf16 16231 122.2 122.2 1.000† 99.2 99.2 1.000† 64.1 64.2 0.999
5120→1024 wgrad bf16 16384 121.7 121.7 1.000† 99.8 99.8 1.000† 64.6 64.6 0.999
5120→1024 wgrad f32 16231 127.8 127.8 1.000† 100.3 100.3 1.000† 77.8 77.8 1.001
5120→1024 wgrad f32 16384 122.9 122.9 1.000† 101.2 101.2 1.000† 78.1 78.3 0.997
5120→5120 dgrad bf16 16231 486.2 486.2 1.000† 424.6 424.6 1.000† 225.4 229.0 0.984
5120→5120 dgrad bf16 16384 488.0 488.0 1.000† 425.3 425.3 1.000† 225.6 229.8 0.982
5120→5120 fwd bf16 16231 487.1 487.1 1.000† 424.5 424.5 1.000† 226.1 227.6 0.993
5120→5120 fwd bf16 16384 487.6 487.6 1.000† 425.4 425.4 1.000† 229.7 233.2 0.985
5120→5120 wgrad bf16 16231 489.3 489.3 1.000† 430.3 430.3 1.000† 232.7 234.8 0.991
5120→5120 wgrad bf16 16384 493.5 493.5 1.000† 435.5 435.5 1.000† 236.8 234.7 1.009
5120→5120 wgrad f32 16231 512.2 510.0 1.004 444.3 445.0 0.999 240.5 236.5 1.017‡
5120→5120 wgrad f32 16384 519.0 520.7 0.997 446.9 446.5 1.001 242.4 238.5 1.016‡
5120→17408 dgrad bf16 16231 1596.3 1596.3 1.000† 1458.5 1458.5 1.000† 906.9 905.9 1.001
5120→17408 dgrad bf16 16384 1599.5 1599.5 1.000† 1476.8 1476.8 1.000† 912.8 909.7 1.003
5120→17408 fwd bf16 16231 1579.8 1579.8 1.000† 1436.7 1436.7 1.000† 887.4 884.3 1.004
5120→17408 fwd bf16 16384 1589.6 1589.6 1.000† 1442.5 1442.5 1.000† 886.2 892.8 0.993
5120→17408 wgrad bf16 16231 1890.0 1882.7 1.004 1598.6 1593.0 1.004 891.2 889.7 1.002
5120→17408 wgrad bf16 16384 1924.0 1918.9 1.003 1604.4 1610.5 0.996 901.1 899.1 1.002
5120→17408 wgrad f32 16231 1915.4 1914.5 1.000 1575.5 1586.4 0.993 902.3 906.0 0.996
5120→17408 wgrad f32 16384 1938.0 1938.5 1.000 1597.4 1593.4 1.002 906.2 903.5 1.003
indexer_hw (6144→32) dgrad bf16 16172 34.1 34.2 0.996 32.7 32.8 0.998 22.8 22.8 1.000
indexer_hw (6144→32) dgrad bf16 16231 34.2 34.3 0.996 32.8 32.9 0.998 22.8 22.8 1.000
indexer_hw (6144→32) dgrad f32 16172 61.3 61.3 1.000† 58.4 58.4 1.000† 38.6 38.6 0.999
indexer_hw (6144→32) dgrad f32 16231 61.6 61.6 1.000† 58.5 58.5 1.000† 38.7 38.7 1.001
indexer_hw (6144→32) fwd bf16 16172 36.4 36.4 1.000† 37.5 37.5 1.000† 26.7 26.8 0.997
indexer_hw (6144→32) fwd bf16 16231 36.4 36.4 1.000† 37.5 37.5 1.000† 26.8 26.9 0.997
indexer_hw (6144→32) wgrad bf16 16172 37.6 37.8 0.996 39.4 39.5 0.997 27.6 27.7 0.994
indexer_hw (6144→32) wgrad bf16 16231 37.7 37.8 0.997 39.5 39.6 0.998 27.6 27.9 0.992
indexer_hw (6144→32) wgrad f32 16172 37.5 37.6 0.996 39.3 39.5 0.995 27.5 27.8 0.990
indexer_hw (6144→32) wgrad f32 16231 37.6 37.7 0.997 39.6 39.7 0.997 27.7 27.8 0.997
indexer_k (6144→128) dgrad bf16 16172 35.7 35.8 0.998 34.0 34.1 0.997 24.7 24.8 0.998
indexer_k (6144→128) dgrad bf16 16231 35.7 35.8 0.997 34.1 34.2 0.997 24.8 24.8 0.998
indexer_k (6144→128) dgrad f32 16172 64.4 64.5 0.998 61.9 62.0 0.998 41.8 41.8 1.000
indexer_k (6144→128) dgrad f32 16231 64.6 64.7 0.998 61.9 62.0 0.998 41.7 41.8 0.999
indexer_k (6144→128) fwd bf16 16172 37.9 37.9 1.000† 39.0 39.0 1.000† 27.3 27.4 0.996
indexer_k (6144→128) fwd bf16 16231 37.8 37.8 1.000† 38.8 38.8 1.000† 27.3 27.4 0.997
indexer_k (6144→128) wgrad bf16 16172 38.5 38.6 0.997 40.7 40.9 0.996 27.9 28.0 0.997
indexer_k (6144→128) wgrad bf16 16231 38.5 38.7 0.996 40.7 40.8 0.997 27.5 27.5 0.999
indexer_k (6144→128) wgrad f32 16172 38.2 38.4 0.995 40.7 40.8 0.997 27.9 28.1 0.995
indexer_k (6144→128) wgrad f32 16231 38.4 38.6 0.995 41.0 41.1 0.997 27.5 27.7 0.994
kv_a (6144→576) dgrad bf16 16172 82.5 82.5 1.000† 64.3 64.3 1.000† 42.7 42.7 0.998
kv_a (6144→576) dgrad bf16 16231 82.5 82.5 1.000† 64.4 64.4 1.000† 42.7 42.8 0.999
kv_a (6144→576) dgrad f32 16172 109.7 109.8 0.999 92.2 92.4 0.998 57.1 57.2 0.999
kv_a (6144→576) dgrad f32 16231 110.2 110.4 0.999 92.5 92.5 1.000 57.0 57.1 0.999
kv_a (6144→576) fwd bf16 16172 92.3 92.5 0.997 71.0 71.1 0.999 42.0 42.0 1.000
kv_a (6144→576) fwd bf16 16231 93.4 93.5 0.998 71.3 71.4 0.998 42.1 42.1 1.000
kv_a (6144→576) wgrad bf16 16172 82.1 82.2 0.999 64.0 64.1 0.998 44.1 44.1 1.000
kv_a (6144→576) wgrad bf16 16231 81.2 81.3 0.998 63.8 63.9 0.998 43.8 43.8 1.001
kv_a (6144→576) wgrad f32 16172 81.6 81.8 0.998 64.3 64.3 1.000 45.1 45.2 0.999
kv_a (6144→576) wgrad f32 16231 83.3 83.4 0.998 64.5 64.6 0.998 44.9 44.8 1.001
q_a (6144→2048) dgrad bf16 16172 249.8 249.8 1.000 211.3 211.1 1.001 111.1 111.2 1.000
q_a (6144→2048) dgrad bf16 16231 250.2 249.9 1.001 213.0 213.1 0.999 111.1 111.2 0.999
q_a (6144→2048) dgrad f32 16172 262.7 262.9 0.999 226.6 226.5 1.000 115.3 115.3 1.000
q_a (6144→2048) dgrad f32 16231 264.5 265.4 0.997 225.0 224.9 1.000 115.3 115.3 1.000
q_a (6144→2048) fwd bf16 16172 250.1 250.2 1.000 210.9 210.9 1.000 115.4 115.4 1.000
q_a (6144→2048) fwd bf16 16231 252.0 252.1 1.000 212.6 212.5 1.000 115.4 115.6 0.999
q_a (6144→2048) wgrad bf16 16172 260.3 260.3 1.000† 218.6 218.6 1.000† 115.6 115.9 0.998
q_a (6144→2048) wgrad bf16 16231 261.3 261.3 1.000† 218.8 218.8 1.000† 116.1 116.3 0.999
q_a (6144→2048) wgrad f32 16172 261.4 261.4 1.000† 219.5 219.5 1.000† 116.7 116.9 0.998
q_a (6144→2048) wgrad f32 16231 262.5 262.5 1.000† 220.0 220.0 1.000† 117.2 117.4 0.998
shared_gate_up (6144→2048) dgrad bf16 16172 250.8 249.9 1.003 214.3 214.3 1.000 111.1 111.2 0.999
shared_gate_up (6144→2048) dgrad bf16 16231 252.5 252.5 1.000 212.5 212.4 1.001 111.1 111.2 0.999
shared_gate_up (6144→2048) dgrad f32 16172 262.8 263.4 0.998 223.2 223.2 1.000 115.2 115.3 0.999
shared_gate_up (6144→2048) dgrad f32 16231 263.0 262.9 1.000 224.9 224.9 1.000 115.3 115.2 1.000
shared_gate_up (6144→2048) fwd bf16 16172 252.0 251.8 1.000 214.1 213.9 1.001 115.3 115.3 1.000
shared_gate_up (6144→2048) fwd bf16 16231 252.3 253.2 0.996 212.7 212.8 1.000 115.5 115.6 0.999
shared_gate_up (6144→2048) wgrad bf16 16172 260.4 260.4 1.000† 218.8 218.8 1.000† 115.6 115.8 0.998
shared_gate_up (6144→2048) wgrad bf16 16231 261.0 261.0 1.000† 219.4 219.4 1.000† 116.1 116.3 0.998
shared_gate_up (6144→2048) wgrad f32 16172 261.3 261.3 1.000† 219.4 219.4 1.000† 116.7 116.9 0.998
shared_gate_up (6144→2048) wgrad f32 16231 262.0 262.0 1.000† 220.5 220.5 1.000† 117.2 117.4 0.998
dense_gate_up (6144→12288) wgrad bf16 16172 1345.0 1345.0 1.000† 1222.3 1222.3 1.000† 728.2 728.3 1.000
dense_gate_up (6144→12288) wgrad bf16 16231 1319.5 1319.5 1.000† 1226.5 1226.5 1.000† 717.0 720.3 0.996
dense_gate_up (6144→12288) wgrad f32 16172 1586.7 1585.3 1.001 1327.6 1329.4 0.999 728.4 732.7 0.994
dense_gate_up (6144→12288) wgrad f32 16231 1608.1 1598.9 1.006 1329.8 1331.9 0.998 726.9 725.5 1.002
8192→128 dgrad bf16 16231 46.2 46.3 0.998 43.7 43.8 0.998 31.2 31.3 0.996
8192→128 dgrad bf16 16384 46.6 46.8 0.997 43.8 43.8 0.999 31.4 31.4 0.998
8192→128 dgrad f32 16231 91.4 91.4 1.000† 85.4 85.4 1.000† 58.4 55.2 1.059‡
8192→128 dgrad f32 16384 90.5 90.5 1.000† 84.7 84.7 1.000† 59.0 56.3 1.048‡
8192→128 fwd bf16 16231 47.9 47.9 1.000† 48.7 48.7 1.000† 34.8 34.9 0.999
8192→128 fwd bf16 16384 48.5 48.5 1.000† 49.3 49.3 1.000† 35.0 35.1 0.999
8192→128 wgrad bf16 16231 44.6 44.6 1.000† 45.6 45.6 1.000† 34.9 34.8 1.003
8192→128 wgrad bf16 16384 44.7 44.7 1.000† 45.9 45.9 1.000† 35.3 35.1 1.005
8192→128 wgrad f32 16231 44.5 44.5 1.000† 45.6 45.6 1.000† 34.9 34.8 1.004
8192→128 wgrad f32 16384 44.7 44.7 1.000† 46.0 46.0 1.000† 35.2 35.1 1.005
12288→4096 dgrad bf16 16231 895.7 895.7 1.000† 816.5 816.5 1.000† 482.2 488.6 0.987
12288→4096 dgrad bf16 16384 898.1 898.1 1.000† 816.5 816.5 1.000† 483.1 484.2 0.998
12288→4096 fwd bf16 16231 939.3 939.3 1.000† 830.1 830.1 1.000† 495.0 495.1 1.000
12288→4096 fwd bf16 16384 924.7 924.7 1.000† 827.4 827.4 1.000† 499.8 495.4 1.009
12288→4096 wgrad bf16 16231 915.0 915.0 1.000† 823.2 823.2 1.000† 490.4 496.5 0.988
12288→4096 wgrad bf16 16384 901.1 901.1 1.000† 826.3 826.3 1.000† 498.1 495.1 1.006
12288→4096 wgrad f32 16231 1048.9 1051.3 0.998 882.8 884.9 0.998 490.9 489.8 1.002
12288→4096 wgrad f32 16384 1054.1 1054.4 1.000 886.9 886.9 1.000 494.7 495.3 0.999
dense_down (12288→6144) wgrad bf16 16172 1335.5 1335.5 1.000† 1210.0 1210.0 1.000† 716.5 719.7 0.996
dense_down (12288→6144) wgrad bf16 16231 1344.7 1344.7 1.000† 1217.6 1217.6 1.000† 720.2 720.0 1.000
dense_down (12288→6144) wgrad f32 16172 1566.4 1573.0 0.996 1328.6 1335.2 0.995 727.1 722.8 1.006
dense_down (12288→6144) wgrad f32 16231 1572.2 1584.4 0.992 1341.9 1338.9 1.002 728.7 731.5 0.996
o_proj (16384→6144) dgrad bf16 16172 2195.7 2204.7 0.996 1833.9 1823.8 1.006 1016.5 1010.0 1.006
o_proj (16384→6144) dgrad bf16 16231 2215.1 2210.9 1.002 1838.9 1829.8 1.005 1005.7 1017.7 0.988
o_proj (16384→6144) dgrad f32 16172 2273.1 2281.7 0.996 1865.6 1865.3 1.000 1025.2 1025.4 1.000
o_proj (16384→6144) dgrad f32 16231 2276.4 2268.5 1.003 1852.6 1852.6 1.000 1019.7 1023.7 0.996
o_proj (16384→6144) fwd bf16 16172 2189.2 2179.0 1.005 1841.4 1837.3 1.002 982.9 980.1 1.003
o_proj (16384→6144) fwd bf16 16231 2194.0 2199.6 0.997 1842.6 1840.8 1.001 987.0 987.1 1.000
o_proj (16384→6144) wgrad bf16 16172 2158.0 2141.5 1.008 1798.7 1791.5 1.004 969.1 970.4 0.999
o_proj (16384→6144) wgrad bf16 16231 2156.5 2145.4 1.005 1781.8 1783.8 0.999 976.2 986.1 0.990
o_proj (16384→6144) wgrad f32 16172 2155.7 2137.7 1.008 1809.9 1805.0 1.003 1029.8 1021.2 1.008
o_proj (16384→6144) wgrad f32 16231 2170.1 2174.8 0.998 1824.4 1819.9 1.002 1023.6 1025.8 0.998
17408→5120 dgrad bf16 16231 1601.9 1601.9 1.000† 1446.6 1446.6 1.000† 880.2 889.2 0.990
17408→5120 dgrad bf16 16384 1567.0 1567.0 1.000† 1449.8 1449.8 1.000† 888.6 886.7 1.002
17408→5120 fwd bf16 16231 1593.2 1593.2 1.000† 1455.2 1455.2 1.000† 897.0 902.6 0.994
17408→5120 fwd bf16 16384 1574.6 1574.6 1.000† 1453.7 1453.7 1.000† 903.6 896.7 1.008
17408→5120 wgrad bf16 16231 1875.8 1864.2 1.006 1585.0 1574.9 1.006 888.8 884.3 1.005
17408→5120 wgrad bf16 16384 1900.3 1895.7 1.002 1586.7 1595.5 0.994 893.8 896.5 0.997
17408→5120 wgrad f32 16231 1900.7 1907.9 0.996 1570.7 1576.0 0.997 892.2 893.9 0.998
17408→5120 wgrad f32 16384 1929.8 1926.2 1.002 1585.6 1583.8 1.001 905.0 907.3 0.997
qabs (None→None) dgrad bf16 16172 235.8 236.0 0.999 227.1 227.2 1.000 153.2 153.3 1.000
qabs (None→None) dgrad bf16 16231 237.5 237.5 1.000 227.3 227.3 1.000 153.6 153.6 1.000
qabs (None→None) fwd bf16 16172 265.0 265.1 1.000 238.0 237.7 1.001 158.4 158.5 1.000
qabs (None→None) fwd bf16 16231 266.9 267.0 1.000 240.3 240.2 1.000 159.1 159.1 1.000
qabs (None→None) wgrad bf16 16172 216.0 215.9 1.000 218.1 218.1 1.000 145.0 144.9 1.000
qabs (None→None) wgrad bf16 16231 217.4 217.4 1.000 218.2 218.2 1.000 145.3 145.2 1.001
vproj (None→None) dgrad bf16 16172 286.2 286.1 1.000 262.6 262.6 1.000 172.7 172.7 1.000
vproj (None→None) dgrad bf16 16231 286.8 286.8 1.000 262.6 262.7 1.000 172.9 173.0 1.000
vproj (None→None) fwd bf16 16172 264.9 264.9 1.000 250.4 250.5 1.000 164.9 164.9 1.000
vproj (None→None) fwd bf16 16231 266.8 267.0 0.999 251.8 251.8 1.000 165.5 165.5 1.000
vproj (None→None) wgrad bf16 16172 236.0 236.0 1.000 236.2 236.2 1.000 158.4 158.2 1.001
vproj (None→None) wgrad bf16 16231 236.5 236.3 1.001 237.7 237.7 1.000 158.7 158.7 1.000
router 4096→256 dgrad f32 16231 187.9 187.9 1.000† 173.1 173.1 1.000† 107.2 107.2 1.001
router 4096→256 fwd f32 16231 192.5 192.5 1.000† 166.5 166.5 1.000† 102.7 102.6 1.000
router 4096→256 wgrad f32 16231 187.9 187.9 1.000† 163.6 163.6 1.000† 98.1 98.1 1.000
router 5120→256 dgrad f32 16231 235.6 235.6 1.000† 225.3 225.3 1.000† 131.7 131.7 1.000
router 5120→256 fwd f32 16231 226.1 226.1 1.000† 195.9 195.9 1.000† 121.0 121.0 1.000
router 5120→256 wgrad f32 16231 235.6 235.6 1.000† 209.4 209.4 1.000† 135.8 135.8 1.000
router 6144→256 dgrad f32 16172 270.2 270.2 1.000† 260.4 260.4 1.000† 151.5 151.5 1.000
router 6144→256 dgrad f32 16231 272.6 272.6 1.000† 262.8 262.8 1.000† 152.0 152.1 1.000
router 6144→256 fwd f32 16172 260.9 260.9 1.000† 228.1 228.1 1.000† 139.7 139.7 1.000
router 6144→256 fwd f32 16231 259.4 259.4 1.000† 226.4 226.4 1.000† 138.9 138.9 1.000
router 6144→256 wgrad f32 16172 248.7 248.7 1.000† 216.1 216.1 1.000† 137.4 137.3 1.000
router 6144→256 wgrad f32 16231 248.9 248.9 1.000† 216.5 216.5 1.000† 137.5 137.4 1.000

Footnotes: † identical code on this GPU (program and plan unchanged): previous µs = this PR µs from the previous-version record, speedup 1.000 by construction. ‡ Rubin rule rows: the new program was measured against the previous version before the prologue (rendezvous-first) change; the shipped program is the same template with the reordered prologue. Rubin = 216-SM sm_107a boards; GB200 = sm_100a, GB300 = sm_103a, 152 SMs; CUDA 13.4.

…row rules

In 49 CTA-pair kernel units the 2-CTA tensor memory allocation
(tcgen05.alloc.cta_group::2) ran before the cluster barrier, so the
allocation handshake could reach a peer CTA that had not started yet and
the kernel could hang. These units now pass the cluster barrier first and
allocate after it.

Add 34 sm_107a ROW_RULES entries: 17 ragged rules (nine FP32 input
gradients on the single-CTA program, five BF16 forwards, the 4096 x 4096
BF16 and FP32 and 5120 x 5120 FP32 weight gradients) and 17 T-exact FP32
weight-gradient rules at T 16231 and 16384. Every planned output stays
bitwise equal to the previous plan's.

Update the rule-count test and keep the exact-wave gate test on ragged keys.
@coderabbitai

coderabbitai Bot commented Nov 2, 2026 •

Copy link
Copy Markdown
Contributor

Warning

Review limit reached

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Next included review available in 2 minutes.

Check out review usage here.

View limit details

Limit details: You’ve used all 8 included reviews currently available.

Learn how review limits work.

Review configuration:

Run configuration
  • Source mode: Auto
  • Repository: Unavailable
  • Branch: Unavailable
  • File: No configuration file from this repository applied
  • Commit: Unavailable
  • Layers: defaults
  • Review profile: CHILL
  • Plan: Advanced
  • Run ID: 27e46e21-eae9-4ea5-b8f3-409ebafe22a5
Commits

Reviewing files that changed from the base of the PR and between b9417a2 and f7b571f.

Files selected for processing: 53
  • flashinfer/experimental/dense_projection_det_gemm/README.md
  • flashinfer/experimental/dense_projection_det_gemm/build.json
  • flashinfer/experimental/dense_projection_det_gemm/cake_backend.py
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0618ae0aec1af120afd0_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_087ba4f157914949c26a_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0a71a49da71c7d34ab37_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_0db65f3ddd780ed42d4d_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_13d9310f37868a18eace_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_14192244441c1003db74_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_295aef2f2e0332de2372_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_3088b8056e0b7f9ca32c_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_31be7faa868eabcedada_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_438841f42a0743b63fb5_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_446497e750658dfb057b_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_505fbe05ce8b62f0a589_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5226b6c854d7db395d7e_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5445e0f16ad84868fa01_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5568657a526e2040f0bc_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_56be8af43d5d0822da87_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5a06c5e2f33772a33174_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_5e77ef909e88a99c4357_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_66df6446dd2cca3106d3_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_765ece06955b70a95b29_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_7d403fa7b167188ccc24_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_806927d101fe72fc1317_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_82f698f67a1cca6e4291_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_895f9ba7159bea07aedd_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_8e1f6805342f620cca8c_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9aa9aa60b0c544bfab06_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_9cc6363566760a38b4fe_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b26ac3f3ed56a78cb052_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_b58b50619ef5ed46e109_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ca9a846738b1443f7318_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d2a5e70bd49083888179_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_d6ebaab2415ea7b00de9_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_ddff2b911137449c1dfc_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_de7492eeb6e24b31687a_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f2e6d4e79cf21105fa1e_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_f31e1dcb1df00973dc0c_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_fb4bcdf13915dfb96c59_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n128_pro7__hen_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_m256_cg_or_pro7__ht_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kk_n256_pro7__cg_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_bg8_pro7__sol_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_kn_n256_tma1_pro7__f32_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_hen_t_skx_pro7__f32_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_pro7__f32_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_m64_hen_t_skx_pro15__f32_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n128_t_pro7__f32_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n192_hee_t_tma1_pro7__f32_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_f32_tma1_pro7__ht_kernel.cu
  • flashinfer/experimental/dense_projection_det_gemm/csrc/cake_dense_projection_det_gemm/cake_dense_projection_det_gemm_tpl_nn_n256_m256_pro7__ht_kernel.cu
  • tests/experimental/test_cake_dense_projection_det_gemm.py
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

@yyihuang

yyihuang commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/experimental/test_cake_dense_projection_det_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py

@yzh119 yzh119 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@yyihuang
yyihuang merged commit b8b72ac into flashinfer-ai:main Nov 2, 2026
30 of 31 checks passed
yyihuang added a commit that referenced this pull request Nov 2, 2026
…e sm_107a token counts its new plans win (#8363)

## Summary

The Cake dense projection GEMM family covers BF16-input dense and
batched projection GEMMs, with BF16 or FP32 weight-gradient output. It
also covers an FP32 router GEMM (256 experts, BF16x3 emulation). In
`flashinfer.experimental.dense_projection_gemm`, a keyword-only
`deterministic` flag selects the variant. `deterministic=True` (the
default) calls `flashinfer.experimental.dense_projection_det_gemm`.
`deterministic=False` runs the split-K / stream-K programs (#8190,
#8238). The family ships on `sm_100a` (GB200), `sm_103a` (GB300) and
`sm_107a` (R200).

This PR changes only the planner. It affects only `deterministic=False`,
the projection weight gradient (K1), FP32 output rows, and `sm_107a`
parts with 208 SMs, 212 SMs, or 216 or more SMs. #8303 gives the
deterministic package new `sm_107a` plans for single token counts (16231
and 16384 tokens). At some of these points, the deterministic program is
now as fast as or faster than the split-K program. This PR runs the
deterministic program at those points.

The split-K program keeps a point only if both conditions hold in eager
and graph mode:

- It is faster than the deterministic program in every run at the op
level.
- The whole backward and checkpointed step are at most 1.3 % slower.

Otherwise, the point runs the deterministic program. The measured token
counts are 16172, 16231 and 16384. 20 measured points change (208 SMs:
10, 212 SMs: 10). The other 46 measured points keep their split-K plans.
Other token counts, other SM counts, `sm_100a`, `sm_103a` and all
kernels are unchanged.

Changed points (`d_in -> d_out`, FP32 output):

- 208 SMs: `4096 -> 48`, `4096 -> 64`, `4096 -> 512`, `5120 -> 48` at
16231 and 16384 tokens; `4096 -> 1024`, `5120 -> 64` at 16231.
- 212 SMs: `4096 -> 512`, `5120 -> 512` at 16231 and 16384; `4096 ->
1024`, `5120 -> 48`, `5120 -> 64` at 16231; `4096 -> 4096` at every
token count from 16172 to 16384 (measured at 16172, 16231 and 16384).
- 216 or more SMs: `4096 -> 48`, `4096 -> 512` at 16231 and 16384; `4096
-> 1024`, `5120 -> 48`, `5120 -> 64` at 16231.

This PR depends on #8303 and must merge after it. Before #8303, these
points run the earlier deterministic plans, which are slower. The "det
Cake (main)" column of the table shows them.

The entries for 216 or more SMs were not measured on such parts. They
are not in the table.

**Speedups** (graph mode, geomean over the 20 rows of each op, range in
parentheses; ratio = baseline us / new us, above 1 means this PR is
faster):

- wgrad: previous/new 1.031 (0.951 to 1.148); det (main)/new 1.457
(1.122 to 2.213); SOTA/new 1.028 (0.864 to 1.143).
- bwd: previous/new 1.022 (0.976 to 1.103); det (main)/new 1.264 (1.053
to 1.701); SOTA/new 0.971 (0.854 to 1.027).
- 2fwd+bwd: previous/new 1.018 (0.986 to 1.104); det (main)/new 1.146
(1.031 to 1.355); SOTA/new 0.957 (0.863 to 1.012).

In eager mode, the previous/new geomeans are 1.024 (wgrad), 1.023 (bwd)
and 1.018 (2fwd+bwd). The largest graph gains (previous/new) are at
212-SM `4096 -> 4096`: wgrad 1.136 to 1.148, bwd 1.097 to 1.103. 208-SM
`4096 -> 1024` at 16231 gives wgrad 1.139 and 2fwd+bwd 1.104. `4096 ->
512` gives wgrad 1.094 and 1.100 at 208 SMs, and 1.053 and 1.059 at 212
SMs.

In graph mode, this PR is slower than the previous version in 9 of 20
wgrad rows, 11 of 20 bwd rows and 5 of 20 2fwd+bwd rows. The lowest are
212-SM `5120 -> 512` at 16231 and 16384 (wgrad 0.951 and 0.955). There,
the split-K program was faster on most GPUs, but on one of the eight
GPUs the deterministic program was as fast or faster. 212-SM `4096 ->
1024` at 16231 gives wgrad 0.970 and bwd 0.976, but 2fwd+bwd 1.021,
because the checkpointed step of the split-K program is slower. This
point was measured twice; the rule uses the runs of both measurements
together. 208-SM and 212-SM `5120 -> 48`, `5120 -> 64` at 16231 (and
208-SM `5120 -> 48` at 16384) give wgrad 0.976 to 0.986. At these
points, the split-K program was not faster in every run, or its
checkpointed step was more than 1.3 % slower. The rule therefore selects
the deterministic program.

In graph mode, the library baseline is faster than the new plan
(SOTA/new below 1) in 6 of 20 wgrad rows, 13 of 20 bwd rows and 17 of 20
2fwd+bwd rows.

- **Rows:** 120 = 20 points x 3 ops x 2 modes, graph rows first. 208
SMs: 10 points, 212 SMs: 10 points.
- **Ops:** wgrad = weight gradient `[d_out, d_in]` (FP32 output) from
BF16 activations `[T, d_in]` and BF16 output gradients `[T, d_out]`; bwd
= input gradient (BF16) + wgrad; 2fwd+bwd = checkpointed step (2 x
forward + bwd). **Modes:** g = CUDA graph replay, e = eager.
- **Latency:** median us over 4 to 24 runs per row. One run = one GPU
and one process layout (one long-lived process over all rows, or a fresh
process per row). 208 SMs: 2 GPUs, 2 runs each; 212 SMs: 8 GPUs, 1 run
each (`4096 -> 1024` at 16231 tokens: 24 runs on 16 GPUs, from two
measurements). In a run, all arms of a row are interleaved in paired
groups: 15 groups of 20 steps for wgrad, 9 groups for bwd and 2fwd+bwd.
L2 is cold before every iteration. The time is the CUPTI GPU time of the
operation (first kernel begin to last kernel end, all launches).
- **Baselines:** det Cake (main) = the deterministic package with the
plans of current main for every op (forward, input gradient, wgrad),
built from the #8303 source. det Cake (#8303) = `deterministic=True`
with #8303. previous Cake = `deterministic=False` with #8303 and without
this PR (its forward and input gradient are the #8303 deterministic
programs). new = `deterministic=False` with this PR: it runs the
deterministic program at every row of the table, so new us is the det
Cake (#8303) measurement (the `deterministic=False` call that dispatches
to it was not timed separately). SOTA = cuBLAS / cuBLASLt, the vendor
GEMM library for `sm_107a`: for wgrad the fastest correct launch among
`torch.mm` and the cuBLASLt heuristic's top-8 algorithms; for bwd and
2fwd+bwd the fastest correct library launch per op. Ratio = baseline us
/ new us; above 1.000 = this PR is faster.
- **Setup:** R200 (`sm_107a`), 212 SMs (driver 620.62) and 208 SMs
(driver 615.71.10); CUDA 13.4 (nvcc 13.4.92), PyTorch 2.14.1+cu130,
cuBLAS 13.8.0.4.
- **Noise:** per row, the interquartile range of the per-run paired
previous/new ratio has a median of 0.0078, a 90th percentile of 0.0257
and a maximum of 0.0388. Ratios within about 1 % of 1.000 are inside
this spread. Some samples of the larger 212-SM rows show the software
power-cap throttle flag; all arms of a row run interleaved in one
process, so the ratios hold, but the absolute times of those rows can be
a few percent high. For 212-SM `4096 -> 1024` at 16231 tokens, the
per-run paired 2fwd+bwd previous/new is 0.983 to 1.041 (graph) and 0.986
to 1.043 (eager). An identical-plan control in the same runs, det Cake
(#8303) against det Cake (main) on the same row at 16172 and 16384
tokens (both use the same wgrad plan there), gives a per-run 2fwd+bwd
ratio of 0.955 to 1.016 (graph) and 0.950 to 1.022 (eager).

*Table: one row per point, op and mode; R200 (`sm_107a`), CUDA 13.4;
latencies in us; ratios = baseline us / new us.*

| GPU | SMs | `d_in -> d_out` | out | T | op | mode | det Cake (main) us
| det Cake (#8303) us | previous Cake us | SOTA us | new us | det
(main)/new | previous/new | SOTA/new |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|R200|212|4096 ->
512|FP32|16231|wgrad|g|74.70|33.87|35.65|37.90|33.87|2.205|1.053|1.119|
|R200|212|4096 ->
512|FP32|16231|bwd|g|105.07|61.89|65.20|62.52|61.89|1.698|1.053|1.010|
|R200|212|4096 ->
512|FP32|16231|2fwd+bwd|g|165.51|122.30|125.75|112.10|122.30|1.353|1.028|0.917|
|R200|212|4096 ->
512|FP32|16384|wgrad|g|75.18|33.97|35.96|38.33|33.97|2.213|1.059|1.128|
|R200|212|4096 ->
512|FP32|16384|bwd|g|105.51|62.02|65.57|63.59|62.02|1.701|1.057|1.025|
|R200|212|4096 ->
512|FP32|16384|2fwd+bwd|g|166.02|122.55|126.07|112.86|122.55|1.355|1.029|0.921|
|R200|212|4096 ->
1024|FP32|16231|wgrad|g|76.15|54.23|52.59|49.04|54.23|1.404|0.970|0.904|
|R200|212|4096 ->
1024|FP32|16231|bwd|g|132.31|108.78|106.15|95.45|108.78|1.216|0.976|0.877|
|R200|212|4096 ->
1024|FP32|16231|2fwd+bwd|g|228.53|196.41|200.59|192.35|196.41|1.164|1.021|0.979|
|R200|212|4096 ->
4096|FP32|16172|wgrad|g|234.22|160.15|181.96|157.34|160.15|1.463|1.136|0.982|
|R200|212|4096 ->
4096|FP32|16172|bwd|g|461.66|358.15|395.05|367.92|358.15|1.289|1.103|1.027|
|R200|212|4096 ->
4096|FP32|16172|2fwd+bwd|g|841.13|737.75|777.00|745.66|737.75|1.140|1.053|1.011|
|R200|212|4096 ->
4096|FP32|16231|wgrad|g|234.32|159.46|183.05|161.47|159.46|1.469|1.148|1.013|
|R200|212|4096 ->
4096|FP32|16231|bwd|g|463.82|357.92|393.53|363.81|357.92|1.296|1.099|1.016|
|R200|212|4096 ->
4096|FP32|16231|2fwd+bwd|g|852.91|746.09|778.15|752.00|746.09|1.143|1.043|1.008|
|R200|212|4096 ->
4096|FP32|16384|wgrad|g|236.40|161.10|184.37|160.84|161.10|1.467|1.144|0.998|
|R200|212|4096 ->
4096|FP32|16384|bwd|g|470.51|363.01|398.13|370.78|363.01|1.296|1.097|1.021|
|R200|212|4096 ->
4096|FP32|16384|2fwd+bwd|g|853.68|742.19|776.27|750.99|742.19|1.150|1.046|1.012|
|R200|212|5120 ->
48|FP32|16231|wgrad|g|23.54|20.01|19.53|20.03|20.01|1.176|0.976|1.001|
|R200|212|5120 ->
48|FP32|16231|bwd|g|44.14|40.66|40.22|38.60|40.66|1.086|0.989|0.949|
|R200|212|5120 ->
48|FP32|16231|2fwd+bwd|g|81.09|77.94|77.06|74.44|77.94|1.040|0.989|0.955|
|R200|212|5120 ->
64|FP32|16231|wgrad|g|23.45|19.47|19.04|20.40|19.47|1.204|0.978|1.048|
|R200|212|5120 ->
64|FP32|16231|bwd|g|43.51|39.65|39.15|38.35|39.65|1.097|0.987|0.967|
|R200|212|5120 ->
64|FP32|16231|2fwd+bwd|g|80.42|76.99|75.88|74.36|76.99|1.045|0.986|0.966|
|R200|212|5120 ->
512|FP32|16231|wgrad|g|74.96|40.12|38.15|34.67|40.12|1.868|0.951|0.864|
|R200|212|5120 ->
512|FP32|16231|bwd|g|110.91|73.54|73.13|65.91|73.54|1.508|0.994|0.896|
|R200|212|5120 ->
512|FP32|16231|2fwd+bwd|g|183.75|146.07|146.39|126.02|146.07|1.258|1.002|0.863|
|R200|212|5120 ->
512|FP32|16384|wgrad|g|75.31|40.23|38.42|35.00|40.23|1.872|0.955|0.870|
|R200|212|5120 ->
512|FP32|16384|bwd|g|111.10|73.98|73.39|66.87|73.98|1.502|0.992|0.904|
|R200|212|5120 ->
512|FP32|16384|2fwd+bwd|g|184.18|146.41|146.96|127.19|146.41|1.258|1.004|0.869|
|R200|208|4096 ->
48|FP32|16231|wgrad|g|24.95|21.66|21.66|23.14|21.66|1.152|1.000|1.068|
|R200|208|4096 ->
48|FP32|16231|bwd|g|43.45|40.26|40.25|39.11|40.26|1.079|1.000|0.971|
|R200|208|4096 ->
48|FP32|16231|2fwd+bwd|g|83.43|80.06|80.46|78.16|80.06|1.042|1.005|0.976|
|R200|208|4096 ->
48|FP32|16384|wgrad|g|25.02|21.66|21.86|22.99|21.66|1.155|1.009|1.061|
|R200|208|4096 ->
48|FP32|16384|bwd|g|43.57|40.46|40.42|39.25|40.46|1.077|0.999|0.970|
|R200|208|4096 ->
48|FP32|16384|2fwd+bwd|g|83.90|80.42|80.69|78.59|80.42|1.043|1.003|0.977|
|R200|208|4096 ->
64|FP32|16231|wgrad|g|25.14|21.35|21.23|23.78|21.35|1.178|0.994|1.114|
|R200|208|4096 ->
64|FP32|16231|bwd|g|43.30|39.59|39.22|39.46|39.59|1.094|0.991|0.997|
|R200|208|4096 ->
64|FP32|16231|2fwd+bwd|g|83.44|79.14|79.15|78.47|79.14|1.054|1.000|0.992|
|R200|208|4096 ->
64|FP32|16384|wgrad|g|25.33|21.49|21.50|23.85|21.49|1.179|1.000|1.110|
|R200|208|4096 ->
64|FP32|16384|bwd|g|43.36|39.98|39.29|39.54|39.98|1.085|0.983|0.989|
|R200|208|4096 ->
64|FP32|16384|2fwd+bwd|g|83.79|79.68|79.66|78.97|79.68|1.052|1.000|0.991|
|R200|208|4096 ->
512|FP32|16231|wgrad|g|74.47|34.34|37.56|38.86|34.34|2.169|1.094|1.132|
|R200|208|4096 ->
512|FP32|16231|bwd|g|106.39|64.81|68.71|66.10|64.81|1.642|1.060|1.020|
|R200|208|4096 ->
512|FP32|16231|2fwd+bwd|g|172.58|130.20|134.53|117.42|130.20|1.325|1.033|0.902|
|R200|208|4096 ->
512|FP32|16384|wgrad|g|74.83|34.36|37.81|39.29|34.36|2.178|1.100|1.143|
|R200|208|4096 ->
512|FP32|16384|bwd|g|106.70|64.96|68.93|66.31|64.96|1.643|1.061|1.021|
|R200|208|4096 ->
512|FP32|16384|2fwd+bwd|g|173.04|130.51|135.11|117.36|130.51|1.326|1.035|0.899|
|R200|208|4096 ->
1024|FP32|16231|wgrad|g|75.89|54.42|62.00|48.90|54.42|1.395|1.139|0.899|
|R200|208|4096 ->
1024|FP32|16231|bwd|g|130.17|108.05|114.50|92.23|108.05|1.205|1.060|0.854|
|R200|208|4096 ->
1024|FP32|16231|2fwd+bwd|g|233.16|197.62|218.12|190.94|197.62|1.180|1.104|0.966|
|R200|208|5120 ->
48|FP32|16231|wgrad|g|27.22|24.23|23.85|25.47|24.23|1.123|0.984|1.051|
|R200|208|5120 ->
48|FP32|16231|bwd|g|48.70|46.19|45.65|45.24|46.19|1.054|0.988|0.979|
|R200|208|5120 ->
48|FP32|16231|2fwd+bwd|g|96.09|93.05|92.94|91.69|93.05|1.033|0.999|0.985|
|R200|208|5120 ->
48|FP32|16384|wgrad|g|27.50|24.50|23.95|25.69|24.50|1.122|0.978|1.049|
|R200|208|5120 ->
48|FP32|16384|bwd|g|49.16|46.67|45.95|45.40|46.67|1.053|0.985|0.973|
|R200|208|5120 ->
48|FP32|16384|2fwd+bwd|g|96.74|93.84|93.54|92.00|93.84|1.031|0.997|0.980|
|R200|208|5120 ->
64|FP32|16231|wgrad|g|27.34|24.04|23.70|26.10|24.04|1.137|0.986|1.086|
|R200|208|5120 ->
64|FP32|16231|bwd|g|48.50|45.37|44.77|44.72|45.37|1.069|0.987|0.986|
|R200|208|5120 ->
64|FP32|16231|2fwd+bwd|g|95.67|92.04|91.96|91.12|92.04|1.039|0.999|0.990|
|R200|212|4096 ->
512|FP32|16231|wgrad|e|74.76|33.88|35.68|38.78|33.88|2.207|1.053|1.145|
|R200|212|4096 ->
512|FP32|16231|bwd|e|105.79|62.43|65.66|63.28|62.43|1.695|1.052|1.014|
|R200|212|4096 ->
512|FP32|16231|2fwd+bwd|e|168.11|125.32|128.52|113.52|125.32|1.341|1.026|0.906|
|R200|212|4096 ->
512|FP32|16384|wgrad|e|75.17|34.02|35.94|39.10|34.02|2.210|1.056|1.149|
|R200|212|4096 ->
512|FP32|16384|bwd|e|106.15|62.60|65.94|63.89|62.60|1.696|1.053|1.021|
|R200|212|4096 ->
512|FP32|16384|2fwd+bwd|e|168.74|125.55|129.04|113.95|125.55|1.344|1.028|0.908|
|R200|212|4096 ->
1024|FP32|16231|wgrad|e|76.08|54.28|52.68|48.93|54.28|1.402|0.971|0.901|
|R200|212|4096 ->
1024|FP32|16231|bwd|e|133.81|109.43|108.14|97.00|109.43|1.223|0.988|0.886|
|R200|212|4096 ->
1024|FP32|16231|2fwd+bwd|e|231.77|203.49|199.57|194.07|203.49|1.139|0.981|0.954|
|R200|212|4096 ->
4096|FP32|16172|wgrad|e|245.40|161.16|179.81|158.29|161.16|1.523|1.116|0.982|
|R200|212|4096 ->
4096|FP32|16172|bwd|e|465.09|363.96|402.73|362.73|363.96|1.278|1.107|0.997|
|R200|212|4096 ->
4096|FP32|16172|2fwd+bwd|e|860.52|730.20|789.30|740.34|730.20|1.178|1.081|1.014|
|R200|212|4096 ->
4096|FP32|16231|wgrad|e|246.17|161.34|179.84|158.05|161.34|1.526|1.115|0.980|
|R200|212|4096 ->
4096|FP32|16231|bwd|e|467.52|361.68|400.13|363.28|361.68|1.293|1.106|1.004|
|R200|212|4096 ->
4096|FP32|16231|2fwd+bwd|e|861.32|733.91|786.43|749.93|733.91|1.174|1.072|1.022|
|R200|212|4096 ->
4096|FP32|16384|wgrad|e|250.13|164.60|181.71|159.68|164.60|1.520|1.104|0.970|
|R200|212|4096 ->
4096|FP32|16384|bwd|e|470.40|367.67|402.45|367.98|367.67|1.279|1.095|1.001|
|R200|212|4096 ->
4096|FP32|16384|2fwd+bwd|e|865.40|735.59|791.04|752.38|735.59|1.176|1.075|1.023|
|R200|212|5120 ->
48|FP32|16231|wgrad|e|23.61|20.13|19.61|20.19|20.13|1.173|0.974|1.003|
|R200|212|5120 ->
48|FP32|16231|bwd|e|44.50|41.31|40.68|39.57|41.31|1.077|0.985|0.958|
|R200|212|5120 ->
48|FP32|16231|2fwd+bwd|e|83.87|80.67|79.95|77.05|80.67|1.040|0.991|0.955|
|R200|212|5120 ->
64|FP32|16231|wgrad|e|23.54|19.70|19.10|20.75|19.70|1.195|0.970|1.053|
|R200|212|5120 ->
64|FP32|16231|bwd|e|44.13|40.39|39.66|39.37|40.39|1.093|0.982|0.975|
|R200|212|5120 ->
64|FP32|16231|2fwd+bwd|e|83.40|79.88|79.03|76.99|79.88|1.044|0.989|0.964|
|R200|212|5120 ->
512|FP32|16231|wgrad|e|75.14|40.11|38.04|34.79|40.11|1.873|0.948|0.867|
|R200|212|5120 ->
512|FP32|16231|bwd|e|111.58|74.07|73.63|66.59|74.07|1.506|0.994|0.899|
|R200|212|5120 ->
512|FP32|16231|2fwd+bwd|e|186.06|148.91|148.22|126.11|148.91|1.249|0.995|0.847|
|R200|212|5120 ->
512|FP32|16384|wgrad|e|75.35|40.19|38.36|34.98|40.19|1.875|0.954|0.870|
|R200|212|5120 ->
512|FP32|16384|bwd|e|111.67|74.44|74.04|67.28|74.44|1.500|0.995|0.904|
|R200|212|5120 ->
512|FP32|16384|2fwd+bwd|e|186.73|149.74|148.87|128.19|149.74|1.247|0.994|0.856|
|R200|208|4096 ->
48|FP32|16231|wgrad|e|25.19|21.90|21.70|23.50|21.90|1.150|0.991|1.073|
|R200|208|4096 ->
48|FP32|16231|bwd|e|44.30|40.86|40.76|39.94|40.86|1.084|0.998|0.977|
|R200|208|4096 ->
48|FP32|16231|2fwd+bwd|e|86.38|82.50|82.38|80.49|82.50|1.047|0.999|0.976|
|R200|208|4096 ->
48|FP32|16384|wgrad|e|25.12|21.82|21.90|23.39|21.82|1.151|1.004|1.072|
|R200|208|4096 ->
48|FP32|16384|bwd|e|44.42|41.00|40.95|40.22|41.00|1.083|0.999|0.981|
|R200|208|4096 ->
48|FP32|16384|2fwd+bwd|e|86.75|82.86|82.79|81.02|82.86|1.047|0.999|0.978|
|R200|208|4096 ->
64|FP32|16231|wgrad|e|25.46|21.66|21.19|24.12|21.66|1.175|0.978|1.114|
|R200|208|4096 ->
64|FP32|16231|bwd|e|44.10|40.34|39.84|40.67|40.34|1.093|0.988|1.008|
|R200|208|4096 ->
64|FP32|16231|2fwd+bwd|e|86.68|81.91|81.61|80.90|81.91|1.058|0.996|0.988|
|R200|208|4096 ->
64|FP32|16384|wgrad|e|25.51|21.77|21.53|24.09|21.77|1.172|0.989|1.107|
|R200|208|4096 ->
64|FP32|16384|bwd|e|44.50|40.59|40.13|40.55|40.59|1.096|0.989|0.999|
|R200|208|4096 ->
64|FP32|16384|2fwd+bwd|e|87.09|82.38|81.97|81.42|82.38|1.057|0.995|0.988|
|R200|208|4096 ->
512|FP32|16231|wgrad|e|74.58|34.34|37.61|38.98|34.34|2.172|1.095|1.135|
|R200|208|4096 ->
512|FP32|16231|bwd|e|106.79|65.38|69.26|66.72|65.38|1.633|1.059|1.020|
|R200|208|4096 ->
512|FP32|16231|2fwd+bwd|e|175.26|133.36|137.68|119.10|133.36|1.314|1.032|0.893|
|R200|208|4096 ->
512|FP32|16384|wgrad|e|74.86|34.44|37.82|39.38|34.44|2.174|1.098|1.143|
|R200|208|4096 ->
512|FP32|16384|bwd|e|107.19|65.39|69.62|66.86|65.39|1.639|1.065|1.022|
|R200|208|4096 ->
512|FP32|16384|2fwd+bwd|e|175.72|133.43|138.01|119.22|133.43|1.317|1.034|0.894|
|R200|208|4096 ->
1024|FP32|16231|wgrad|e|75.98|54.47|62.02|48.98|54.47|1.395|1.139|0.899|
|R200|208|4096 ->
1024|FP32|16231|bwd|e|131.39|110.30|115.32|94.92|110.30|1.191|1.046|0.861|
|R200|208|4096 ->
1024|FP32|16231|2fwd+bwd|e|230.27|200.47|220.74|191.65|200.47|1.149|1.101|0.956|
|R200|208|5120 ->
48|FP32|16231|wgrad|e|27.22|24.28|24.06|25.87|24.28|1.121|0.991|1.065|
|R200|208|5120 ->
48|FP32|16231|bwd|e|49.42|46.74|46.38|45.62|46.74|1.057|0.992|0.976|
|R200|208|5120 ->
48|FP32|16231|2fwd+bwd|e|98.89|96.04|95.89|93.56|96.04|1.030|0.998|0.974|
|R200|208|5120 ->
48|FP32|16384|wgrad|e|27.48|24.56|24.20|26.03|24.56|1.119|0.985|1.060|
|R200|208|5120 ->
48|FP32|16384|bwd|e|49.87|47.07|46.77|45.97|47.07|1.059|0.994|0.977|
|R200|208|5120 ->
48|FP32|16384|2fwd+bwd|e|99.64|96.78|96.46|93.72|96.78|1.030|0.997|0.968|
|R200|208|5120 ->
64|FP32|16231|wgrad|e|27.50|24.19|23.83|26.46|24.19|1.137|0.985|1.094|
|R200|208|5120 ->
64|FP32|16231|bwd|e|49.34|46.11|45.59|45.76|46.11|1.070|0.989|0.992|
|R200|208|5120 ->
64|FP32|16231|2fwd+bwd|e|98.70|95.39|94.92|93.51|95.39|1.035|0.995|0.980|

---------

Co-authored-by: Yingyi Huang <averyh@nvidia.com>
yyihuang added a commit that referenced this pull request Nov 3, 2026
…192-wide 4096 -> 1024 FP32 and six more deterministic token counts (#8438)

## Summary

The Cake dense projection GEMM family covers BF16-input dense and
batched projection GEMMs, with BF16 or FP32 weight-gradient output, and
an FP32 router GEMM. In `flashinfer.experimental.dense_projection_gemm`,
the keyword-only `deterministic` flag selects the variant:
`deterministic=True` (the default) calls
`flashinfer.experimental.dense_projection_det_gemm`, and
`deterministic=False` runs the split-K / stream-K programs (#8190,
#8238). The family ships on `sm_100a` (GB200), `sm_103a` (GB300) and
`sm_107a` (R200).

This PR changes only the planner of `deterministic=False`. The change
applies to the projection weight gradient (K1) on `sm_107a` parts with
216 or more SMs. There are no kernel changes. #8363 shipped the 216-SM
entries of the deterministic token counts without a 216-SM measurement.
These entries are now measured on two 216-SM parts at 16172, 16231 and
16384 tokens.

This PR uses the rule of #8363. A point keeps the split-K program only
if all three conditions hold:

- (a) At the op level, the split-K program was faster than the
deterministic program in every run. This requires a one-sided paired
sign test p < 0.05 and a median time ratio below 0.99, in eager launches
and in CUDA graph replays.
- (b) Next to a kernel holding 8, 16 or 32 SMs on another stream, it was
no more than 3 % slower than the deterministic program, in both launch
orders.
- (c) It left the backward pass and the checkpointed step at most 1.3 %
slower than the deterministic program. This is the ratio of the medians
over the runs, eager and graph.

Otherwise, the point runs the deterministic program.

Changes:

- `4096 -> 1024` FP32 on exactly 216 SMs, 16172 to 16384 tokens: the
192-wide row-major program at one split (176 CTAs, as on 212 SMs since
#8275) replaces 256-wide tiles at one split (128 CTAs). 16231 tokens
keep the deterministic program.
- The deterministic program at 6 more points, from 216 SMs up: FP32
`4096 -> 32` at 16231 tokens; `4096 -> 64` at 16231 and 16384; `5120 ->
48` and `5120 -> 64` at 16384; BF16 `5120 -> 1024` at 16231.
- The 7 existing 216-SM points are confirmed. Other points, SM counts
and architectures are unchanged.

**Accuracy:** At `4096 -> 1024` FP32 (16172 and 16384 tokens), the
output of the 192-wide program is bitwise equal to the output of the
previous plan (256-wide tiles, one split) on the same inputs. The 6 new
deterministic points return the output of the deterministic program.

**Speedups** (ratio = baseline us / new us; above 1 means this PR is
faster):

- `4096 -> 1024` FP32 at 16172 / 16384 tokens, graph mode:
- previous/new: wgrad 1.160 / 1.164, bwd 1.106 / 1.122, 2fwd+bwd 1.064 /
1.070.
- det/new: wgrad 1.680 / 1.418, bwd 1.420 / 1.270, 2fwd+bwd 1.185 /
1.071.
- SOTA/new: wgrad 0.919 / 0.923, bwd 0.892 / 0.893, 2fwd+bwd 0.972 /
0.960.
- Eager mode previous/new: wgrad 1.177 / 1.179, bwd 1.106 / 1.108,
2fwd+bwd 1.052 / 1.016.
- The 6 new deterministic points against the previous plan
(previous/new; below 1 means the deterministic program is slower):
- Graph: wgrad 0.980 to 1.010 (geomean 0.996), bwd 0.988 to 1.009
(geomean 0.995), 2fwd+bwd 0.999 to 1.007 (geomean 1.003).
- Eager: wgrad 0.973 to 1.007 (geomean 0.990), bwd 0.991 to 1.015
(geomean 0.999), 2fwd+bwd 0.994 to 1.001 (geomean 0.998).
- The lowest value is `4096 -> 32` at 16231 tokens, eager wgrad 0.973.
There, the deterministic program is 2.7 % slower at the op.
- At the five FP32 points, the split-K program missed part (a) or (b) of
the rule in at least one of the 8 run-and-mode cells. It missed 1 of 8
at `4096 -> 32` and 5 to 8 of 8 at the other four.
- At BF16 `5120 -> 1024`, 16231 tokens, the split-K program passed (a)
and (b) in all 8 cells. The deterministic program is 2.0 % slower at the
op there (wgrad previous/new 0.980). However, the eager backward pass
with the split-K program was 1.5 % slower than with the deterministic
program (ratio of the medians over the runs; eager bwd previous/new
1.015). This is above the 1.3 % limit of (c), so that point runs the
deterministic program.

Definitions:

- **Rows:** 48 = 8 points x 3 ops x 2 modes, graph rows first; all on
216 SMs.
- **Ops:** wgrad = weight gradient `[d_out, d_in]` (output dtype as in
the `out` column) from BF16 activations `[T, d_in]` and BF16 output
gradients `[T, d_out]`. bwd = input gradient (BF16) + wgrad. 2fwd+bwd =
checkpointed step (2 x forward + bwd).
- **Modes:** g = CUDA graph replay, e = eager.
- **Latency:** Median us over 4 runs per row. The runs use two 216-SM
GPUs. Each GPU has one long-lived process over all rows and a fresh
process per row, with one arm order per GPU. In a run, all arms of a row
are interleaved in paired groups: 15 groups of 20 steps for wgrad, 9
groups for bwd and 2fwd+bwd. L2 is cold before every iteration. The time
is the CUPTI GPU time of the operation (first kernel begin to last
kernel end, all launches).
- **Baselines:**
- det Cake = `deterministic=True` at FlashInfer main eabf03f (the
deterministic package with the #8303 plans; the measured deterministic
package is the same tree).
- previous Cake = `deterministic=False` at FlashInfer main eabf03f,
without this PR. It uses the same kernel programs as this PR. Its
forward and input gradient are the deterministic programs.
- new = `deterministic=False` with this PR. At the 6 new deterministic
points, it runs the deterministic program, so new us is the det Cake
measurement. The `deterministic=False` call that dispatches to it was
not timed separately.
- SOTA = cuBLAS / cuBLASLt, the vendor GEMM library for `sm_107a` and
the strongest non-Cake GEMM available there. For wgrad, it is the
fastest correct launch among `torch.mm` and the cuBLASLt heuristic's
top-8 algorithms. For bwd and 2fwd+bwd, it is the fastest correct
library launch per op.
- Ratio = baseline us / new us, computed from the medians in the table;
above 1.000 = this PR is faster.
- **Setup:** R200 (`sm_107a`), 216 SMs, driver 615.71.10; CUDA 13.4
(nvcc 13.4.92), PyTorch 2.14.1+cu130, cuBLAS 13.8.0.4.
- **Noise:**
- Per-run paired previous/new: `4096 -> 1024` wgrad 1.156 to 1.196, bwd
1.086 to 1.124, 2fwd+bwd 1.044 to 1.061; the 6 deterministic points
wgrad 0.934 to 1.018, bwd 0.966 to 1.004, 2fwd+bwd 0.985 to 1.029.
- The two GPUs differ in speed by a few percent. Thus, a median over 4
runs (2 per GPU) can differ from the per-run paired ratios. At BF16
`5120 -> 1024`, every per-run paired bwd previous/new is 0.981 to 0.995,
while the ratio of the medians is 1.015 in eager mode.
- Some samples show the software power-cap throttle flag. All arms of a
row run interleaved in one process, so the ratios hold, but absolute
times can be a few percent high.

*Table: one row per point, op and mode; R200 (`sm_107a`), 216 SMs, CUDA
13.4; latencies in us (median over 4 runs); ratios = baseline us / new
us.*

| GPU | SMs | `d_in -> d_out` | out | T | op | mode | det Cake us | SOTA
us | previous Cake us | new us | det/new | SOTA/new | previous/new |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|R200|216|4096 ->
1024|FP32|16172|wgrad|g|90.10|49.30|62.20|53.62|1.680|0.919|1.160|
|R200|216|4096 ->
1024|FP32|16172|bwd|g|144.93|91.06|112.83|102.06|1.420|0.892|1.106|
|R200|216|4096 ->
1024|FP32|16172|2fwd+bwd|g|237.76|194.87|213.32|200.56|1.185|0.972|1.064|
|R200|216|4096 ->
1024|FP32|16384|wgrad|g|76.52|49.82|62.82|53.98|1.418|0.923|1.164|
|R200|216|4096 ->
1024|FP32|16384|bwd|g|128.83|90.58|113.85|101.47|1.270|0.893|1.122|
|R200|216|4096 ->
1024|FP32|16384|2fwd+bwd|g|217.54|194.93|217.28|203.04|1.071|0.960|1.070|
|R200|216|4096 ->
32|FP32|16231|wgrad|g|20.49|20.56|20.20|20.49|1.000|1.003|0.986|
|R200|216|4096 ->
32|FP32|16231|bwd|g|38.01|35.91|37.54|38.01|1.000|0.945|0.988|
|R200|216|4096 ->
32|FP32|16231|2fwd+bwd|g|75.78|73.70|75.74|75.78|1.000|0.973|0.999|
|R200|216|4096 ->
64|FP32|16231|wgrad|g|20.58|22.63|20.79|20.58|1.000|1.100|1.010|
|R200|216|4096 ->
64|FP32|16231|bwd|g|38.34|38.01|38.21|38.34|1.000|0.991|0.997|
|R200|216|4096 ->
64|FP32|16231|2fwd+bwd|g|75.30|75.64|75.82|75.30|1.000|1.005|1.007|
|R200|216|4096 ->
64|FP32|16384|wgrad|g|20.94|22.73|20.83|20.94|1.000|1.085|0.995|
|R200|216|4096 ->
64|FP32|16384|bwd|g|38.64|38.34|38.42|38.64|1.000|0.992|0.994|
|R200|216|4096 ->
64|FP32|16384|2fwd+bwd|g|76.14|75.92|76.34|76.14|1.000|0.997|1.003|
|R200|216|5120 ->
48|FP32|16384|wgrad|g|23.24|24.74|23.41|23.24|1.000|1.065|1.007|
|R200|216|5120 ->
48|FP32|16384|bwd|g|44.85|43.47|44.55|44.85|1.000|0.969|0.993|
|R200|216|5120 ->
48|FP32|16384|2fwd+bwd|g|90.12|88.88|90.29|90.12|1.000|0.986|1.002|
|R200|216|5120 ->
64|FP32|16384|wgrad|g|23.28|25.22|23.19|23.28|1.000|1.083|0.996|
|R200|216|5120 ->
64|FP32|16384|bwd|g|43.95|43.82|43.57|43.95|1.000|0.997|0.991|
|R200|216|5120 ->
64|FP32|16384|2fwd+bwd|g|89.14|88.88|89.50|89.14|1.000|0.997|1.004|
|R200|216|5120 ->
1024|BF16|16231|wgrad|g|64.54|50.46|63.28|64.54|1.000|0.782|0.980|
|R200|216|5120 ->
1024|BF16|16231|bwd|g|133.20|115.45|134.46|133.20|1.000|0.867|1.009|
|R200|216|5120 ->
1024|BF16|16231|2fwd+bwd|g|257.66|236.33|258.33|257.66|1.000|0.917|1.003|
|R200|216|4096 ->
1024|FP32|16172|wgrad|e|89.33|49.39|62.30|52.95|1.687|0.933|1.177|
|R200|216|4096 ->
1024|FP32|16172|bwd|e|145.58|91.82|113.44|102.60|1.419|0.895|1.106|
|R200|216|4096 ->
1024|FP32|16172|2fwd+bwd|e|239.96|195.78|215.59|204.90|1.171|0.955|1.052|
|R200|216|4096 ->
1024|FP32|16384|wgrad|e|76.44|49.93|62.94|53.39|1.432|0.935|1.179|
|R200|216|4096 ->
1024|FP32|16384|bwd|e|129.36|91.78|114.95|103.77|1.247|0.884|1.108|
|R200|216|4096 ->
1024|FP32|16384|2fwd+bwd|e|219.94|196.50|209.68|206.32|1.066|0.952|1.016|
|R200|216|4096 ->
32|FP32|16231|wgrad|e|20.73|20.92|20.17|20.73|1.000|1.009|0.973|
|R200|216|4096 ->
32|FP32|16231|bwd|e|38.46|36.70|38.13|38.46|1.000|0.954|0.991|
|R200|216|4096 ->
32|FP32|16231|2fwd+bwd|e|78.14|75.71|77.73|78.14|1.000|0.969|0.995|
|R200|216|4096 ->
64|FP32|16231|wgrad|e|20.86|23.15|20.80|20.86|1.000|1.110|0.997|
|R200|216|4096 ->
64|FP32|16231|bwd|e|38.94|39.08|38.90|38.94|1.000|1.004|0.999|
|R200|216|4096 ->
64|FP32|16231|2fwd+bwd|e|78.22|77.99|78.26|78.22|1.000|0.997|1.001|
|R200|216|4096 ->
64|FP32|16384|wgrad|e|21.20|23.25|20.89|21.20|1.000|1.097|0.985|
|R200|216|4096 ->
64|FP32|16384|bwd|e|39.37|39.28|39.06|39.37|1.000|0.998|0.992|
|R200|216|4096 ->
64|FP32|16384|2fwd+bwd|e|78.97|78.28|78.64|78.97|1.000|0.991|0.996|
|R200|216|5120 ->
48|FP32|16384|wgrad|e|23.46|25.15|23.63|23.46|1.000|1.072|1.007|
|R200|216|5120 ->
48|FP32|16384|bwd|e|45.25|44.25|45.30|45.25|1.000|0.978|1.001|
|R200|216|5120 ->
48|FP32|16384|2fwd+bwd|e|93.17|91.17|93.30|93.17|1.000|0.979|1.001|
|R200|216|5120 ->
64|FP32|16384|wgrad|e|23.42|25.70|23.30|23.42|1.000|1.097|0.995|
|R200|216|5120 ->
64|FP32|16384|bwd|e|44.64|44.46|44.41|44.64|1.000|0.996|0.995|
|R200|216|5120 ->
64|FP32|16384|2fwd+bwd|e|92.55|91.19|92.42|92.55|1.000|0.985|0.999|
|R200|216|5120 ->
1024|BF16|16231|wgrad|e|64.57|50.54|63.30|64.57|1.000|0.783|0.980|
|R200|216|5120 ->
1024|BF16|16231|bwd|e|133.64|115.54|135.58|133.64|1.000|0.865|1.015|
|R200|216|5120 ->
1024|BF16|16231|2fwd+bwd|e|260.66|236.04|259.15|260.66|1.000|0.906|0.994|

Other Cake variants of this family (the deterministic package, the
batched MLA head rows, the FP32 router GEMM, BF16 and FP32
weight-gradient outputs) and the `sm_100a` / `sm_103a` planners are
unchanged; the family ships on `sm_100a`, `sm_103a` and `sm_107a`.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **Performance**
* Dense projection GEMM now uses refined execution plans for selected
GPU configurations and input sizes, including 216-SM devices.
* The updated plans cover specific FP32 and BF16 workloads and token
counts, with exclusions where a tuned plan does not apply.
* Execution selection remains deterministic for the supported
configurations.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Yingyi Huang <averyh@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants