[yyihuang]

This PR switches the sm_107a mid forward of the non-deterministic package of Cake's SiLU-gated RMSNorm training kernel (flashinfer.cake_rmsnorm_silu_gate_train with deterministic=False) to cake_rmsnorm_silu_gate_train_87f0173e9583ad4fdd2a. That is the L2-prefetching module sm_100a and sm_103a already run. In it, thread 0 also warms two later row blocks in L2, with one bulk prefetch per tensor, wherever the inputs x and z hold their rows packed. The two blocks are the CTA's next block and the CTA's first block one resident wave later. The sm_107a routing thresholds and grid do not change: the mid forward serves inputs from 800 rows per SM up to the long forward's 1740. The five shapes in Table 1 fall in that range.

The module this replaces, cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850, is the same 128-thread build without the prefetch. Nothing routes to it anymore, so this PR removes it. It is the "previous Cake" baseline below. Every other module and route is FlashInfer main's code: the backward, the short and long forwards, and every sm_100a and sm_103a route run the same code as before.

The PR is one commit on FlashInfer main after PR #8622. The A/B compared two FlashInfer trees: one whose sm_107a mid forward runs without the prefetch, and the same tree plus this PR's sm_107a route. In the measured tree, the new module's sources are byte-identical to this PR's.

A speedup is baseline us / new us, so a value above 1 means this PR is faster. The table gives ranges (min-max) over the five shapes of Table 1 on R200 (sm_107a, 212 SMs), CUDA 13.4. For stepg_f32 and fwd + bwd, the ranges cover the twelve shapes of Table 2: those five, plus seven shapes that run the same code in both trees.

op	det Cake / new	strongest non-Cake / new	previous Cake / new (A/B)
fwd	1.119-1.141	1.071-1.080	1.113-1.136
fwdg	1.131-1.142	1.047-1.077	1.111-1.136
bwdg	1.081-1.163	1.097-1.119	0.9998-1.0002
ckptg	1.115-1.151	1.070-1.084	1.046-1.060
stepg_f32 (12 shapes)	1.068-1.218	-	0.9996-1.061
fwd + bwd (12 shapes)	1.062-1.197	-	0.998-1.045

In the pre-registered A/B, all 20 cells whose code changes (fwd, fwdg, ckptg and stepg_f32 at the five shapes) are faster than the previous Cake route. bwdg runs the same code in both trees. Neither the strongest non-Cake kernel (Table 1) nor the deterministic Cake kernel (Tables 1 and 2) is faster than new on any row. Every cell below 1x is a same-code cell against the previous Cake route, each by less than 0.3 %, and every one is listed under "Cells below 1x".

The kernel's other Cake variants are the deterministic package flashinfer.cake_rmsnorm_silu_gate_det_train (the same API with deterministic=True) and the FP32 weight-gradient output of the non-deterministic package. Both packages ship on sm_100a, sm_103a and sm_107a. sm_100a and sm_103a have run this mid forward since PR #8622, and this PR changes only the non-deterministic package's mid forward on sm_107a.

How the tables were measured

• Latency: CUPTI kernel time with cold L2, taken with FlashInfer's GPU timer in flashinfer.testing. Inputs are BF16, and all times are in us.
• Table 1: the deterministic Cake kernel, the strongest non-Cake kernel and new were measured in the same process. Each value is the median of the per-round medians over 10 rounds.
• A/B values: the previous Cake and new (A/B) columns and all of Table 2 are medians of the pre-registered interleaved A/B, which ran 16 sequences and 160 processes. stepg_f32 has no same-process measurement, so its deterministic Cake values also come from the A/B.
• Strongest non-Cake kernel per row: the fastest of the following, with torch 2.14.1+cu130 and Triton 3.8.0.
  • FLA 0.6.0 rms_norm_gated (fla.modules.fused_norm_gate).
  • mamba_ssm v2.3.2 rmsnorm_fn (layernorm_gated).
  • torch.compile in default and max-autotune-no-cudagraphs modes, best of 5 fresh compiles each. #k names the fastest compile.
  • The inductor-generated persistent-reduction Triton kernel of the FP32 formula, launched directly.
• Deterministic Cake: flashinfer.cake_rmsnorm_silu_gate_det_train, unchanged by this PR.
• Operations:
  • fwd is one eager forward call, and fwdg is the forward replayed from a CUDA graph.
  • bwdg is the backward replayed from a CUDA graph.
  • ckptg is the checkpointed step (2 x forward + backward) in one CUDA graph, and stepg_f32 is the same step with an FP32 dw.
  • fwd + bwd is fwdg + bwdg, derived.
