Repository navigation
Conversation
Re-vendor both Cake RMSNorm packages from flashinfer-ai/flashinfer c51a7540f042d7607f23bc29417147ca7a921828 (main, the merge of #8848). Since the previous commit, eb3eb68d, these changes touched them: - flashinfer-ai/flashinfer#8810: on SM103a the non-deterministic backward of H = 2048 ([T, H] rows, BF16 and FP32 weight) runs the deterministic package's program from T 58497 to 61952. FlashInfer measured the pair it ran there before above the weight-gradient budget (every element within atol 1e-4 + rtol 1e-4 of an FP64 sum) at T 58560. - flashinfer-ai/flashinfer#8827: new non-deterministic backward routes on SM100a and SM107a. With the FP32 weight the op passes, the backward runs one kernel where FlashInfer measured the pair it ran before with a larger weight-gradient error than the eager FP32 reduction of the same inputs: at H = 2048 on SM100a from T 337 to 4096 and on SM107a from T 1025 to 16384, at H = 4096 on SM100a from T 1793 to 3584, at H = 5120 on SM100a from T 65 to 6144 and on SM107a from T 4097 to 19968. On SM100a, H = 5120 runs one kernel above T 56256, where FlashInfer measured the programs it replaces above the budget, and H = 2048 one kernel from T 45057 to 58880; above T 65536, H = 2048 (SM100a) and H = 4096 (SM100a and SM107a) run compensated programs. On SM107a, H = 2048 from T 65 to 127 and H = 4096 from T 65 to 256 run another single kernel than before. Other widths get such routes too. - flashinfer-ai/flashinfer#8813: the deterministic backward of the smallest token counts on SM103a and SM107a (at H = 2048: T 3 to 171 on SM103a, up to 256 on SM107a) runs a new column-task kernel whose weight-gradient column sums use FP32 compensated arithmetic in a fixed row order. On SM103a it replaces the FP64 fold of the previous column-task kernel; on SM107a it replaces a main kernel and its reduce, so det's dw there no longer depends on the SM count. det's bits change at those token counts and stay reproducible run to run. On SM107a the deterministic backward of H = 256 also runs, from T 1 to 70656, the TwoSum main kernel and dw reduce it ran above T 70656 (the form SM100a runs), so det's H = 256 bits on SM107a change there too. On SM107a the FP32-weight forward of H = 5120 above 49152 rows also takes a new grid. - flashinfer-ai/flashinfer#8464: docstring formatting. - flashinfer-ai/flashinfer#8849: on SM103a the non-deterministic FP32-weight backward of H = 4096 above T 262144 and of H = 5120 above T 196608 runs a single kernel with TwoSum per-thread weight-gradient sums, in place of the kernel FlashInfer measured above the budget there (1.038 times it at T 393216 for H = 4096, 1.004 times at T 262144 for H = 5120). The BF16-weight routes, which the op does not use, switch at T 262144 and 393216. - flashinfer-ai/flashinfer#8848: on SM100a the non-deterministic backward runs a single kernel with blocked FP32 weight-gradient sums up to T 131072, from T 65537 at H = 2048, 4096 and 6144 (also with the residual), from T 56257 at H = 5120 and, at H = 1536, from T 75777 (FP32 weight) or 78785 (BF16 weight). Above T 131072 the TwoSum programs #8827 put in these windows still run. FlashInfer measured the new backward 1.029 to 1.266 times as fast as those programs on B200, with no measured dw element above the budget. At H = 5120 it runs the kernels SM107a runs above T 49152, whose sources gain a branch for SM100a only; their SM107a SASS does not change. FlashInfer's open pull request #8835 on the deterministic package is not included, nor #8899, which merged into main after c51a7540 and changes only SM103a's non-deterministic backward at H = 1536, 2048 and 6144 up to T 131072 (speed). The packages' entry points do not change, so the op code does not. cake_rmsnorm_det_train_backend goes from 582 to 634 files, cake_rmsnorm_train_backend from 581 to 653. Prebuilding all seven widths builds 146 libraries on SM100a (was 142), 160 on SM103a (was 154) and 138 on SM107a (was 133). The programs above the 232448 B shared-memory opt-in are unchanged: only det's backward reduce at 4096 on SM107a and its copy in the non-deterministic package (271360 B, CUDA 13.4's oversized mode), so the op's CUDA-runtime check covers the same programs; every binding above the opt-in sets that mode. The op's docstring names the column-task backward of det's smallest token counts on every architecture. Tests: - test_dw_accuracy_gpu.py (GPU): FlashInfer's gates at the token counts these changes moved. At H = 2048, on every architecture, T 128, 337, 1025, 4096, 16384, 58497, 58560, 61952 and 65536 with det and fast (fast also 98304); fast at one token count per window #8827, #8848 and #8849 changed, on that window's architectures, at H = 2048 (T 131073), H = 4096 (T 256, 3584, 65537, 131073, 262145) and H = 5120 (T 1792, 6144, 19968, 56257, 131073, 196609). An FP32 and a BF16 weight parameter, two inputs each: rstd within relative L2 1e-5 of FP64; y and dx within FlashInfer's tolerance; every dw element within atol 1e-4 + rtol 1e-4 of the FP64 sum at the forward's rstd and dw within relative L2 1e-4 (a BF16 gradient widened by its rounding); where fast's backward runs one kernel, the FP32 dw's relative L2 error at most the eager FP32 reduction's (or 2^-23 if larger); det bit for bit twice; over the FP32-parameter cases, the median ratio of the dw error to the eager FP32 reduction's at most 1, over det and over each backward form of fast. det stops at T 65536, where FlashInfer's deterministic suite checks the same element bound. These checks hold the installed build to the budget; test_fixed_windows.py is what fails on a build without these routes. - test_fixed_windows.py (no GPU): the windows above in the packages' registry: on SM103a fast runs det's programs from T 58497 to 61952 (H = 2048); fast runs one kernel in the #8827 windows; above a cut another single kernel than at the cut, the same one up to the next cut: on SM100a above T 65536 (H = 2048 and 4096) and T 131072 (H = 2048, 4096 and 5120), on SM107a above T 65536 (H = 4096), on SM103a above T 262144 (H = 4096) and T 196608 (H = 5120); on SM107a fast's T 65 to 127 at H = 2048 and T 65 to 256 at H = 4096 run another single kernel than the one they ran before; det's smallest token counts run a column-task kernel on every architecture and every width but 256. - test_rmsnorm_gpu.py: det's repeatability at the smallest token counts also covers the second column-task window of SM103a at H = 2048. python tools/vendor_cake.py fi --checkout ../flashinfer --commit c51a7540f042d7607f23bc29417147ca7a921828 \ --op rmsnorm --package flashinfer/cake_rmsnorm_det_train_backend \ --package flashinfer/cake_rmsnorm_train_backend Co-Authored-By: Claude <noreply@anthropic.com>
yyihuang
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
On main,
fast's FP32 weight gradient on B200 goes over FlashInfer's per-element dw budget at large token counts. The budget requires every element to be within atol 1e-4 + rtol 1e-4 of an FP64 sum; we measured single elements at 1.295 to 1.364 times that budget at H 4096 T 393216 and 524288 and at H 5120 T 524288. This PR re-vendors both Cake RMSNorm packages from flashinfer-ai/flashinfer@c51a7540, whose backward routes bring the worst element at those points down to 0.075–0.086 of the budget, and adds flashinfer-ai/flashinfer#8848's blocked-FP32 B200 backward for large T up to 131072. The cost isfastbackward time in some windows (main/PR by CUPTI kernel time as low as 0.72 on R200, 0.83 on B200 and 0.92 on GB300), whiledetat the smallest token counts gets faster on GB300 and R200.Changes
Both packages,
cake_rmsnorm_det_train_backend(det) andcake_rmsnorm_train_backend(fast), now come from flashinfer-ai/flashinfer@c51a7540f042d7607f23bc29417147ca7a921828(main, the merge of flashinfer-ai/flashinfer#8848); before this PR they were ateb3eb68d. The table lists the FlashInfer pull requests in that range that touch the packages, limited to the FP32-weight[T, H]routes the op uses:fastbackward at H 2048 runsdet's program from T 58497 to 61952. FlashInfer measured the pair that ran there before above the dw budget at T 58560.fastbackward runs one kernel in windows where FlashInfer measured the previous pair with a larger dw error than the eager FP32 reduction: H 2048 T 337–4096 (sm_100a) and T 1025–16384 (sm_107a); H 4096 T 1793–3584 (sm_100a); H 5120 T 65–6144 (sm_100a) and T 4097–19968 (sm_107a). On sm_100a, H 5120 runs one kernel above T 56256, where FlashInfer measured the replaced programs above the budget, and H 2048 runs one kernel from T 45057 to 58880. Above T 65536, H 2048 (sm_100a) and H 4096 (sm_100a, sm_107a) run compensated (TwoSum) programs; on sm_100a, flashinfer-ai/flashinfer#8848 replaces them up to T 131072. On sm_107a, H 2048 T 65–127 and H 4096 T 65–256 run a different single kernel. Other widths get such routes too.detbackward at the smallest token counts (at H 2048: T 3–171 on sm_103a, up to 256 on sm_107a) runs a new column-task kernel. It sums dw columns with FP32 compensated arithmetic in a fixed row order. On sm_103a it replaces the FP64 fold of the previous column-task kernel. On sm_107a it replaces a main kernel and its reduce, sodet's dw there no longer depends on the SM count.det's bits change at those token counts and stay reproducible run to run. On sm_107a,det's H 256 backward at T ≤ 70656 also moves to the TwoSum main kernel and dw reduce it ran above T 70656 (the form sm_100a runs), so its H 256 bits change there too. On sm_107a the H 5120 forward above 49152 rows also uses a new grid.fastbackward of H 4096 above T 262144 and of H 5120 above T 196608 runs one kernel with TwoSum per-thread dw sums. It replaces the kernel FlashInfer measured above the budget there (1.038 times the budget at H 4096 T 393216, 1.004 times at H 5120 T 262144).fastbackward runs one kernel with blocked FP32 dw sums up to T 131072: from T 65537 at H 2048, 4096 and 6144, from T 56257 at H 5120, and from T 75777 at H 1536. Above T 131072 the TwoSum programs from flashinfer-ai/flashinfer#8827 still run. FlashInfer measured the new backward at 1.029–1.266 times the speed of those programs on B200, with no measured dw element above the budget. At H 5120 it runs the kernels that sm_107a runs above T 49152. Their sources gain an SM100a-only branch, and their sm_107a SASS does not change.c51a7540and is not included. It changes only the GB300 (sm_103a)fastbackward speed at H 1536, 2048 and 6144 up to T 131072; it can follow in a separate update.det's column-task backward at the smallest token counts on every architecture.cake_rmsnorm_det_train_backendgoes from 582 to 634,cake_rmsnorm_train_backendfrom 581 to 653. Prebuilding all seven widths now builds 146 libraries on sm_100a (was 142), 160 on sm_103a (was 154) and 138 on sm_107a (was 133).det's backward reduce at H 4096 on sm_107a and its copy infast(271360 B, CUDA 13.4's oversized mode), so the op's CUDA-runtime check still covers the same programs.Vendoring command:
Accuracy
worst= max over inputs and elements of |dw − ref| / (1e-4 + 1e-4·|ref|). The reference is FP64 at the forward's rstd, and the budget is 1.0. All rows use the FP32 weight.B200 (sm_100a). 32 inputs per point, T up to 524288. Measured on the previous revision of this branch, which runs this PR's programs at these points.
fastpointfast's program changes. Inside the fix(cake_rmsnorm): sm_100a and sm_107a dw-accuracy backward routes flashinfer-ai/flashinfer#8827 window H 2048 T 45057–58880, the worst element rises at some points even though relL2 falls. The 0.707 at T 58880 is the highest PRfastworst in this measurement.fastrelL2 5.86e-08–6.15e-08 (main 7.61e-08–8.57e-08) and worst ≤ 0.312 (main ≤ 0.399). As in the #8827 window, the worst element rises at some of these points while relL2 falls (for example H 4096 T 131072: 0.269 → 0.312). The #8848 windows at H 1536 and H 6144 were not measured.detdoes not change on sm_100a; see Known limitation.GB300 (sm_103a). Measured on the previous revision of this branch; the programs are identical for this architecture. On our GB300 no main case went over the budget, including the points where FlashInfer measured main's kernel above it:
fastpointR200 (sm_107a). Max over 2 seeds. Measured on the previous revision of this branch; the programs are identical for this architecture. No case on main or on this PR went over the budget.
fastH 2048, T 1025fastH 4096, T 131072fastH 5120, T 19968detH 2048, T 128Performance
All tables show backward time. main/PR > 1 means this PR is faster. The "FI PR" column gives the flashinfer-ai/flashinfer pull request that changed the route.
B200 (sm_100a, 148 SMs),
fastbackwardCUPTI kernel time (sum of the op's kernels and memsets, median of per-process medians), FP32 weight, L2 flushed before every timed op, main and PR builds interleaved across processes. Rows marked † were measured on this PR's head. The other rows were measured on the previous revision of this branch, whose sm_100a programs are the same as this PR's outside the 8848 windows.
GB300 (sm_103a, 152 SMs)
CUPTI kernel time (sum of the op's kernels and memsets, median of per-process medians), FP32 weight, L2 flushed before every timed op, main and PR builds interleaved across processes. Measured on the previous revision of this branch; the programs are identical for this architecture.
detdetdetfastfastfastfastR200 (sm_107a, 212 SMs)
CUPTI kernel time (sum of the op's kernels and memsets, median of per-process medians), FP32 weight, L2 flushed before every timed op, main and PR builds interleaved across processes. Measured on the previous revision of this branch; the programs are identical for this architecture. One run of 6 processes per build; an earlier 4-process run of the same builds measured H 2048 T 1025 at 0.715. In the same run,
det's backward at H 2048 T 1025–16384, whose SASS is identical in both builds, measured 1.043–1.150, so R200 differences within about 15 % at small T are not attributable to the code.fastfastfastfastfastfastfastfastfastdetdetTrade-off
fastbackward in windows where FlashInfer moved to more accurate programs:fast: H 2048 T 127 (1.481) and H 4096 T 256 (1.705).det's column-task backward at the smallest token counts: GB300 (1.188–1.254) and R200 (1.241–1.358 at H 2048 T 127 and 256 and H 4096 T 256).detand the forward run the same programs. Their measured ratios are 0.988–1.062; the 1.062 isdetat H 5120 T 65, 5.7 vs 5.4 µs.fastat the 8810 rows and at H 4096 T 262145 measured 0.996–1.002.Tests
Test changes:
test_fixed_windows.py(no GPU): checks the routes above in the packages' registry. This covers the cuts, the single-kernel windows, anddet's column-task backward at the smallest token counts on every architecture and every width except 256. This is the test that fails on a build without these routes: 35 passed on this PR, while on main 23 failed and 12 passed.test_dw_accuracy_gpu.py(GPU, 63 cases): FlashInfer's gates at the token counts these changes moved. 31 points, each run as an FP32-weight and a BF16-weight case with two inputs (62 cases), plus one aggregate test.detandfast: T 128, 337, 1025, 4096, 16384, 58497, 58560, 61952 and 65536, plus T 98304 forfast.fastpoints at H 4096 and 5120 (and H 2048 above T 131072), one token count per changed window, on that window's architectures: H 2048 T 131073; H 4096 T 256, 3584, 65537, 131073, 262145; H 5120 T 1792, 6144, 19968, 56257, 131073, 196609.fastruns one kernel, FP32 dw relL2 at most that of the eager FP32 reduction (or 2^-23 if larger);detidentical bit for bit across two runs.detand for each backward form offast.detcases stop at T 65536.test_rmsnorm_gpu.py: thedetrepeatability test at the smallest token counts now also covers sm_103a's second column-task window at H 2048._windows.py: helpers that read from the packages' registry which backward programs a token count runs and in which form (one kernel, column task or pair). The two new modules andtest_rmsnorm_gpu.pyuse it.Runs:
--upstream;tests/cake/rmsnorm(notmulti_gpu): 206 passed, 8 skipped (accuracy cases for other architectures);tests/cake/_runtime: 316 passed;detdigests bitwise equal across 3 processes;detand forward programs are identical.not multi_gpu, empty cache): 190 passed, 14 skipped (accuracy cases for other architectures);_runtime: 316/316;test_dw_accuracy_gpu.py: 43 passed, 14 skipped (cases for other architectures; the previous revision's 57-case file, and the 6 cases this PR adds are sm_100a-only);detdigests bitwise equal across 4 processes on 3 GPUs;test_dw_accuracy_gpu.py: 45 passed, 12 skipped, the same result as on main (the previous revision's 57-case file; the 6 added cases are sm_100a-only).Known limitation
On B200, the
detbackward from FlashInfer's deterministic package goes over the element budget in single elements at very large token counts (32 FP32 inputs per point):main and this PR run the same
detprograms at these points, so this PR does not change the behaviour.test_dw_accuracy_gpu.pychecksdetonly up to T 65536, which is where FlashInfer's deterministic suite checks the same element bound. FlashInfer also measureddetabove the budget at large token counts on sm_103a (flashinfer-ai/flashinfer#8849) and on sm_100a (flashinfer-ai/flashinfer#8827); our GB300 and R200 runs did not reach those token counts. The upstream draft flashinfer-ai/flashinfer#8835 targets all three architectures and is not included here.🤖 Generated with Claude Code