Repository navigation
fix(cake_rmsnorm_det): tighten two large-row weight-gradient accuracy edges - #8293
Conversation
…y edges The compensated backward programs now start at lower row counts on two routes: the H = 4096 in-chunk compensated form above 44032 rows on SM100 (was 52736) and the compensated-dot H = 256 ring above 99840 rows on SM103 (was 114688). Below the old edges the plain programs' weight gradient came closer than the 0.8 routing margin to the per-element bound. cake_jit.py moves the window records of the two edges, cake_windows.py follows and records the measured class of the accuracy edges that stay, and the pairs test covers both sides of the new edges. No program changes.
|
@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_aliases.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py |
|
No actionable comments were generated in the recent review.
|
| Layer / File(s) | Summary |
|---|---|
Update backward route windows flashinfer/cake_rmsnorm_det_train_backend/cake_jit.py, flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py |
SM100 H=4096 boundaries change from 52,736 to 44,032 rows. SM103 H=256 boundaries change from 114,688 to 99,840 rows. SM100 H=6144 uses the bulk-copy program above 4,096 rows. |
Update edge metadata and test expectations flashinfer/cake_rmsnorm_det_train_backend/cake_windows.py, tests/norm/test_cake_rmsnorm_det_train_pairs.py |
Several backward edges change from unmeasured to same. The displayed windows and accuracy-edge expectations reflect the updated boundaries, and the test row set includes additional edge values. |
Estimated code review effort: 2 (Simple) | ~10 minutes
Change: Bug fix
Possibly related issues
flashinfer-ai/flashinfer#7045— The issue concerns row-dependent RMSNorm training program selection and accuracy changes; this PR updates backward row thresholds and accuracy-edge expectations.
Suggested reviewers: yzh119
Merge Risk: 
Minimal · up to a2ad7
-
Selected backward routes switch to compensated kernels earlier. No actionable merge-blocking risk was identified.
Pre-merge checks: 
5
-
Full details: Title check
-
Explanation
The title names the main change: tightening two large-row weight-gradient accuracy edges in cake_rmsnorm_det. The raw summary and PR objectives both describe moving the H=4096 SM100 and H=256 SM103 row thresholds, which matches this title. The title follows the conventional commit format and is specific enough for a reviewer scanning history.
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 3 files.
Full details: Description check
-
Explanation
The description clearly explains the routing changes, rationale, affected hardware and shapes, performance impact, and scope. It omits the related-issues section, checklist confirmations, explicit test results, and experimental-track section, but the core pull request information is complete and on topic.
-
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.
yzh119
left a comment
There was a problem hiding this comment.
LGTM. Edge renames are consistent across cake_jit, cake_windows and the tests. The PR description is only "Speedup table follows." and there is no table yet, so please add it.
|
@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_aliases.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py |
…#8253) Follow-up to #8208 for the opt-in non-deterministic weight gradient (deterministic=False) of flashinfer.cake_rmsnorm_train. Forward: on every architecture, route and token count the package now runs the program the deterministic package selects (its programs and windows rendered into this package), so y, rstd and the residual output are bitwise equal to deterministic=True. In #8208 the H = 256 routes and, on SM100 and SM103, the H = 5120 routes ran forward windows and programs that differ from the deterministic package's. Why the backward windows move: in #8208 dw elements of five routed deterministic pairs besides the SM107 H = 5120 pair (below) exceed the per-element budget (atol 1e-4 + rtol 1e-4 against FP64 at the kernel's FP32 rstd) at sampled token counts up to 65536: on GB300 (SM103) H = 4096 at T 50688 (1.032 of the budget) and the FP32-weight H = 6144 route at T 56320 (1.010; 1.085 at T 54273 against FP64 at torch's FP32 rstd), and on a 212-SM R200 (SM107) both H = 6144 routes at T 31040 (1.057) and the FP32-weight H = 1536 route at T 61952 (1.020). Up to T 65536, six more routed pair windows reach more than 0.85 of the budget with inputs of their own weight dtype: on B200 (SM100) H = 2048 (0.883 at T 59392), H = 4096 (0.923 at T 52737) and its FP32-weight route (0.909 at T 44544 and 0.916 at T 65024; 0.983 at T 63488 against FP64 at torch's FP32 rstd), on GB300 the FP32-weight H = 2048 (0.904 at T 62464) and H = 5120 (0.888 at T 62464) routes, and on a 212-SM R200 H = 1536 (0.980 at T 62976). On a 212-SM R200 the BF16-weight form of the FP32-weight H = 2048 pair, the same dw computation, reaches 0.900 of the budget at T 51200 (0.873 at T 55296 on a 216-SM R200), and on GB300 the FP32-weight form of the BF16-weight H = 1536 pair, the same dw computation, reaches 0.890 at T 47104. The worst dw error of seven FP32-weight pairs exceeds the FP32 eager reduction's: on a 212-SM R200 H = 512 at T 1088, 1664, 1920 and 2368, H = 4096 at T 24576 (0.7328 against 0.7241), H = 6144 at T 30784 (0.8671 against 0.8066) and H = 256 at 2048 rows (512 tokens of 4 heads; 0.1435 against 0.1158), on B200 H = 1536 at T 33280 (0.6840 against 0.6219), and on GB300 H = 1536 at T 38400 (0.7399 against 0.7092) and H = 4096 at T 9728 (0.3383 against 0.3332). On a 216-SM R200 the worst dw error of the SM107 FP32-weight H = 6144 pair exceeds the eager reduction's at T 29184 (0.8224 against 0.8121), and that of the SM107 FP32-weight H = 4096 pair that ran from T 1 to 256 at T 64 (0.0256 against 0.0233). On a 212-SM R200 the median dw error of that H = 4096 pair also exceeds the eager reduction's at T 4 (0.0024 against 0.0022), and its worst at T 8 and 16 (0.0079 against 0.0078 and 0.0109 against 0.0106) and, on the input family for which the BF16-weight H = 4096 pair computes the same dw (measured with that pair), at T 5, 6, 9 and 20 (0.0043 against 0.0037, 0.0063 against 0.0049, 0.0071 against 0.0070 and 0.0119 against 0.0108). Backward windows: the single kernel (FP64 dw column sums; twosum per-thread sums on SM103 and on SM107 FP32-weight H = 2048) now runs above - SM100: T 58880 on H = 2048, T 44032 on H = 4096 and T 32768 on FP32-weight H = 1536; - SM103: T 49152 on H = 6144, T 61952 on H = 5120 and H = 2048, T 46592 on BF16-weight and T 37888 on FP32-weight H = 1536, and T 50176 on BF16-weight and T 9216 on FP32-weight H = 4096; - SM107: T 30720 on BF16-weight and T 28672 on FP32-weight H = 6144, T 53248 on H = 1536, T 24064 on FP32-weight H = 4096, T 50176 on FP32-weight H = 2048 (where the BF16-weight route already ends) and 1536 rows on FP32-weight H = 256; FP32-weight H = 512 runs it above T 1024 up to T 4096, and FP32-weight H = 4096 also from T 1 to 256 (the single kernel of that route's top window). Each move covers both weight dtypes unless a weight is named: the BF16- and FP32-weight deterministic pairs compute the same dw for equal inputs, so a failing input of one is a failing input of the other. A worst error above the FP32 eager reduction's (read on the FP32-weight, H = 5120 and H = 256 routes) moves only the route on which it occurs. SM103 H = 2048 (its twosum form) and SM107 H = 1536 get new single kernel programs, builds of the same kernel sources; SM107 FP32-weight H = 512 runs a single kernel program this package ships for SM100, now also built for SM107; the other moved windows run single kernel programs this package already ships for their architecture. The main kernels of the SM100 FP32-weight H = 1536 pair that ran from T 32768 to 36864 and of the SM107 FP32-weight H = 4096 pair that ran from T 1 to 256 are no longer shipped. Small token counts: a survey of the routed programs of every backward route whose lowest window ran a deterministic pair in #8208, other than the SM107 FP32-weight H = 4096 route (above), at every T from 1 to 64 and T 96, 127 to 129, 192, 255 and 256 on B200, GB300 and a 212-SM R200, and on the R200 also at every T up to 256 on H = 512, H = 2048 and H = 256 ([T, 256] input), with 60 inputs of each of three families and more token counts sampled near and above each cut, finds on all 35 of these routes that the pair's median dw error over all inputs, or on the FP32-weight, H = 5120 and H = 256 routes also its worst dw error, exceeds the FP32 eager reduction's at some sampled T (on four of them only at T 3 and by less than 0.0001 of the budget). On every one of these routes a single kernel now runs from T 1 up to a cut and the pair above it: the single kernel is no worse than the eager reduction by those readings at every sampled T up to the cut, and the pair at every sampled T above it (on SM103 FP32-weight H = 1536, over the 60 inputs of each family, the pair's worst dw error exceeds the eager reduction's at T 320, 0.0596 against 0.0583 of the budget, so its cut is 320). On 31 of these routes the single kernel is one that the route also runs in a higher window (on 9 of them the twosum form). On FP32-weight H = 512 of SM100 and SM103 no single kernel the route ran passes there (SM103 ran none): over the 60 inputs of each family an FP64-accumulating single kernel of the same tiling, on SM100 the one that route runs above T 960, exceeds the eager reduction's worst dw error at T 4 (0.0033 against 0.0032), so they run a twosum single kernel, new programs; on BF16-weight H = 512 SM103 and SM107 ran no single kernel, and they run an FP64-accumulating one, a new program on SM103 and on SM107 the program this package runs on SM100 H = 512. The cuts (in tokens; in rows on H = 256; each cut is set for its weight dtype alone): - SM100: BF16 weight 14 on H = 1536 and 26 on H = 512; FP32 weight 64 on H = 1536, 22 on H = 5120 and 292 on H = 512; - SM103: BF16 weight 14 on H = 1536, 12 on H = 2048, 1080 rows on H = 256, 22 on H = 5120, 26 on H = 512 and 3 on H = 4096 and H = 6144; FP32 weight 320 on H = 1536, 42 on H = 2048, 1080 rows on H = 256, 15 on H = 4096, 22 on H = 5120, 20 on H = 6144, 6 on residual H = 6144 and 292 on H = 512; - SM107: BF16 weight 14 on H = 1536, 6 on H = 2048, 1056 rows on H = 256, 6 on H = 4096, 18 on H = 5120, 77 on H = 512 and 3 on H = 6144 and residual H = 6144; FP32 weight 64 on H = 1536, 127 on H = 2048, 1296 rows on H = 256, 229 on H = 512, 13 on H = 5120, 20 on H = 6144 and 3 on residual H = 6144. Cost: the single kernel's median backward time against the deterministic pair it replaces, over the timing sessions and their pooled medians (CUDA graph, L2 flushed before every replay): - B200 (148 SMs): 1.001-1.012x on H = 2048 and 0.997-1.005x on FP32-weight H = 2048 (T 58881 to 65536); 0.983-1.064x on H = 4096 and 0.980-1.043x on FP32-weight H = 4096 (T 44033 to 65536); 0.941-0.965x on FP32-weight H = 1536 (T 32769 to 36864); - GB300 (152 SMs): 1.297-1.311x on BF16-weight H = 1536 (T 46593 to 53248) and 1.322-1.336x on FP32-weight H = 1536 (T 37889 to 45056); 1.118-1.137x on H = 2048 and 1.183-1.200x on FP32-weight H = 2048 (T 61953 to 65536); 1.019-1.034x on BF16-weight H = 4096 (T 50177 to 63488) and 1.013-1.093x on FP32-weight H = 4096 (T 9217 to 23552); 1.136-1.154x on H = 5120 and 1.167-1.183x on FP32-weight H = 5120 (T 61953 to 65536); 1.051-1.061x on H = 6144 (T 49153 to 51200) and 1.062-1.079x on FP32-weight H = 6144 (T 49153 to 60416); - R200 (212 SMs): 1.586-1.600x on H = 6144 (T 30721 to 31744) and 1.543-1.566x on FP32-weight H = 6144 (T 30721 to 31744; T 28673 to 30720 not timed); 1.406-1.463x on H = 1536 and 1.415-1.478x on FP32-weight H = 1536 (T 53249 to 65536); 1.340-1.385x on FP32-weight H = 4096 (T 24576 to 28672; T 1 to 256 not timed), 1.313-1.423x on FP32-weight H = 512 (T 1088 to 4096), 1.233-1.267x on FP32-weight H = 2048 (T 50688 to 56320) and 1.102-1.165x on FP32-weight H = 256 (1537 to 3584 rows). Small-T cost: the single kernel's median backward time against the deterministic pair's in each of two timing sessions, and on GB300 also over the samples of both sessions (CUDA graph, L2 flushed before every replay, memory clock checked at the board maximum), at T 1, 4, 8 and 16 where those are at most the cut and at the cut (on H = 256 also at the cut in the 4-head and 24-head views; on the T 3 cuts at T 1 to 3): - B200 (148 SMs): 1.205-1.367x on H = 1536 and 0.952-1.150x on H = 512; with FP32 weight 1.157-1.304x on H = 1536, 1.044-1.132x on H = 512 and 1.610-1.872x on H = 5120. - GB300 (152 SMs): 1.889-2.088x on H = 1536, 1.890-2.105x on H = 2048, 1.014-1.276x on H = 256, 2.672-32.000x on H = 5120, 1.493-1.657x on H = 512, 2.502-2.592x on H = 4096 and 12.004-35.083x on H = 6144 (six sessions on three GPUs); with FP32 weight 1.904-2.019x on H = 1536 at T 1, 4, 8, 16 and 64 (four sessions at T 1 to 16) and 2.494-2.563x at T 320, 1.941-2.224x on H = 2048, 0.995-1.286x on H = 256, 1.929-2.139x on H = 512, 2.416-2.590x on H = 4096, 2.636-2.805x on H = 5120, 2.844-32.895x on H = 6144 and 6.149-34.422x on residual H = 6144. The largest ratios are at T 1, where the single kernels take 174.640 us (H = 5120), 212.128 to 212.353 us (H = 6144, session medians), 209.503 us (FP32-weight H = 6144) and 207.376 us (FP32-weight residual H = 6144) against 5.872, 6.048 to 6.336, 6.464 and 6.144 us for the pairs; the median time of every timed single kernel is at most 0.22 ms. - R200: not timed. Backward, SM107 H = 5120 (both weight dtypes): the compensated deterministic pair that #8208 runs from T 19968 now runs up to T 49152 and the FP64-accumulating single kernel above it. In #8208 the pair ran up to T 65536 (BF16 weight) and T 57344 (FP32 weight). On a 212-SM R200 one BF16-weight dw element of that pair exceeds the budget at T 61440 (1.110 of it); on a 216-SM R200 its worst element reaches 0.867 of the budget at T 53248 and 0.891 at T 61440 with BF16 weight, and 0.905 at T 57344 with FP32 weight (worst of 200 inputs of each of three families). Above T 49152 the single kernels stay within 0.703 (BF16 weight) and 0.646 (FP32 weight) of the budget at the sampled token counts up to 65537. Their median backward time is 1.80-1.87x the pair's on a 212-SM R200 (T 53248 to 65536). Since #8251 the deterministic package runs other programs on SM107 H = 5120 from T 19968, so from T 19968 to 49152 deterministic=True and deterministic=False run different deterministic pairs. Backward, SM103 H = 256 (both weight dtypes): above 5120 rows each row count runs a deterministic pair instead of this package's single kernels, except from 93184 to 139264 rows, where the twosum single kernel of #8208 keeps running. The pair's main kernel is the one the deterministic package runs for that row count on a device with 152 or more SMs; its dw reduce is that package's up to 40960 rows and above 1572864 rows, and from 40960 to 93184 and from 139264 to 1572864 rows the one that package ran there before #8319 (the measured pairs). From #8293 to #8319 the deterministic package ran the pair of 40960 to 93184 rows up to 99840 rows; that pair reaches 0.9807 of the budget at 93696 rows (3904 tokens of 24 heads, FP32 weight, worst of 60 inputs of each of three families). In that time it ran another pair from 99840 rows, whose worst dw element stays within 0.353 (BF16 weight) and 0.373 (FP32 weight) of the budget at the sampled row counts up to 139264 rows (worst of 200 inputs of each of three families); this package keeps the twosum single kernel up to 139264 rows. From 139264 to 1572864 rows, where this package now runs that pair in place of the twosum single kernel, the twosum single kernel's median backward time on GB300 is 1.058-1.095x (1.060-1.093x pooled) that pair's at every timed row count; below 139264 rows the two are not timed against each other. Where the pairs run, their worst dw element stays within 0.837 of the budget at every sampled row count, and on GB300 their median backward time is 0.820-0.989x (0.829-0.972x pooled) that of the single kernels they replace. From 93184 to 139264 rows the twosum single kernel's median backward time is 1.130-1.188x that of the pair of 40960 to 93184 rows on GB300. Up to 1080 rows the single kernel runs (small token counts, above); from 1081 to 5120 rows nothing changes. Accuracy of dw with deterministic=False at this head, against FP64 at the kernel's FP32 rstd and at torch's FP32 rstd with the per-element budget atol 1e-4 + rtol 1e-4, worst over 60 or more inputs of each of three input families, BF16 and FP32 weight, at every sampled token count up to 65536, the top of the qualified range (every sampled row count of the H = 256 views): the routed deterministic pairs stay within 0.850 of the budget (B200 0.842; GB300 0.850; R200 0.832 with 212 SMs and 0.844 with 216 SMs) and the single kernels within 0.751 (B200 0.713; GB300 0.572, the twosum forms 0.0006; R200 0.727 with 212 SMs and 0.751 with 216 SMs). Over the same samples the median dw error over the inputs is no worse than the FP32 eager reduction's at every sampled token count where both were recorded; above T 256 the median over the inputs of each family is no worse either, except one family of the SM103 FP32-weight H = 512 route at T 1088 (0.0395 against 0.0376). On the H = 5120, H = 256 and FP32-weight routes the worst dw error over those inputs is also no worse than the FP32 eager reduction's at every sampled token count on B200, GB300 and R200 (212 and 216 SMs). The window tables record the measured outcome at each of the 74 new backward edges (10 on SM100, 37 on SM103 and 27 on SM107; dx is bitwise identical across every one) and, for each of the 130 new backward windows (19, 64 and 47), the smallest side allocation (100 or 132 KB per CTA) beside which it waits for a kernel resident on every SM, or that it waits beside neither, measured on B200, GB300 and a 212-SM R200. Beside 100 KB per CTA the SM103 H = 6144 single kernels above T 49152 wait, where the pairs they replace waited from 132 KB. Other new windows that wait beside a smaller allocation than the program that ran their token counts in #8208: SM100 bwd_w32_h5120 (0, 22]; SM103 none; SM107 none. The registry's window-reason comments state measured values, except the two SM100 H = 5120 comments above T 65536, which name the top of the qualified range; the H = 5120 test comment states measured values. The window notes name the small-T windows and the twosum windows from T 1. The docstrings say that deterministic=True reproduces dw bitwise for a given token count on devices with the same SM count. The mixed-form workspace test in test_cake_rmsnorm_train_pairs.py takes the first BF16-weight route without residual of SM100, SM103 or SM107 whose window at 257 tokens runs the deterministic pair, and discounts the route's own single-kernel records up to 256 tokens: no route runs the deterministic pair at 1 token now.
This PR changes the backward pass of the deterministic Cake RMSNorm training package (
cake_rmsnorm_det_train), which covers BF16 and FP32 weights, residual and non-residual inputs, row widths H 256 (per-head view) / 512 / 1536 / 2048 / 4096 / 5120 / 6144, and both forward and backward passes. It adds no kernels and lowers two accuracy edges: sm_100a H 4096 from 52736 to 44032 rows, and sm_103a H 256 from 114688 to 99840 rows. Rows between the new and old edges now use the compensated dw program that the package already used above the old edges, because at some row counts just below the old edges the previous program came too close to the per-element dw tolerance. On sm_100a H 4096, the moved rows run at 0.994x-1.019x of the previously shipped Cake package, within the measurement noise band of about 1.4%. On sm_103a H 256 (contiguous rows, key [T,4,256] and query [T,24,256] views), the moved rows are slower than the previously shipped Cake package, at 0.911x-0.943x (worst 0.911x, key view with FP32 weight at 107264 rows), and rows at the new edge itself are unchanged (0.997x-1.007x). Every row stays faster than torch.compile, the fastest non-Cake baseline: 1.158x-1.367x on sm_100a H 4096 and 1.054x-1.206x on sm_103a H 256.NVIDIA B200 (148 SMs), sm_100a; CUDA 13.4 (kernels built with nvcc 13.4), PyTorch 2.14.1+cu130, Triton 3.8.0. Backward pass. Latency = CUPTI kernel-active time per call, cold L2, median of 4 x 20 calls per process. Previous Cake = the previously shipped Cake package; Cake = this package; both through the package's public backward entry point in the same processes (counterbalanced). Each Cake value = fastest measured process over the row's formal sessions (4 sessions on 7 rows, 6 sessions on 3 rows; each session a fresh process set with fresh caches). torch.compile = fastest measured per-shape static compile (FP32 formula, dynamic=False; minimum over every process of the formal sessions and a tc-only sweep, each with a fresh inductor / Triton cache). Speedup = baseline / Cake.
NVIDIA GB300 (152 SMs), sm_103a; CUDA 13.4 (kernels built with nvcc 13.4), PyTorch 2.14.1+cu130, Triton 3.8.0. Backward pass. Latency = CUPTI kernel-active time per call, cold L2, median of 4 x 20 calls per process. Previous Cake = the previously shipped Cake package; Cake = this package; both through the package's public backward entry point in the same processes (counterbalanced). Each Cake value = fastest measured process over the row's formal sessions (4 sessions; each session a fresh process set with fresh caches). torch.compile = fastest measured per-shape static compile (FP32 formula, dynamic=False; minimum over every process of the formal sessions and a tc-only sweep, each with a fresh inductor / Triton cache). Speedup = baseline / Cake.
The other Cake variants of this family are unchanged: forward, residual backward, every other row width and route, and sm_107a; the family ships on sm_100a, sm_103a and sm_107a.