• In Table 2, the first five shapes run the new module. 260x48, 1031x48, 2048x48, 3533x48, 7685x48, 10600x48 and 16384x48 run the same code in both trees.

Cells below 1x

• vs the strongest non-Cake kernel (Table 1): none.
• vs the deterministic Cake kernel: none.
• vs the previous Cake route (A/B), only at cells that run the same code in both trees:
  • bwdg (Table 1): 8192x32 0.9998, 7684x48 0.99999
  • Table 2: 1031x48 stepg_f32 0.9997, 1031x48 fwd + bwd 0.998, 2048x48 stepg_f32 0.9996, 3533x48 ckptg 0.9996, 3533x48 stepg_f32 0.9999, 7685x48 stepg_f32 0.9999, 7685x48 fwd + bwd 0.9998, 10600x48 stepg_f32 0.99996, 16384x48 fwd + bwd 0.9999

Results on R200 (sm_107a, 212 SMs)

• GPU: R200 (VR200 node; nvidia-smi name 'NVIDIA Graphics Device'), sm_107a, compute capability 10.7, 212 SMs, driver 620.43.
• CUDA: 13.4 (nvcc Build cuda_13.4.r13.4/compiler.38855100_0); PyTorch 2.14.1+cu130 (CUDA 13.0); Triton 3.8.0; Python 3.14.8.

Table 1 (R200, sm_107a, CUDA 13.4): new vs the deterministic Cake kernel, the strongest non-Cake kernel and the previous Cake route

T x H	op	det Cake us	strongest non-Cake us	strongest non-Cake kernel	new us	previous Cake us (A/B)	new us (A/B)	det / new	non-Cake / new	previous / new (A/B)
3534x48	fwd	18.192	17.408	torch.compile default #1	16.256	18.272	16.241	1.119	1.071	1.125
3534x48	fwdg	18.256	16.896	torch.compile default #3	16.144	18.088	16.096	1.131	1.047	1.124
3534x48	bwdg	38.528	37.072	FLA rms_norm_gated	33.136	33.088	33.088	1.163	1.119	1.00001
3534x48	ckptg	72.831	68.273	FLA rms_norm_gated	63.264	66.264	63.048	1.151	1.079	1.051
3534x48	stepg_f32	72.576 (A/B)	not measured	-	-	66.392	63.392	1.145 (A/B)	n/a	1.047
4097x48	fwd	20.416	19.328	FLA rms_norm_gated	17.889	19.904	17.888	1.141	1.080	1.113
4097x48	fwdg	20.192	19.104	torch.compile max-autotune-no-cudagraphs #2	17.825	19.872	17.880	1.133	1.072	1.111
4097x48	bwdg	40.064	40.689	FLA rms_norm_gated	37.057	37.128	37.120	1.081	1.098	1.0002
4097x48	ckptg	78.544	75.376	FLA rms_norm_gated	70.464	73.808	70.536	1.115	1.070	1.046
4097x48	stepg_f32	78.440 (A/B)	not measured	-	-	73.816	70.849	1.107 (A/B)	n/a	1.042
8192x32	fwd	26.304	24.928	FLA rms_norm_gated	23.072	25.961	22.944	1.140	1.080	1.131
8192x32	fwdg	26.272	24.784	FLA rms_norm_gated	23.009	25.961	22.928	1.142	1.077	1.132
8192x32	bwdg	50.352	50.608	FLA rms_norm_gated	46.048	45.992	46.000	1.093	1.099	0.9998
8192x32	ckptg	100.273	96.256	FLA rms_norm_gated	89.616	94.256	89.256	1.119	1.074	1.056
8192x32	stepg_f32	99.885 (A/B)	not measured	-	-	94.256	89.425	1.117 (A/B)	n/a	1.054
6000x48	fwd	28.096	26.624	torch.compile default #1	24.673	28.064	24.704	1.139	1.079	1.136
6000x48	fwdg	28.048	26.417	torch.compile default #3	24.640	27.872	24.544	1.138	1.072	1.136
6000x48	bwdg	54.928	54.529	FLA rms_norm_gated	49.328	49.184	49.184	1.114	1.105	1.000003
6000x48	ckptg	108.096	103.552	FLA rms_norm_gated	95.904	101.608	95.824	1.127	1.080	1.060
6000x48	stepg_f32	108.769 (A/B)	not measured	-	-	101.760	95.945	1.134 (A/B)	n/a	1.061
7684x48	fwd	35.008	33.168	FLA rms_norm_gated	30.720	34.712	30.816	1.140	1.080	1.126
7684x48	fwdg	35.008	32.961	torch.compile default #5	30.688	34.537	30.656	1.141	1.074	1.127
7684x48	bwdg	66.208	65.648	FLA rms_norm_gated	59.824	59.720	59.720	1.107	1.097	0.99999
7684x48	ckptg	133.120	128.496	FLA rms_norm_gated	118.560	125.296	118.361	1.123	1.084	1.059
7684x48	stepg_f32	132.496 (A/B)	not measured	-	-	125.569	118.705	1.116 (A/B)	n/a	1.058

