Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yyihuang merged 2 commits into from
Nov 2, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
1e02d80
perf(cake_dense_projection_gemm): run the deterministic program at th…
averyhNV Nov 2, 2026
91c6b64
perf(cake_dense_projection_gemm): keep the split program for 5120 -> …
averyhNV Nov 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 45 additions & 13 deletions flashinfer/experimental/dense_projection_gemm/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,12 +68,13 @@ than the deterministic package's shipped plan of the row in every measurement ru
row that missed; a row without an entry was not measured. Both run the deterministic program.

The K1 verdicts were measured at 16172 to 16384 tokens. Those of `K1_SPLIT_RULES` hold at every token count their split
count fits (`num_splits * 64 <= T + 63`), except, on `sm_100a` / `sm_103a` parts of 152 SMs or more, the token counts of
`K1_SPLIT_RULE_EXCLUDED_TOKENS` (the K1 token table below); the tables of `K1_SPLIT_RULES_BY_SM` hold only at 16172 to
16384 tokens (`K1_SPLIT_RULES_BY_SM_TOKENS`). The router verdicts hold only inside token windows (`ROUTER_RULES`; the
last table below) and pass a stricter rule: measured in one long-lived process and in a fresh process per token count,
and no more than 3 % slower than the deterministic plan next to such a kernel both in the same launch order and with the
router partition launched before that kernel against the deterministic plan launched after it. Every other token count
count fits (`num_splits * 64 <= T + 63`), except the token counts of `K1_SPLIT_RULE_EXCLUDED_TOKENS` (the K1 token
tables below: `sm_100a` / `sm_103a` parts of 152 SMs or more, `sm_107a` parts of 208, 212 and 216 SMs); the tables of
`K1_SPLIT_RULES_BY_SM` hold only at 16172 to 16384 tokens (`K1_SPLIT_RULES_BY_SM_TOKENS`), except those token counts.
The router verdicts hold only inside token windows (`ROUTER_RULES`; the last table below) and pass a stricter rule:
measured in one long-lived process and in a fresh process per token count, and no more than 3 % slower than the
deterministic plan next to such a kernel both in the same launch order and with the router partition launched before
that kernel against the deterministic plan launched after it. Every other token count
of a router row runs the deterministic program. The K1 verdicts hold from the floor of `K1_MIN_SM_COUNT` /
`K1_RULE_MIN_SM_COUNT` up (a part with more SMs than the measured ones takes them without a measurement, a part with
fewer runs the deterministic program); the tables of `K1_SPLIT_RULES_BY_SM` hold for exactly 208 and 212 SMs. The
Expand All @@ -89,7 +90,7 @@ of `"waves"` (the deterministic plan's last wave below 80 % full). The measured
| `K1_MIN_SM_COUNT` | `arch` | the fewest SMs for the K1 verdicts of the architecture (`sm_100a` / `sm_103a`: 148, `sm_107a`: 216) |
| `K1_RULE_MIN_SM_COUNT` | a key of `K1_SPLIT_RULES` | the fewest SMs for that row's verdict where it is above `K1_MIN_SM_COUNT` (`sm_103a` `4096 -> 48`: 152) |
| `K1_SPLIT_RULES_BY_SM` | `(arch, sm_count)` | the K1 verdicts of the parts with exactly `sm_count` SMs below a row's floor (keys and values as in `K1_SPLIT_RULES`), at the token counts of `K1_SPLIT_RULES_BY_SM_TOKENS` only; a part without a table, and every other token count, runs every K1 row below the floor on the deterministic program (`sm_107a`: 208 and 212 SMs) |
| `K1_SPLIT_RULE_EXCLUDED_TOKENS` | a key of `K1_SPLIT_RULES` on `sm_100a` / `sm_103a` | `((min_sm_count, min_tokens, max_tokens), ...)`: on a part of at least `min_sm_count` SMs the row runs the deterministic program at a token count from `min_tokens` to `max_tokens` (`None`: no upper bound); every other token count keeps the verdict |
| `K1_SPLIT_RULE_EXCLUDED_TOKENS` | a key of `K1_SPLIT_RULES` or of a table of `K1_SPLIT_RULES_BY_SM` | `((sm_count, min_tokens, max_tokens), ...)`: the row runs the deterministic program at a token count from `min_tokens` to `max_tokens` (`None`: no upper bound) on a part of at least `sm_count` SMs (`sm_count` at or above the row's SM-count floor) or of exactly `sm_count` SMs (below it); every other token count and part keeps the verdict |
| `K1_SPLIT_RULES_BY_SM_TOKENS` | -- | `(min_tokens, max_tokens)`: the token counts (`K` of the weight gradient) the tables of `K1_SPLIT_RULES_BY_SM` hold for, `(16172, 16384)` (where they were measured) |
| `ROUTER_RULES` | `(arch, layout, hidden, experts)`: `layout` `"kk"` (forward, `hidden = K`) or `"nn_t"` (weight gradient, `hidden` = the launch's `M`), `experts` = the launch's `N` | `None` (the deterministic program), token windows `((min_tokens, max_tokens, spec), ...)` (the spec of the window that holds the token count -- `K` of the weight gradient, `M` of the forward; every other token count: the deterministic program) or a spec alone (every token count: a plan a test or a measurement forces). A spec is `"waves"`, `"H<f>x<k>"`, `"H<f>x<k>p<h>"`, `"forward"`, `"pairs"`, `"tiles"`, `"S<n>"` or a unit count |
| `ROUTER_SM_COUNT` | `arch` | the SM count of the router verdicts of the architecture (`sm_100a` / `sm_103a`: 152, `sm_107a`: 216); a device with any other SM count runs the deterministic program, except in a `"waves"` window (the condition of `router_waves`) or with a table of its own in `ROUTER_RULES_BY_SM` |
Expand All @@ -101,19 +102,19 @@ of `"waves"` (the deterministic plan's last wave below 80 % full). The measured
|---|---|---|---|
| `sm_100a` (>= 148 SMs) | `4096 -> 32, 48, 64`; `5120 -> 32, 48, 64`; `6144 -> 32`: (64, 2). `4096 -> 128`: deterministic (measured, missed) | deterministic (4 splits) | 256 experts, in the token windows of the last table: on 152 SMs hidden 4096 `"H1x5p44"` and hidden 5120 `"H0.5x3"`; hidden 5120 `"waves"` on every SM count where it applies. Every other token count and SM count, and hidden 6144: deterministic (6 splits) |
| `sm_103a` (>= 148 SMs) | `4096 -> 32`; `5120 -> 32, 48, 64`; `6144 -> 32`: (64, 2); `4096 -> 48` (>= 152 SMs): (64, 2); `5120 -> 128`, `6144 -> 128`: (128, 2). `4096 -> 64`, `4096 -> 128`: deterministic (measured, missed) | deterministic (4 splits) | 256 experts, in the token windows of the last table: on 152 SMs hidden 5120 `"H0.5x3"`; hidden 5120 `"waves"` on every SM count where it applies. Every other token count and SM count, hidden 4096 (`"H1x4p44"` up to 7.6 % slower when launched before a resident kernel) and hidden 6144: deterministic (6 splits) |
| `sm_107a` (>= 216 SMs) | `4096 -> 32` BF16 (64, 4), FP32 (64, 3); `4096 -> 48` FP32 (64, 3); `4096 -> 64` FP32 (64, 3); `5120 -> 32` (64, 3); `5120 -> 48`, `5120 -> 64` FP32 (64, 3); `6144 -> 32` (64, 3); `8192 -> 128` (128, 2); `4096 -> 512` FP32 (128, 1) `p = 1`; `5120 -> 512` FP32 (128, 1) `p = 1`; `4096 -> 1024` FP32 (256, 1) `p = 1`; `5120 -> 1024` (256, 1) `p = 1`. `4096 -> 64` BF16, `4096 -> 128`, `4096 -> 512` BF16, `4096 -> 4096` FP32: deterministic (measured, missed) | deterministic (4 splits) | 256 experts, in the token window of the last table: on 216 SMs hidden 5120 `"H0.5x4"` (the head/tail program). Every other token count and SM count, and hidden 4096 / 6144: deterministic (6 splits): at 4096 `"H0.5x4"` loses next to a kernel holding 16 SMs that starts during it, at 6144 it leaves the backward pass no faster |
| `sm_107a` (208 SMs, 16172 to 16384 tokens) | `4096 -> 32, 48, 64`, `5120 -> 32, 48, 64` FP32: (64, 3); `4096 -> 512`, `5120 -> 512` FP32 (128, 1) `p = 1`; `4096 -> 1024`, `5120 -> 1024` FP32: (256, 1) `p = 1`. `4096 -> 32, 64` BF16, `5120 -> 32` BF16, `6144 -> 32`: deterministic (no verdict: its stream-K program did not complete next to a resident kernel); `4096 -> 512` BF16, `5120 -> 1024` BF16, `4096 -> 4096` FP32, `8192 -> 128`: deterministic (measured, missed) | deterministic (4 splits) | 256 experts, hidden 5120 at 16172 to 16384 tokens: `"H0.5x4"` (the head/tail program). Every other token count, and hidden 4096 / 6144: deterministic (6 splits): next to a resident kernel the head/tail program missed at both hidden sizes on the parts at full SM clock |
| `sm_107a` (212 SMs, 16172 to 16384 tokens) | `4096 -> 32, 48, 64` FP32 (64, 3); `5120 -> 32, 48, 64` FP32 (64, 3); `8192 -> 128` (128, 2); `4096 -> 512`, `5120 -> 512` FP32 (128, 1) `p = 1`; `4096 -> 1024` FP32: (192, 1) `p = 1` (row-major-store program, 176 CTAs); `5120 -> 1024`, `4096 -> 4096` FP32: (256, 1) `p = 1`. `4096 -> 32, 64` BF16, `5120 -> 32` BF16, `6144 -> 32`: deterministic (no verdict, as on 208 SMs); `4096 -> 512` BF16: deterministic (measured, missed); `5120 -> 1024` BF16: deterministic (1.3 to 2.5 % faster here, missed on 208 SMs) | deterministic (4 splits) | 256 experts, hidden 5120 at 16172 to 16384 tokens: `"H0.5x4"` (the head/tail program). Every other token count, and hidden 4096 / 6144: deterministic (6 splits): next to a resident kernel `"H0.5x4"` missed at hidden 4096 on all four parts and at hidden 6144 on two of them |
| `sm_107a` (>= 216 SMs) | `4096 -> 32` BF16 (64, 4), FP32 (64, 3); `4096 -> 48` FP32 (64, 3); `4096 -> 64` FP32 (64, 3); `5120 -> 32` (64, 3); `5120 -> 48`, `5120 -> 64` FP32 (64, 3); `6144 -> 32` (64, 3); `8192 -> 128` (128, 2); `4096 -> 512` FP32 (128, 1) `p = 1`; `5120 -> 512` FP32 (128, 1) `p = 1`; `4096 -> 1024` FP32 (256, 1) `p = 1`; `5120 -> 1024` (256, 1) `p = 1`, except the token counts of the `sm_107a` K1 token table. `4096 -> 64` BF16, `4096 -> 128`, `4096 -> 512` BF16, `4096 -> 4096` FP32: deterministic (measured, missed) | deterministic (4 splits) | 256 experts, in the token window of the last table: on 216 SMs hidden 5120 `"H0.5x4"` (the head/tail program). Every other token count and SM count, and hidden 4096 / 6144: deterministic (6 splits): at 4096 `"H0.5x4"` loses next to a kernel holding 16 SMs that starts during it, at 6144 it leaves the backward pass no faster |
| `sm_107a` (208 SMs, 16172 to 16384 tokens) | `4096 -> 32, 48, 64`, `5120 -> 32, 48, 64` FP32: (64, 3); `4096 -> 512`, `5120 -> 512` FP32 (128, 1) `p = 1`; `4096 -> 1024`, `5120 -> 1024` FP32: (256, 1) `p = 1`; except the token counts of the `sm_107a` K1 token table. `4096 -> 32, 64` BF16, `5120 -> 32` BF16, `6144 -> 32`: deterministic (no verdict: its stream-K program did not complete next to a resident kernel); `4096 -> 512` BF16, `5120 -> 1024` BF16, `4096 -> 4096` FP32, `8192 -> 128`: deterministic (measured, missed) | deterministic (4 splits) | 256 experts, hidden 5120 at 16172 to 16384 tokens: `"H0.5x4"` (the head/tail program). Every other token count, and hidden 4096 / 6144: deterministic (6 splits): next to a resident kernel the head/tail program missed at both hidden sizes on the parts at full SM clock |
| `sm_107a` (212 SMs, 16172 to 16384 tokens) | `4096 -> 32, 48, 64` FP32 (64, 3); `5120 -> 32, 48, 64` FP32 (64, 3); `8192 -> 128` (128, 2); `4096 -> 512`, `5120 -> 512` FP32 (128, 1) `p = 1`; `4096 -> 1024` FP32: (192, 1) `p = 1` (row-major-store program, 176 CTAs); `5120 -> 1024` FP32: (256, 1) `p = 1`; except the token counts of the `sm_107a` K1 token table. `4096 -> 32, 64` BF16, `5120 -> 32` BF16, `6144 -> 32`: deterministic (no verdict, as on 208 SMs); `4096 -> 512` BF16, `4096 -> 4096` FP32: deterministic (measured, missed; `4096 -> 4096` FP32 against the deterministic package's 160-wide tiles of the row); `5120 -> 1024` BF16: deterministic (1.3 to 2.5 % faster here, missed on 208 SMs) | deterministic (4 splits) | 256 experts, hidden 5120 at 16172 to 16384 tokens: `"H0.5x4"` (the head/tail program). Every other token count, and hidden 4096 / 6144: deterministic (6 splits): next to a resident kernel `"H0.5x4"` missed at hidden 4096 on all four parts and at hidden 6144 on two of them |

*Table: the rows that run this package's programs. `4096 -> 128` misses on `sm_100a` / `sm_103a` next to a kernel
holding 32 SMs (its 128 CTAs no longer fit the free SMs) and on `sm_107a` at 16384 tokens; on `sm_107a`,
`4096 -> 512` BF16 (192 CTAs) and `4096 -> 4096` FP32 (216 CTAs) miss next to a kernel holding 32 SMs the same way.
`4096 -> 64` (`sm_103a`; `sm_107a` BF16) and `4096 -> 48` on 148-SM `sm_103a` parts are not faster than the
deterministic plan in every run. The `sm_107a` one-split rows store whole tiles: they win where the deterministic
package's `sm_107a` program of the row is slower. Every K1 choice also needs `num_splits * 64 <= T + 63` (shorter token
counts run the deterministic program); on `sm_100a` / `sm_103a` parts of 152 SMs or more the token counts of the K1
token table below run the deterministic program; the 208-SM and 212-SM choices hold only at 16172 to 16384 tokens
(`K1_SPLIT_RULES_BY_SM_TOKENS`).*
counts run the deterministic program); the token counts of the K1 token tables below run the deterministic program
(`sm_100a` / `sm_103a`: parts of 152 SMs or more; `sm_107a`: 208, 212 and 216 SMs); the 208-SM and 212-SM choices hold
only at 16172 to 16384 tokens (`K1_SPLIT_RULES_BY_SM_TOKENS`).*

| K1 row `d_in -> d_out` | `sm_100a` BF16 | `sm_100a` FP32 | `sm_103a` BF16 | `sm_103a` FP32 |
|---|---|---|---|---|
Expand Down Expand Up @@ -154,6 +155,37 @@ slower there: BF16 output at 1024 tokens on `sm_103a` `4096 -> 48`, `6144 -> 32`
`6144 -> 32`; FP32 output at 1024 tokens on `sm_100a` `6144 -> 32`; FP32 output at 8192 tokens on `sm_103a`
`4096 -> 48`, `5120 -> 128`. At 65536 tokens neither program meets the BF16 error budget.*

| K1 row `d_in -> d_out` (FP32 out) | 208 SMs | 212 SMs | 216 SMs (not measured) |
|---|---|---|---|
| `4096 -> 32` | -- | -- | -- |
| `4096 -> 48` | 16231; 16384 | -- | 16231; 16384 |
| `4096 -> 64` | 16231; 16384 | -- | -- |
| `4096 -> 512` | 16231; 16384 | 16231; 16384 | 16231; 16384 |
| `4096 -> 1024` | 16231 | 16231 | 16231 |
| `5120 -> 32` | -- | -- | -- |
| `5120 -> 48` | 16231; 16384 | 16231 | 16231 |
| `5120 -> 64` | 16231 | 16231 | 16231 |
| `5120 -> 512` | -- | 16231; 16384 | -- |
| `5120 -> 1024` | -- | -- | -- |

*Table: the K1 token counts that run the deterministic program on `sm_107a` parts (`K1_SPLIT_RULE_EXCLUDED_TOKENS`;
`--`: none). They are token counts of the deterministic package's `sm_107a` rules for single token counts (16231 and
16384 tokens; `4096 -> 1024` at 16231 only), which give the row a faster deterministic plan there; every other token
count keeps the row's verdict. On 212 SMs `4096 -> 4096` FP32 runs the deterministic program at every token count
(`K1_SPLIT_RULES_BY_SM`): the deterministic package's rule for that row holds at every token count. On 208 and 212 SMs
the rows were measured at 16172, 16231 and 16384 tokens against those plans: 15 paired groups of 20 steps per row and
token count with the L2 cache flushed before every launch, in a long-lived process per GPU and in a fresh process per
row, in both arm orders, eager launches and CUDA graph replays, next to a kernel holding 8, 16 or 32 SMs in both launch
orders, with the backward pass and the checkpointed step timed (9 paired groups); 212 SMs: eight parts in two nodes
(`4096 -> 1024`: sixteen parts in five nodes), 208 SMs: two parts (`5120 -> 1024`: three). A measured token count
runs the deterministic program where the row's program failed the K1 rule above in some run, or left the backward
pass or the checkpointed step more than 1.3 % slower than the deterministic program (ratio of the medians over the
runs, or median of the paired ratios, in eager launches or in graph replays). `4096 -> 1024` on 212 SMs and
`5120 -> 1024` on 208 SMs were measured twice; the rule takes the runs of both measurements together. The 216-SM
entries were not measured on 216-SM parts: they are the rule token counts where the deterministic package's speed-up
of the row on a 216-SM part (its own measurement against its earlier plan) leaves this package's program less than
2 % faster than the deterministic program; the other rule token counts keep the 216-SM verdicts.*

`"waves"` (`router_waves`) launches whole waves of the device's CTA pairs (`pairs = sm_count // 2`): `c * pairs`
units with `c = max(2, round(tiles * k_iters / (68 * pairs))) + 1`, and only where the deterministic plan's six
splits fill their last wave below 80 % (`w / ceil(w)`, `w = 6 * tiles / pairs`); otherwise the row runs the
Expand Down
Loading
Loading