[yyihuang]

Related issue: #8066

flashinfer.cake_rmsnorm_train is RMSNorm for training. The forward returns y (BF16) and rstd (FP32), and the backward returns dx (BF16) and dw (FP32). The two modes, deterministic=True and deterministic=False, share the forward and differ only in the backward. Both modes support:

• [T, H] rows with BF16 or FP32 weight
• a fused residual add (H6144)
• per-head views [T,24,256] (strided) and [T,4,256] (contiguous), with a shared w[256]

This PR follows #8577, #8626, #8655, #8704, #8810, #8827, #8848 and #8849. It changes only the deterministic=False backward on sm_103a (GB300). The change covers H1536, H2048 and H6144, with BF16 or FP32 weight, plain rows, in the windows listed below. FlashInfer main runs a TwoSum single in all of these windows, the same program that stays above T 131072. With this PR, the backward in each window runs one of two program kinds:

• kb: a single kernel with blocked FP32 dw sums, the kernel kind that perf(cake_rmsnorm): blocked FP32 dw sums for the large-T GB300 and R200 backward #8655 ships on GB300
• det: the same programs as deterministic=True

The windows are:

• H1536, BF16 weight: kb on T 46593 to 53184 and on T 76353 to 131072; det on T 53185 to 76352
• H1536, FP32 weight: kb on T 37889 to 40448 and on T 76353 to 131072; det on T 40449 to 76352
• H2048, both weight dtypes: det on T 61953 to 68160 (the programs that fix(cake_rmsnorm_train): end the sm_103a H2048 det-pair backward windows below their dw budget breach #8810 routes on T 58497 to 61952); kb on T 68161 to 131072
• H6144, BF16 weight: det on T 49153 to 57344; kb on T 57345 to 131072
• H6144, FP32 weight: det on T 49153 to 83392; kb on T 83393 to 131072

In every ratio below, a value under 1 means the baseline is faster than new. The ranges cover all rows of the table, except that the SOTA ranges leave out the rows marked !.

• Against main. main is the TwoSum single this PR replaces. new is faster on every bwd and 2xfwd+bwd row where main was timed (20 shapes, run A). bwd is 1.0279 to 1.2870 and 2xfwd+bwd is 1.0128 to 1.1468. These ratios are det-anchored (see Measurement).
• Against det on the 20 det shapes. new runs det's programs there, and the table repeats det's time, so det/new is 1.0000.
• Against det on the 27 kb shapes. det is faster than new on 25 of 27 bwd rows (det/new 0.9426 to 0.9989). new is faster on the other 2: H6144 with BF16 weight at T 98304 and 131072 (1.0060 and 1.0115). On 2xfwd+bwd, det/new is 0.9699 to 1.0058, and det is faster on 24 of 27 rows. On 23 bwd rows and 18 2xfwd+bwd rows (marked -), det is faster than new in all 4 sessions, each time by more than the row's A/A floor in that session. new is faster in the same sense on 3 rows (marked +): bwd of H6144 with BF16 weight at T 98304 and 131072, and 2xfwd+bwd at T 131072.
• Against the strongest in-budget non-Cake SOTA. On the 43 shapes that have an in-budget SOTA, bwd is 0.9896 to 1.5171 and 2xfwd+bwd is 0.9964 to 1.3069. SOTA is faster than new on 1 bwd row (QA, H1536 T 131072 with FP32 weight) and on 1 2xfwd+bwd row (TC, H2048 T 114688 with BF16 weight).
  • On 4 shapes, no non-Cake candidate meets the dw budget (marked !): H6144 T 131072 with BF16 weight, and H6144 T 114688, 122880 and 131072 with FP32 weight.
  • On these shapes, the fastest candidate that has these outputs is QA. Its worst dw element is 1.543 to 2.976 times the budget. It is faster than new: bwd 0.9210 to 0.9630 and 2xfwd+bwd 0.9666 to 0.9829.
• Forward. This PR does not change the forward, and the fwd values are as measured. main/new is 0.9961 to 1.0018. det/new on kb shapes is 0.9956 to 1.0030. SOTA/new is 0.9893 to 1.1306 and is below 1 on 7 rows: every H6144 row with BF16 weight, and H2048 T 131072 with BF16 weight.

Measurement

• Latency: the median of the CUPTI span of one CUDA-graph replay with L2 flushed (cold L2). Each run uses the 4 GPUs of one GB300 node, one session per GPU. Each session times 201 replays per implementation and op, split over 3 interleaved rounds, and the median pools the 4 sessions. Each implementation's own setup, reduction, cast and copy kernels are inside the span. In each session, det, new and every non-Cake candidate run in one process, and the implementation order rotates. Over the table's rows, the largest |1 - A/A| in a session is 0.0066 to 0.0341 across the 12 sessions of the three runs. In the 4 sessions of main's process, it is 0.0080 to 0.0277.
• Context: on small rows, Cake gains somewhat more than the other non-compile libraries from sharing the process with torch.compile. On GB300, SOTA/new vs the non-compile candidates is up to about 1.10x higher in the one-process measurement than in a process without torch.compile; on large rows the difference is <= 1.02x. Every row here is a large row (T >= 37889). Ratios vs det are unaffected.
• Runs: A, B and C ran on three different GB300 nodes. Each row comes from one run (column run): from A where A measured the shape, otherwise from B or C. A and one other run both measured 8 shapes. On those shapes, the other run's det/new differs from A's by at most 0.0017, and its SOTA/new by at most 0.0084.
• Builds: det and new were timed on an earlier build of this branch. On kb rows, that build runs the same kb program as this PR's head. It also ran kb in the det windows, so on det rows new us is det's time (see Key). Its det programs in these windows are the same as on FlashInfer main at this PR's base, and this PR does not change them. main, and the det it is anchored to, were timed on an earlier FlashInfer main build. In these windows, that build's main and det programs are the same as on FlashInfer main at this PR's base.
• det and new: det = deterministic Cake (deterministic=True), which this PR does not change. new = this PR's non-deterministic Cake (deterministic=False).
• SOTA: the fastest non-Cake candidate in the run that has these outputs and meets the accuracy budget. The budget is dx atol = rtol = 1e-2 and dw per element atol 1e-4 + rtol 1e-4 vs FP64, taking the worst of 5 seeds. new meets the dw budget in every changed window. On the 8 rows marked !, no candidate meets the budget (see Key).
  • Candidates: QuACK 0.6.5 (heuristic config and per-shape autotuned), torch.compile (torch 2.14.1+cu130; max-autotune-no-cudagraphs and default mode) of the explicit FP32-statistics reference, Liger-Kernel 0.8.5, and Transformer Engine 2.20.2 (run as the FP32 problem).
  • On BF16-weight rows, library APIs that return dw in BF16 do not have these outputs, so they are never chosen as SOTA.
• main: this backend on FlashInfer main before this PR. Its deterministic=False backward runs the TwoSum single in these windows.
  • main was timed only in run A, in its own process together with det.
  • main/new = (main / det of that process) / (new / det of new's process).
  • main us is main's latency in that process. main/new is det-anchored, so it can differ slightly from main us / new us. n/a means main was not timed at that shape.
• Speedup: baseline latency / new latency (> 1: new is faster; < 1: the baseline is faster).

Key:

• kb / det at the end of the workload: the program new runs at that T. kb = the single kernel with blocked FP32 dw sums. det = the programs of deterministic=True. On det rows, new us is det's measurement from the same run, so det/new is 1.0000 by construction.
• op 2xfwd+bwd: checkpointed step (forward, forward recompute, backward).
• + / - after det/new: in each of the 4 sessions of that run, det/new differs from 1 by more than the row's A/A floor in that session, and all 4 sessions agree that new (+) or det (-) is faster. Each session also times a second copy of det and of new. A row's A/A floor in a session is the larger |1 - A/A| of its det and new copies. Unmarked rows do not meet that test.
• ! after the SOTA impl and after SOTA/new: no non-Cake candidate meets the dw budget. The table shows the fastest candidate that has these outputs, with its worst dw element as a multiple of the budget in parentheses, and computes SOTA/new against it.
• SOTA names:
  • Q = QuACK, heuristic config
  • QA = QuACK, per-shape autotuned
  • TC = torch.compile, inductor mode max-autotune-no-cudagraphs
  • TCd = torch.compile, inductor default mode

GB300 (sm_103a): new non-deterministic Cake vs deterministic Cake, the strongest non-Cake SOTA and main, every measured shape

• GPU: GB300, sm_103a, 152 SMs, compute capability 10.3.
• Software: CUDA toolkit 13.4; torch 2.14.1+cu130 (CUDA runtime 13.0).

GPU / arch	workload [T, H] x, w, form, program	op	det Cake us	SOTA us (impl)	main us	new ND us	det/new	SOTA/new	main/new	run
GB300 / sm_103a	[46593, 1536] BF16, w BF16, plain, kb	fwd	39.9360	41.9520 QA	40.0320	40.0000	0.9984	1.0488	1.0008	A
GB300 / sm_103a	[46593, 1536] BF16, w BF16, plain, kb	bwd	75.2320	83.0730 QA	95.6480	76.9920	0.9771	1.0790	1.2413	A
GB300 / sm_103a	[46593, 1536] BF16, w BF16, plain, kb	2xfwd+bwd	153.6650	164.1610 QA	176.9450	157.8570	0.9734	1.0399	1.1205	A
GB300 / sm_103a	[53248, 1536] BF16, w BF16, plain, det	fwd	45.2800	47.6160 QA	n/a	45.2800	1.0000	1.0516	n/a	B
GB300 / sm_103a	[53248, 1536] BF16, w BF16, plain, det	bwd	84.2240	95.3930 TC	n/a	84.2240	1.0000	1.1326	n/a	B
GB300 / sm_103a	[53248, 1536] BF16, w BF16, plain, det	2xfwd+bwd	172.9930	194.2090 TC	n/a	172.9930	1.0000	1.1226	n/a	B
GB300 / sm_103a	[65536, 1536] BF16, w BF16, plain, det	fwd	56.4480	58.3690 QA	56.3530	56.4480	1.0000	1.0340	0.9994	A
GB300 / sm_103a	[65536, 1536] BF16, w BF16, plain, det	bwd	101.2160	108.1600 QA	128.9290	101.2160	1.0000	1.0686	1.2734	A
GB300 / sm_103a	[65536, 1536] BF16, w BF16, plain, det	2xfwd+bwd	212.2890	221.8730 QA	243.3765	212.2890	1.0000	1.0451	1.1468	A
GB300 / sm_103a	[75776, 1536] BF16, w BF16, plain, det	fwd	64.4480	67.0080 QA	n/a	64.4480	1.0000	1.0397	n/a	C
GB300 / sm_103a	[75776, 1536] BF16, w BF16, plain, det	bwd	114.8170	121.3765 QA	n/a	114.8170	1.0000	1.0571	n/a	C
GB300 / sm_103a	[75776, 1536] BF16, w BF16, plain, det	2xfwd+bwd	241.5050	251.5530 QA	n/a	241.5050	1.0000	1.0416	n/a	C
GB300 / sm_103a	[81920, 1536] BF16, w BF16, plain, kb	fwd	69.9840	72.2560 QA	n/a	69.9840	1.0000	1.0325	n/a	B
GB300 / sm_103a	[81920, 1536] BF16, w BF16, plain, kb	bwd	123.1040	134.9130 TC	n/a	125.4730	0.9811 -	1.0752	n/a	B
GB300 / sm_103a	[81920, 1536] BF16, w BF16, plain, kb	2xfwd+bwd	260.6720	277.6645 TC	n/a	266.7195	0.9773 -	1.0410	n/a	B
GB300 / sm_103a	[98304, 1536] BF16, w BF16, plain, kb	fwd	83.8090	86.4320 QA	n/a	83.9685	0.9981	1.0293	n/a	B
GB300 / sm_103a	[98304, 1536] BF16, w BF16, plain, kb	bwd	146.3360	163.1680 TC	n/a	149.0560	0.9818 -	1.0947	n/a	B
GB300 / sm_103a	[98304, 1536] BF16, w BF16, plain, kb	2xfwd+bwd	311.5210	333.9850 TC	n/a	317.6815	0.9806 -	1.0513	n/a	B
GB300 / sm_103a	[131072, 1536] BF16, w BF16, plain, kb	fwd	112.3530	114.7210 QA	112.3520	112.2890	1.0006	1.0217	1.0003	A
GB300 / sm_103a	[131072, 1536] BF16, w BF16, plain, kb	bwd	190.5610	200.9930 TC	241.4575	194.0810	0.9819 -	1.0356	1.2441	A
GB300 / sm_103a	[131072, 1536] BF16, w BF16, plain, kb	2xfwd+bwd	412.7705	426.9460 TC	467.4585	420.2735	0.9821 -	1.0159	1.1123	A
GB300 / sm_103a	[61953, 2048] BF16, w BF16, plain, det	fwd	70.8480	70.9920 Q	70.8480	70.8480	1.0000	1.0020	1.0018	A
GB300 / sm_103a	[61953, 2048] BF16, w BF16, plain, det	bwd	122.1130	137.6000 TC	138.1130	122.1130	1.0000	1.1268	1.1319	A
GB300 / sm_103a	[61953, 2048] BF16, w BF16, plain, det	2xfwd+bwd	263.7780	281.9050 QA	280.9290	263.7780	1.0000	1.0687	1.0658	A
GB300 / sm_103a	[64000, 2048] BF16, w BF16, plain, det	fwd	73.0240	73.4080 Q	n/a	73.0240	1.0000	1.0053	n/a	C
GB300 / sm_103a	[64000, 2048] BF16, w BF16, plain, det	bwd	126.5930	141.7280 TC	n/a	126.5930	1.0000	1.1196	n/a	C
GB300 / sm_103a	[64000, 2048] BF16, w BF16, plain, det	2xfwd+bwd	272.5150	291.0400 TC	n/a	272.5150	1.0000	1.0680	n/a	C
GB300 / sm_103a	[66560, 2048] BF16, w BF16, plain, det	fwd	75.9040	76.1930 Q	n/a	75.9040	1.0000	1.0038	n/a	C
GB300 / sm_103a	[66560, 2048] BF16, w BF16, plain, det	bwd	130.5925	142.9770 TC	n/a	130.5925	1.0000	1.0948	n/a	C
GB300 / sm_103a	[66560, 2048] BF16, w BF16, plain, det	2xfwd+bwd	281.9860	297.6650 TC	n/a	281.9860	1.0000	1.0556	n/a	C
GB300 / sm_103a	[69632, 2048] BF16, w BF16, plain, kb	fwd	79.5190	79.7440 Q	n/a	79.4560	1.0008	1.0036	n/a	C
GB300 / sm_103a	[69632, 2048] BF16, w BF16, plain, kb	bwd	135.9040	146.6570 TC	n/a	140.7040	0.9659 -	1.0423	n/a	C
GB300 / sm_103a	[69632, 2048] BF16, w BF16, plain, kb	2xfwd+bwd	294.2420	307.4890 TC	n/a	300.0340	0.9807 -	1.0248	n/a	C
GB300 / sm_103a	[81920, 2048] BF16, w BF16, plain, kb	fwd	93.4730	93.5360 Q	93.6005	93.4730	1.0000	1.0007	1.0010	A
GB300 / sm_103a	[81920, 2048] BF16, w BF16, plain, kb	bwd	158.2885	166.4330 TC	177.2170	163.2330	0.9697 -	1.0196	1.0840	A
GB300 / sm_103a	[81920, 2048] BF16, w BF16, plain, kb	2xfwd+bwd	345.4720	354.6250 TC	366.0505	351.6020	0.9826 -	1.0086	1.0405	A
GB300 / sm_103a	[98304, 2048] BF16, w BF16, plain, kb	fwd	111.8400	111.8720 Q	n/a	111.8720	0.9997	1.0000	n/a	B
GB300 / sm_103a	[98304, 2048] BF16, w BF16, plain, kb	bwd	186.3040	212.0655 TC	n/a	194.6240	0.9573 -	1.0896	n/a	B
GB300 / sm_103a	[98304, 2048] BF16, w BF16, plain, kb	2xfwd+bwd	409.9200	436.1130 TC	n/a	419.7600	0.9766 -	1.0390	n/a	B
GB300 / sm_103a	[114688, 2048] BF16, w BF16, plain, kb	fwd	130.7520	130.9115 Q	n/a	130.8480	0.9993	1.0005	n/a	B
GB300 / sm_103a	[114688, 2048] BF16, w BF16, plain, kb	bwd	214.7850	225.1200 TC	n/a	223.8720	0.9594 -	1.0056	n/a	B
GB300 / sm_103a	[114688, 2048] BF16, w BF16, plain, kb	2xfwd+bwd	475.8085	487.3610 TC	n/a	489.1210	0.9728 -	0.9964	n/a	B
GB300 / sm_103a	[131072, 2048] BF16, w BF16, plain, kb	fwd	152.3850	151.4250 TC	153.0570	153.0570	0.9956	0.9893	0.9994	A
GB300 / sm_103a	[131072, 2048] BF16, w BF16, plain, kb	bwd	242.6260	259.4095 TC	273.6330	252.8975	0.9594 -	1.0257	1.0820	A
GB300 / sm_103a	[131072, 2048] BF16, w BF16, plain, kb	2xfwd+bwd	541.3790	558.6915 TC	577.4275	556.5780	0.9727 -	1.0038	1.0373	A
GB300 / sm_103a	[49153, 6144] BF16, w BF16, plain, det	fwd	171.8410	170.7850 TC	171.8410	171.8410	1.0000	0.9939	1.0000	A
GB300 / sm_103a	[49153, 6144] BF16, w BF16, plain, det	bwd	277.5530	329.5060 TC	294.3060	277.5530	1.0000	1.1872	1.0615	A
GB300 / sm_103a	[49153, 6144] BF16, w BF16, plain, det	2xfwd+bwd	614.5470	667.9390 TC	632.1640	614.5470	1.0000	1.0869	1.0289	A
GB300 / sm_103a	[57344, 6144] BF16, w BF16, plain, det	fwd	199.6480	199.1690 QA	n/a	199.6480	1.0000	0.9976	n/a	C
GB300 / sm_103a	[57344, 6144] BF16, w BF16, plain, det	bwd	322.0815	336.8980 TC	n/a	322.0815	1.0000	1.0460	n/a	C
GB300 / sm_103a	[57344, 6144] BF16, w BF16, plain, det	2xfwd+bwd	713.7625	747.2995 TC	n/a	713.7625	1.0000	1.0470	n/a	C
GB300 / sm_103a	[65536, 6144] BF16, w BF16, plain, kb	fwd	227.9370	227.4260 Q	227.9370	227.7770	1.0007	0.9985	1.0010	A
GB300 / sm_103a	[65536, 6144] BF16, w BF16, plain, kb	bwd	366.1775	389.7300 TC	384.8975	366.5950	0.9989	1.0631	1.0501	A
GB300 / sm_103a	[65536, 6144] BF16, w BF16, plain, kb	2xfwd+bwd	815.2520	859.1410 TC	834.3090	814.8520	1.0005	1.0544	1.0241	A
GB300 / sm_103a	[73728, 6144] BF16, w BF16, plain, kb	fwd	255.3610	254.4010 QA	n/a	255.3450	1.0001	0.9963	n/a	C
GB300 / sm_103a	[73728, 6144] BF16, w BF16, plain, kb	bwd	409.2650	569.8910 TCd	n/a	410.9620	0.9959 -	1.3867	n/a	C
GB300 / sm_103a	[73728, 6144] BF16, w BF16, plain, kb	2xfwd+bwd	914.2270	1112.1000 TCd	n/a	914.8355	0.9993	1.2156	n/a	C
GB300 / sm_103a	[98304, 6144] BF16, w BF16, plain, kb	fwd	339.8420	338.7540 Q	339.8100	339.7940	1.0001	0.9969	1.0001	A
GB300 / sm_103a	[98304, 6144] BF16, w BF16, plain, kb	bwd	542.4825	602.4350 TC	564.9945	539.2350	1.0060 +	1.1172	1.0477	A
GB300 / sm_103a	[98304, 6144] BF16, w BF16, plain, kb	2xfwd+bwd	1214.8545	1298.5990 TC	1239.5595	1212.8055	1.0017	1.0707	1.0211	A
GB300 / sm_103a	[131072, 6144] BF16, w BF16, plain, kb	fwd	451.8905	449.7940 Q	451.7470	451.7785	1.0002	0.9956	1.0000	A
GB300 / sm_103a	[131072, 6144] BF16, w BF16, plain, kb	bwd	717.2845	682.9315 QA! (1.693)	746.1810	709.1390	1.0115 +	0.9630!	1.0514	A
GB300 / sm_103a	[131072, 6144] BF16, w BF16, plain, kb	2xfwd+bwd	1616.0885	1579.4000 QA! (1.693)	1645.1775	1606.8245	1.0058 +	0.9829!	1.0237	A
GB300 / sm_103a	[37889, 1536] BF16, w FP32, plain, kb	fwd	32.9920	37.4400 QA	33.2480	33.1200	0.9961	1.1304	0.9990	A
GB300 / sm_103a	[37889, 1536] BF16, w FP32, plain, kb	bwd	60.6560	69.4725 QA	80.7685	64.3530	0.9426 -	1.0796	1.2561	A
GB300 / sm_103a	[37889, 1536] BF16, w FP32, plain, kb	2xfwd+bwd	124.8650	141.7610 QA	145.0570	128.7335	0.9699 -	1.1012	1.1259	A
GB300 / sm_103a	[40449, 1536] BF16, w FP32, plain, det	fwd	35.0410	39.6160 QA	n/a	35.0410	1.0000	1.1306	n/a	B
GB300 / sm_103a	[40449, 1536] BF16, w FP32, plain, det	bwd	65.4410	72.3520 QA	n/a	65.4410	1.0000	1.1056	n/a	B
GB300 / sm_103a	[40449, 1536] BF16, w FP32, plain, det	2xfwd+bwd	134.1445	149.0560 QA	n/a	134.1445	1.0000	1.1112	n/a	B
GB300 / sm_103a	[46593, 1536] BF16, w FP32, plain, det	fwd	40.3840	43.9675 TC	n/a	40.3840	1.0000	1.0887	n/a	B
GB300 / sm_103a	[46593, 1536] BF16, w FP32, plain, det	bwd	73.9210	80.5120 QA	n/a	73.9210	1.0000	1.0892	n/a	B
GB300 / sm_103a	[46593, 1536] BF16, w FP32, plain, det	2xfwd+bwd	153.1200	167.9360 QA	n/a	153.1200	1.0000	1.0968	n/a	B
GB300 / sm_103a	[53248, 1536] BF16, w FP32, plain, det	fwd	46.1440	50.7200 TC	n/a	46.1440	1.0000	1.0992	n/a	B
GB300 / sm_103a	[53248, 1536] BF16, w FP32, plain, det	bwd	83.3600	126.4640 Q	n/a	83.3600	1.0000	1.5171	n/a	B
GB300 / sm_103a	[53248, 1536] BF16, w FP32, plain, det	2xfwd+bwd	173.6000	226.8805 Q	n/a	173.6000	1.0000	1.3069	n/a	B
GB300 / sm_103a	[65536, 1536] BF16, w FP32, plain, det	fwd	57.4730	60.8000 TC	57.2480	57.4730	1.0000	1.0579	0.9961	A
GB300 / sm_103a	[65536, 1536] BF16, w FP32, plain, det	bwd	100.4490	106.0645 QA	129.2810	100.4490	1.0000	1.0559	1.2870	A
GB300 / sm_103a	[65536, 1536] BF16, w FP32, plain, det	2xfwd+bwd	213.6015	228.5450 QA	241.9855	213.6015	1.0000	1.0700	1.1325	A
GB300 / sm_103a	[75776, 1536] BF16, w FP32, plain, det	fwd	65.8240	69.5680 TC	n/a	65.8240	1.0000	1.0569	n/a	C
GB300 / sm_103a	[75776, 1536] BF16, w FP32, plain, det	bwd	114.1760	119.2965 QA	n/a	114.1760	1.0000	1.0448	n/a	C
GB300 / sm_103a	[75776, 1536] BF16, w FP32, plain, det	2xfwd+bwd	243.1055	259.3610 QA	n/a	243.1055	1.0000	1.0669	n/a	C
GB300 / sm_103a	[81920, 1536] BF16, w FP32, plain, kb	fwd	71.1680	74.2080 TC	n/a	71.1360	1.0004	1.0432	n/a	B
GB300 / sm_103a	[81920, 1536] BF16, w FP32, plain, kb	bwd	122.0800	127.1040 QA	n/a	125.6320	0.9717 -	1.0117	n/a	B
GB300 / sm_103a	[81920, 1536] BF16, w FP32, plain, kb	2xfwd+bwd	261.3440	277.8565 QA	n/a	264.6080	0.9877	1.0501	n/a	B
GB300 / sm_103a	[98304, 1536] BF16, w FP32, plain, kb	fwd	85.4085	88.1600 TC	n/a	85.1525	1.0030	1.0353	n/a	B
GB300 / sm_103a	[98304, 1536] BF16, w FP32, plain, kb	bwd	145.2160	158.1120 TC	n/a	148.7365	0.9763 -	1.0630	n/a	B
GB300 / sm_103a	[98304, 1536] BF16, w FP32, plain, kb	2xfwd+bwd	314.0160	331.1045 TC	n/a	316.2880	0.9928	1.0468	n/a	B
GB300 / sm_103a	[114688, 1536] BF16, w FP32, plain, kb	fwd	99.5850	102.5600 TC	n/a	99.4560	1.0013	1.0312	n/a	B
GB300 / sm_103a	[114688, 1536] BF16, w FP32, plain, kb	bwd	167.2650	176.5460 TC	n/a	171.5200	0.9752 -	1.0293	n/a	B
GB300 / sm_103a	[114688, 1536] BF16, w FP32, plain, kb	2xfwd+bwd	363.6810	375.3610 TC	n/a	367.9040	0.9885	1.0203	n/a	B
GB300 / sm_103a	[131072, 1536] BF16, w FP32, plain, kb	fwd	114.4320	116.8330 TC	114.2730	114.2405	1.0017	1.0227	1.0006	A
GB300 / sm_103a	[131072, 1536] BF16, w FP32, plain, kb	bwd	189.6970	191.7450 QA	242.6735	193.7610	0.9790 -	0.9896	1.2526	A
GB300 / sm_103a	[131072, 1536] BF16, w FP32, plain, kb	2xfwd+bwd	415.6345	428.9300 QA	468.8345	419.6170	0.9905	1.0222	1.1179	A
GB300 / sm_103a	[61953, 2048] BF16, w FP32, plain, det	fwd	72.4800	73.9200 QA	72.5120	72.4800	1.0000	1.0199	0.9985	A
GB300 / sm_103a	[61953, 2048] BF16, w FP32, plain, det	bwd	122.0165	132.7370 QA	145.7280	122.0165	1.0000	1.0879	1.1943	A
GB300 / sm_103a	[61953, 2048] BF16, w FP32, plain, det	2xfwd+bwd	264.0020	275.4900 QA	288.5460	264.0020	1.0000	1.0435	1.0926	A
GB300 / sm_103a	[64000, 2048] BF16, w FP32, plain, det	fwd	74.5600	75.8415 QA	n/a	74.5600	1.0000	1.0172	n/a	C
GB300 / sm_103a	[64000, 2048] BF16, w FP32, plain, det	bwd	126.4640	136.2085 QA	n/a	126.4640	1.0000	1.0771	n/a	C
GB300 / sm_103a	[64000, 2048] BF16, w FP32, plain, det	2xfwd+bwd	272.6090	283.1050 QA	n/a	272.6090	1.0000	1.0385	n/a	C
GB300 / sm_103a	[66560, 2048] BF16, w FP32, plain, det	fwd	77.1840	78.6405 QA	n/a	77.1840	1.0000	1.0189	n/a	C
GB300 / sm_103a	[66560, 2048] BF16, w FP32, plain, det	bwd	130.3680	141.2810 QA	n/a	130.3680	1.0000	1.0837	n/a	C
GB300 / sm_103a	[66560, 2048] BF16, w FP32, plain, det	2xfwd+bwd	282.7195	294.4170 QA	n/a	282.7195	1.0000	1.0414	n/a	C
GB300 / sm_103a	[69632, 2048] BF16, w FP32, plain, kb	fwd	81.0250	82.3370 QA	n/a	81.0560	0.9996	1.0158	n/a	C
GB300 / sm_103a	[69632, 2048] BF16, w FP32, plain, kb	bwd	136.1920	151.1370 TC	n/a	140.1920	0.9715 -	1.0781	n/a	C
GB300 / sm_103a	[69632, 2048] BF16, w FP32, plain, kb	2xfwd+bwd	295.6490	314.0810 TC	n/a	299.9050	0.9858 -	1.0473	n/a	C
GB300 / sm_103a	[81920, 2048] BF16, w FP32, plain, kb	fwd	94.9760	96.2240 TC	95.0730	94.9755	1.0000	1.0131	1.0000	A
GB300 / sm_103a	[81920, 2048] BF16, w FP32, plain, kb	bwd	158.1130	167.1690 TC	187.4255	162.2730	0.9744 -	1.0302	1.1555	A
GB300 / sm_103a	[81920, 2048] BF16, w FP32, plain, kb	2xfwd+bwd	346.1140	356.9940 TC	376.2580	350.0660	0.9887 -	1.0198	1.0743	A
GB300 / sm_103a	[98304, 2048] BF16, w FP32, plain, kb	fwd	113.2800	114.5600 TC	n/a	113.3440	0.9994	1.0107	n/a	B
GB300 / sm_103a	[98304, 2048] BF16, w FP32, plain, kb	bwd	186.3680	202.3365 TC	n/a	194.1445	0.9599 -	1.0422	n/a	B
GB300 / sm_103a	[98304, 2048] BF16, w FP32, plain, kb	2xfwd+bwd	410.3690	428.6400 TC	n/a	417.8245	0.9822 -	1.0259	n/a	B
GB300 / sm_103a	[114688, 2048] BF16, w FP32, plain, kb	fwd	131.6480	132.8955 TC	n/a	131.5200	1.0010	1.0105	n/a	B
GB300 / sm_103a	[114688, 2048] BF16, w FP32, plain, kb	bwd	215.3930	224.3205 TC	n/a	224.1915	0.9608 -	1.0006	n/a	B
GB300 / sm_103a	[114688, 2048] BF16, w FP32, plain, kb	2xfwd+bwd	475.8400	485.4090 TC	n/a	485.2325	0.9806 -	1.0004	n/a	B
GB300 / sm_103a	[131072, 2048] BF16, w FP32, plain, kb	fwd	151.2010	151.8090 TC	151.3290	151.3930	0.9987	1.0027	0.9987	A
GB300 / sm_103a	[131072, 2048] BF16, w FP32, plain, kb	bwd	243.0410	285.8900 TCd	291.0740	252.8010	0.9614 -	1.1309	1.1509	A
GB300 / sm_103a	[131072, 2048] BF16, w FP32, plain, kb	2xfwd+bwd	541.8270	604.2270 TCd	593.1225	553.6670	0.9786 -	1.0913	1.0708	A
GB300 / sm_103a	[49153, 6144] BF16, w FP32, plain, det	fwd	171.7770	172.0970 TC	171.7450	171.7770	1.0000	1.0019	0.9998	A
GB300 / sm_103a	[49153, 6144] BF16, w FP32, plain, det	bwd	278.0810	337.7620 TC	299.3620	278.0810	1.0000	1.2146	1.0780	A
GB300 / sm_103a	[49153, 6144] BF16, w FP32, plain, det	2xfwd+bwd	612.3550	679.0275 TC	637.8920	612.3550	1.0000	1.1089	1.0420	A
GB300 / sm_103a	[57344, 6144] BF16, w FP32, plain, det	fwd	198.8160	199.5840 QA	n/a	198.8160	1.0000	1.0039	n/a	B
GB300 / sm_103a	[57344, 6144] BF16, w FP32, plain, det	bwd	321.1045	343.6170 TC	n/a	321.1045	1.0000	1.0701	n/a	B
GB300 / sm_103a	[57344, 6144] BF16, w FP32, plain, det	2xfwd+bwd	710.6895	751.2810 TC	n/a	710.6895	1.0000	1.0571	n/a	B
GB300 / sm_103a	[65536, 6144] BF16, w FP32, plain, det	fwd	227.1050	228.0010 QA	227.1370	227.1050	1.0000	1.0039	0.9999	A
GB300 / sm_103a	[65536, 6144] BF16, w FP32, plain, det	bwd	366.9785	400.6095 TC	391.1380	366.9785	1.0000	1.0916	1.0660	A
GB300 / sm_103a	[65536, 6144] BF16, w FP32, plain, det	2xfwd+bwd	812.2605	868.2600 TC	839.9080	812.2605	1.0000	1.0689	1.0342	A
GB300 / sm_103a	[81920, 6144] BF16, w FP32, plain, det	fwd	281.3755	282.4960 QA	n/a	281.3755	1.0000	1.0040	n/a	B
GB300 / sm_103a	[81920, 6144] BF16, w FP32, plain, det	bwd	452.8330	505.3130 TCd	n/a	452.8330	1.0000	1.1159	n/a	B
GB300 / sm_103a	[81920, 6144] BF16, w FP32, plain, det	2xfwd+bwd	1008.2255	1085.5705 TCd	n/a	1008.2255	1.0000	1.0767	n/a	B
GB300 / sm_103a	[98304, 6144] BF16, w FP32, plain, kb	fwd	338.4665	339.4095 QA	337.9380	337.9860	1.0014	1.0042	1.0000	A
GB300 / sm_103a	[98304, 6144] BF16, w FP32, plain, kb	bwd	542.3710	620.4510 TC	574.5630	559.4425	0.9695 -	1.1091	1.0279	A
GB300 / sm_103a	[98304, 6144] BF16, w FP32, plain, kb	2xfwd+bwd	1211.4785	1318.5985 TC	1245.5755	1229.9585	0.9850 -	1.0721	1.0128	A
GB300 / sm_103a	[114688, 6144] BF16, w FP32, plain, kb	fwd	391.7115	392.9920 QA	n/a	391.9370	0.9994	1.0027	n/a	B
GB300 / sm_103a	[114688, 6144] BF16, w FP32, plain, kb	bwd	626.9765	598.4330 QA! (2.976)	n/a	649.7455	0.9650 -	0.9210!	n/a	B
GB300 / sm_103a	[114688, 6144] BF16, w FP32, plain, kb	2xfwd+bwd	1404.7860	1380.2270 QA! (2.976)	n/a	1427.8740	0.9838 -	0.9666!	n/a	B
GB300 / sm_103a	[122880, 6144] BF16, w FP32, plain, kb	fwd	419.6010	420.7040 QA	n/a	419.6480	0.9999	1.0025	n/a	B
GB300 / sm_103a	[122880, 6144] BF16, w FP32, plain, kb	bwd	675.8250	639.7455 QA! (2.731)	n/a	693.7920	0.9741 -	0.9221!	n/a	B
GB300 / sm_103a	[122880, 6144] BF16, w FP32, plain, kb	2xfwd+bwd	1509.6185	1477.0905 QA! (2.731)	n/a	1527.4905	0.9883 -	0.9670!	n/a	B
GB300 / sm_103a	[131072, 6144] BF16, w FP32, plain, kb	fwd	449.2815	450.4665 QA	449.0900	449.1540	1.0003	1.0029	0.9999	A
GB300 / sm_103a	[131072, 6144] BF16, w FP32, plain, kb	bwd	715.4435	682.6280 QA! (1.543)	759.1400	727.8110	0.9830 -	0.9379!	1.0440	A
GB300 / sm_103a	[131072, 6144] BF16, w FP32, plain, kb	2xfwd+bwd	1605.1605	1579.2400 QA! (1.543)	1652.2655	1619.4800	0.9912 -	0.9752!	1.0197	A

Cake ships both modes of cake_rmsnorm_train, with all the forms above, on sm_100a, sm_103a and sm_107a. This PR does not change the forward, deterministic=True, the sm_100a or sm_107a routes, or any GB300 window other than the ones above.

Summary by CodeRabbit

• New Features
  • Expanded RMSNorm training support for SM103a GPUs across additional workload sizes and weight precisions.
  • Added prebuilt CUDA 12.9 kernel options for supported configurations.

[yyihuang]

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_aliases.py tests/norm/test_cake_rmsnorm_train.py tests/norm/test_cake_rmsnorm_train_all_sm.py tests/norm/test_cake_rmsnorm_train_api.py tests/norm/test_cake_rmsnorm_train_ckpt.py tests/norm/test_cake_rmsnorm_train_dist.py tests/norm/test_cake_rmsnorm_train_legacy.py tests/norm/test_cake_rmsnorm_train_pairs.py tests/norm/test_cake_rmsnorm_train_stall.py tests/norm/test_cake_rmsnorm_train_vendored.py tests/norm/test_cake_rmsnorm_train_windows.py

[yzh119]

LGTM

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

Adds SM103a CUDA RMSNorm training kernels and TVM-FFI bindings, registers the kernels as AOT libraries, and updates backward program routing and token-count windows for H=1536, H=2048, and H=6144.

Changes: SM103a RMSNorm backward

Layer / File(s)	Summary
CUDA kernels, bindings, and AOT registration flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_*_{binding,kernel}.cu, flashinfer/cake_rmsnorm_train_backend/build.json	Adds CUDA RMSNorm training kernels and TVM-FFI bindings. The bindings validate tensor and launch inputs, configure dynamic shared memory, and launch the kernels. The AOT manifest adds 11 SM103a library records with CUDA 12.9 requirements.
Backward routing and program windows flashinfer/cake_rmsnorm_train_backend/cake_jit.py, flashinfer/cake_rmsnorm_train_backend/cake_windows.py	Adds and orders kblock and deterministic main/reduce program windows for H=1536, H=2048, and H=6144 routes. The two-sum single-kernel selections now begin above 131,072 rows. Updates window edge data, shared-memory wait classifications, notes, and rendered listings.

Estimated code review effort: 4 (Complex) | ~60 minutes

Change: Refactor

Sequence Diagram(s)

sequenceDiagram
  participant CakeJIT
  participant TVMFFIBinding
  participant CUDAKernel
  CakeJIT->>TVMFFIBinding: Select and invoke the routed run entry
  TVMFFIBinding->>CUDAKernel: Validate inputs and launch with tensors and grid

Loading

Possibly related issues

• [CAKE] Activation & Quantization Kernel Progress Tracker #6598: The issue covers Cake RMSNorm kernels, including training, which this change extends with SM103a training kernels and backward routes.

Possibly related PRs

• flashinfer-ai/flashinfer#8810: Extends the same H2048 replacement programs to additional token-count ranges and both weight dtypes.
• flashinfer-ai/flashinfer#8319: Adds an SM103a H1536 FP32-weight kernel that this change includes in the backward route.

Suggested labels: op: misc

Merge Risk: Minimal · up to 402f7

No concrete merge-blocking issue is established for the SM103a backward-routing changes. Merge after normal build and test checks pass.

Pre-merge checks: 4 · 1

Failed checks: 1 inconclusive

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 10.87% 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 describes the blocked FP32 weight-gradient sums added for GB300 backward windows.
Description check		The description covers the change, related issue, accuracy, and performance measurements. It is mostly complete. Add the template headings, pre-commit checklist status, and automated test results. The…
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: Description check

Explanation

The description covers the change, related issue, accuracy, and performance measurements. It is mostly complete. Add the template headings, pre-commit checklist status, and automated test results. The test request does not report results.

Full details: Docstring Coverage

Explanation

Docstring coverage is 10.87% 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 92 functions across 23 files. (2 skipped: 1 unsupported, 1 too large.)

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.

Share

• X
• Mastodon
• Reddit
• LinkedIn

Recent review info

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: dd4a2b4f-4d73-436d-a841-f5f81d850235

Commits

Reviewing files that changed from the base of the PR and between 617a70f and 402f794.

Files selected for processing: 25

• flashinfer/cake_rmsnorm_train_backend/build.json
• flashinfer/cake_rmsnorm_train_backend/cake_jit.py
• flashinfer/cake_rmsnorm_train_backend/cake_windows.py
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_1df502fd0e17091869b7_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_1df502fd0e17091869b7_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_289dc9ab6377fbf87332_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_289dc9ab6377fbf87332_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_2c32330b326a6db3f77b_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_2c32330b326a6db3f77b_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_76e43a7d9c7ea1f59e0d_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_76e43a7d9c7ea1f59e0d_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_87e0b86085e228f738cc_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_87e0b86085e228f738cc_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_882ec880717f382e43eb_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_882ec880717f382e43eb_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_935b8bea9166f29912ff_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_935b8bea9166f29912ff_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_9a47ed21f63c848f2a59_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_9a47ed21f63c848f2a59_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_aa46aa07d86900ad832f_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_aa46aa07d86900ad832f_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_ca6211f3ea2ce8c86a05_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_ca6211f3ea2ce8c86a05_kernel.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_cbeae168d70edd196417_binding.cu
• flashinfer/cake_rmsnorm_train_backend/csrc/cake_rmsnorm_train_cbeae168d70edd196417_kernel.cu

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 4 remain after this review.