Table 2 (R200, sm_107a, CUDA 13.4): full step, previous Cake route vs new (A/B medians)

T x H	rows	op	previous Cake us	new us	previous / new	det Cake us	det / new
3534x48	169632	ckptg	66.264	63.048	1.051	72.609	1.152
3534x48	169632	stepg_f32	66.392	63.392	1.047	72.576	1.145
3534x48	169632	fwd + bwd	51.176	49.184	1.041	56.833	1.156
4097x48	196656	ckptg	73.808	70.536	1.046	78.449	1.112
4097x48	196656	stepg_f32	73.816	70.849	1.042	78.440	1.107
4097x48	196656	fwd + bwd	57.000	55.000	1.036	60.368	1.098
8192x32	262144	ckptg	94.256	89.256	1.056	100.032	1.121
8192x32	262144	stepg_f32	94.256	89.425	1.054	99.885	1.117
8192x32	262144	fwd + bwd	71.953	68.929	1.044	76.656	1.112
6000x48	288000	ckptg	101.608	95.824	1.060	108.872	1.136
6000x48	288000	stepg_f32	101.760	95.945	1.061	108.769	1.134
6000x48	288000	fwd + bwd	77.056	73.728	1.045	83.328	1.130
7684x48	368832	ckptg	125.296	118.361	1.059	132.564	1.120
7684x48	368832	stepg_f32	125.569	118.705	1.058	132.496	1.116
7684x48	368832	fwd + bwd	94.257	90.376	1.043	100.800	1.115
260x48	12480	ckptg	14.704	14.704	1.000	18.064	1.229
260x48	12480	stepg_f32	14.352	14.344	1.001	17.472	1.218
260x48	12480	fwd + bwd	12.664	12.664	1.00002	15.160	1.197
1031x48	49488	ckptg	26.592	26.584	1.0003	29.768	1.120
1031x48	49488	stepg_f32	26.624	26.632	0.9997	29.656	1.114
1031x48	49488	fwd + bwd	22.224	22.272	0.998	26.624	1.195
2048x48	98304	ckptg	43.328	43.328	1.000	46.537	1.074
2048x48	98304	stepg_f32	42.832	42.848	0.9996	45.896	1.071
2048x48	98304	fwd + bwd	33.952	33.952	1.00001	37.968	1.118
3533x48	169584	ckptg	67.192	67.216	0.9996	72.520	1.079
3533x48	169584	stepg_f32	67.152	67.156	0.9999	72.496	1.080
3533x48	169584	fwd + bwd	51.280	51.272	1.0002	56.736	1.107
7685x48	368880	ckptg	121.057	121.009	1.0004	132.776	1.097
7685x48	368880	stepg_f32	121.136	121.152	0.9999	132.832	1.096
7685x48	368880	fwd + bwd	92.017	92.032	0.9998	100.688	1.094
10600x48	508800	ckptg	162.752	162.692	1.0004	173.993	1.069
10600x48	508800	stepg_f32	162.858	162.865	0.99996	173.985	1.068
10600x48	508800	fwd + bwd	123.296	123.280	1.0001	131.025	1.063
16384x48	786432	ckptg	236.545	236.453	1.0004	254.433	1.076
16384x48	786432	stepg_f32	236.737	236.643	1.0004	254.441	1.075
16384x48	786432	fwd + bwd	177.737	177.761	0.9999	188.721	1.062

Summary by CodeRabbit

• Performance
  • Updated RMSNorm-SiLU-gate training execution for newer GPU architectures, including sm_107a. The revised path supports a lower minimum workload threshold on that architecture and applies bulk L2 prefetching across supported architectures.
