Skip to content

perf(cake_dense_projection_gemm): run the deterministic program at the sm_107a token counts its new plans win - #8363

Merged
yyihuang merged 2 commits into
flashinfer-ai:mainfrom
yyihuang:cake-dense-projection-gemm-sm107a-det-token-counts
Nov 2, 2026
Merged

perf(cake_dense_projection_gemm): run the deterministic program at the sm_107a token counts its new plans win#8363
yyihuang merged 2 commits into
flashinfer-ai:mainfrom
yyihuang:cake-dense-projection-gemm-sm107a-det-token-counts

Conversation

@yyihuang

@yyihuang yyihuang commented Nov 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

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 not on all of them. At 16231, the deterministic program was as fast or faster on one of the eight GPUs. At 16384, it was as fast or faster on one GPU, and on a second GPU the split-K program was less than 1 % 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.

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

@averyhNV
perf(cake_dense_projection_gemm): run the deterministic program at th…
1e02d80 
…e sm_107a token counts its new plans win

The deterministic package's sm_107a rules for single token counts (flashinfer-ai#8303: 16231
and 16384 tokens) make several FP32 weight-gradient rows faster there than this
package's split program.  Re-measured on 208-SM and 212-SM parts against those
plans at 16172, 16231 and 16384 tokens (op cells and whole backward /
checkpointed step), this routes the losing token counts to the deterministic
program:

- K1_SPLIT_RULE_EXCLUDED_TOKENS also takes sm_107a rows: an entry holds from its
  SM count up at or above the row's SM-count floor (as before) and on exactly its
  SM count below it (the parts of K1_SPLIT_RULES_BY_SM); k1_plan applies the
  entries to both branches.
- 208 SMs: 4096 -> 48, 4096 -> 64, 4096 -> 512, 5120 -> 48 at 16231 and 16384;
  4096 -> 1024, 5120 -> 64 at 16231; 5120 -> 1024 at 16384.
- 212 SMs: 4096 -> 512, 5120 -> 512 at 16231 and 16384; 4096 -> 1024, 5120 -> 48,
  5120 -> 64 at 16231; 4096 -> 4096 runs the deterministic program (verdict None).
- 216 SMs (not measured on 216-SM parts yet): 4096 -> 48, 4096 -> 512 at 16231
  and 16384; 4096 -> 1024, 5120 -> 48, 5120 -> 64 at 16231.

Every other token count, SM count and row keeps its verdict.  Depends on flashinfer-ai#8303:
before it merges these token counts run the earlier deterministic plans.
@coderabbitai

coderabbitai Bot commented Nov 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Warning

Review limit reached

You'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 1 minute.

Check out review usage here.

View limit details

Limit details: You’ve used all 8 included reviews currently available.

Learn how review limits work.

Review configuration:

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: 9508fd71-8785-4d89-a4b0-6d12b2a7ea0a
Commits

Reviewing files that changed from the base of the PR and between 1c8c8e6 and 91c6b64.

Files selected for processing: 3
  • flashinfer/experimental/dense_projection_gemm/README.md
  • flashinfer/experimental/dense_projection_gemm/cake_backend.py
  • tests/experimental/test_cake_dense_projection_gemm.py
  • 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

@yyihuang

yyihuang commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py

@averyhNV
perf(cake_dense_projection_gemm): keep the split program for 5120 -> …
91c6b64 
…1024 FP32 at 16384 tokens on 208 SMs

A second measurement of 5120 -> 1024 FP32 on a 208-SM part (16172, 16231 and
16384 tokens, four more runs on another part), taken together with the first,
keeps this package's program within the rule at 16384 tokens: backward pass
0.8923 / 0.8886 and checkpointed step 0.9847 / 0.9655 of the deterministic
program (ratio of the medians over the runs, graph replays / eager launches).
The 208-SM entry for that row is removed.

4096 -> 1024 FP32 at 16231 tokens on 212 SMs, measured a second time as well,
keeps its entry: over both measurements its checkpointed step is 1.0213 of the
deterministic program in graph replays.
@yyihuang

yyihuang commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py

@yyihuang
yyihuang marked this pull request as ready for review November 2, 2026 19:00
@yyihuang

yyihuang commented Nov 2, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/experimental/test_cake_dense_projection_gemm.py tests/experimental/test_cake_dense_projection_det_gemm_names.py tests/experimental/test_cake_dense_projection_det_gemm.py

@yyihuang
yyihuang enabled auto-merge (squash) November 2, 2026 19:03
@yyihuang yyihuang mentioned this pull request Nov 2, 2026
Open
1 task done
yzh119
yzh119 approved these changes Nov 2, 2026
View reviewed changes

@yzh119 yzh119 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@yyihuang
yyihuang merged commit dd15ada into flashinfer-ai:main Nov 2, 2026
29 of 38 checks passed
@yyihuang yyihuang mentioned this pull request Nov 3, 2026
Merged
yyihuang added a commit that referenced this pull request Nov 3, 2026
@yyihuang @averyhNV
perf(cake_dense_projection_gemm): measured 216-SM sm_107a K1 routes: …
33e71a5 
…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>
@yyihuang yyihuang mentioned this pull request Nov 5, 2026
Merged
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yzh119 yzh119 yzh119 approved these changes
@flashinfer-bot flashinfer-bot Awaiting requested review from flashinfer-bot flashinfer-bot is a code owner
@bkryu bkryu Awaiting requested review from bkryu bkryu is a code owner
@aleozlx aleozlx Awaiting requested review from aleozlx aleozlx is a code owner
@YangXu1990uiuc YangXu1990uiuc Awaiting requested review from YangXu1990uiuc YangXu1990uiuc is a code owner
@Anerudhan Anerudhan Awaiting requested review from Anerudhan Anerudhan is a code owner
@Jingfan-Sun Jingfan-Sun Awaiting requested review from Jingfan-Sun Jingfan-Sun is a code owner

Assignees

No one assigned

Labels

run-ci

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@yyihuang @yzh119 @flashinfer-bot @averyhNV