Repository navigation
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
Conversation
…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.
Contributor
|
Warning Review limit reachedYou'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. View limit details
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. |
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_dense_projection_det_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py |
yyihuang
marked this pull request as ready for review
November 2, 2026 18:59
yyihuang
enabled auto-merge (squash)
November 2, 2026 18:59
yyihuang
requested review from
Anerudhan,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
flashinfer-bot and
yzh119
as code owners
November 2, 2026 19:00
1 task done
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>
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.
The Cake
dense_projection_det_gemmpackage inflashinfer.experimentalis 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 addssm_107arow 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 onsm_103a, where one more instruction also moves. Outputs are bitwise identical to the previous version.The PR also adds 34
sm_107arow 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 thesm_107acount 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_100aandsm_103a, and the FP32 router GEMM through BF16x3 emulation. It has programs forsm_100a(GB200),sm_103a(GB300) andsm_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.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.