• Documentation
  • Clarified the description of mid-forward L2 prefetching.

[yyihuang]

@flashinfer-bot run tests/norm/test_cake_rmsnorm_silu_gate_det_train.py tests/norm/test_cake_rmsnorm_silu_gate_det_train_aliases.py tests/norm/test_cake_rmsnorm_silu_gate_det_train_api.py tests/norm/test_cake_rmsnorm_silu_gate_det_train_stall.py tests/norm/test_cake_rmsnorm_silu_gate_train.py tests/norm/test_cake_rmsnorm_silu_gate_train_api.py tests/norm/test_cake_rmsnorm_silu_gate_train_stall.py

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

The sm_107a mid-forward route now uses the shared module for sm_100a, sm_103a, and sm_107a. The module adds bulk L2 prefetching and an sm_107a row threshold. Build metadata and generated CUDA sources reflect the module change.

Changes: Mid-forward module consolidation

Layer / File(s)	Summary
Shared mid-forward behavior and routing flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py, flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py	The shared mid-forward module now covers sm_100a, sm_103a, and sm_107a. It prefetches the CTA’s next packed-row block and a block one wave later. The row threshold is 976 for sm_100a and sm_103a, and 800 for sm_107a. The sm_107a route selects this module.
sm_107a artifact transition flashinfer/cake_rmsnorm_silu_gate_train_backend/build.json, flashinfer/cake_rmsnorm_silu_gate_train_backend/csrc/...	Build metadata adds the shared module’s sm_107a artifact and removes the previous artifact record. The previous module’s generated binding and kernel source files are deleted.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Refactor

Merge Risk: Minimal · up to 80acb

No actionable merge-blocking issue was identified in the SM107 module consolidation. Merge readiness remains subject to normal build and test checks.

Pre-merge checks: 5

Check name	Status	Explanation
Description check		The description follows the repository template in substance. It has a clear Description section that explains the change, the reason for it, and the measured results. It covers the removed module and…
Docstring Coverage		No eligible functions found in the changed files. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Title check		The title names the main change: the SM107 mid-forward path of the Cake RMSNorm-SiLU-gate training kernel now uses the L2-prefetching module, as the raw summary (cake_jit.py now serves sm_107a from th…

Full details: Description check

Explanation

The description follows the repository template in substance. It has a clear Description section that explains the change, the reason for it, and the measured results. It covers the removed module and the unchanged routes. It includes tables of measurements, a list of cells below 1x, and the hardware and software versions. The template's checklist items (pre-commit, tests) are not ticked, and the Related Issues section is empty. The author did request specific test targets in the PR comments, but the description does not state whether the tests were run or passed. The experimental section is correctly left untouched because this is not an experimental PR. The description is mostly complete, so the check passes, though the missing Related Issues link and the unchecked test and pre-commit items are noted for the author.

Full details: Docstring Coverage

Explanation

No eligible functions found in the changed files. Skipping docstring coverage check. 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 0 functions across 2 files. (1 skipped: 1 unsupported.)

Full details: Title check

Explanation

The title names the main change: the SM107 mid-forward path of the Cake RMSNorm-SiLU-gate training kernel now uses the L2-prefetching module, as the raw summary (cake_jit.py now serves sm_107a from the shared module, and the old sm_107a-only module is removed) and the PR objectives both describe. The 'perf' type prefix matches the stated goal of faster measured timings. The title is specific and short, and it does not describe anything missing from the changeset.

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: 3fb69628-d646-4d68-861d-fc7c73e18de7

Commits

Reviewing files that changed from the base of the PR and between bc4e12b and 80acbaf.

Files selected for processing: 5

• flashinfer/cake_rmsnorm_silu_gate_train_backend/__init__.py
• flashinfer/cake_rmsnorm_silu_gate_train_backend/build.json
• flashinfer/cake_rmsnorm_silu_gate_train_backend/cake_jit.py
• flashinfer/cake_rmsnorm_silu_gate_train_backend/csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_binding.cu
• flashinfer/cake_rmsnorm_silu_gate_train_backend/csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_kernel.cu

Files with no reviewable changes: 2

• flashinfer/cake_rmsnorm_silu_gate_train_backend/csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_binding.cu
• flashinfer/cake_rmsnorm_silu_gate_train_backend/csrc/cake_rmsnorm_silu_gate_train_b0eb3df9844d5b734850_kernel.cu

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