diff --git a/flashinfer/experimental/dense_projection_gemm/README.md b/flashinfer/experimental/dense_projection_gemm/README.md index ebb7d1fecc9..7d9eb3c869c 100644 --- a/flashinfer/experimental/dense_projection_gemm/README.md +++ b/flashinfer/experimental/dense_projection_gemm/README.md @@ -43,7 +43,7 @@ deterministic already.* | property | contract | |---|---| | operands and accumulation | BF16 operands with FP32 accumulation (K1); FP32 operands split exactly into three BF16 terms with FP32 accumulation (K2), except in a K2 SIMT window (input and weight gradient): the FP32 operands as they are on the CUDA cores with compensated sums (per output a sum and its compensation: Kahan adds of short FMA chains, or in the Dot2 form each product's rounding error kept; the slices merged by TwoSum with their compensations), no BF16x3 split. No silent TF32 or BF16 for the router | -| accuracy | the deterministic package's budgets against FP64 of the exact operands, worst over repeated runs: BF16 out `atol = rtol = 1e-2`; FP32 out relative Frobenius error `<= 5e-5` and max error `<= 1e-3 * max abs(ref)`; router `<= 2e-6` relative Frobenius error and no row whose top-8 experts differ from FP64 (outside near ties of the 8th and 9th logits) | +| accuracy | the deterministic package's budgets against FP64 of the exact operands, worst over repeated runs: BF16 out `atol = rtol = 1e-2`; FP32 out relative Frobenius error `<= 5e-5` and max error `<= 1e-3 * max abs(ref)`; router `<= 2e-6` relative Frobenius error and no row whose top-8 experts differ from FP64 (outside near ties of the 8th and 9th logits). At the largest token counts some thin K1 rows miss the BF16 (FP32) budget with the deterministic program and with this package's program alike; there the row keeps its route. The deterministic program: BF16 output of `sm_100a` `4096 -> 32`, `4096 -> 48`, `4096 -> 64`, `5120 -> 32`, `5120 -> 48` and of `sm_103a` `4096 -> 32`, `4096 -> 48`, `5120 -> 32`, `5120 -> 48`, `5120 -> 64`, `5120 -> 128` (from 45058, 49153, 65536 or 65537 tokens up by row: the K1 token tables below), and FP32 output of the rows of the odd-block K1 table at an odd number of K blocks from 98305 tokens up. This package's program: BF16 output of `sm_100a` `4096 -> 64` from about 43009 to 65535 tokens and `5120 -> 32` at 47105 to 65536 tokens, of `sm_103a` `4096 -> 32` from about 49153 to 65535 tokens, and on 148-SM parts of `sm_103a` `5120 -> 48` from 49152 tokens, `5120 -> 64` at 47104 and from 49153 tokens and `5120 -> 128` from 65537 tokens | | rounding | each output element is rounded once, from its complete FP32 sum | | subnormals | kept: K1 adds the partials with IEEE round-to-nearest adds, K2 reduces with the TMA reduce-add; a K2 SIMT window (input and weight gradient) accumulates with round-to-nearest FMAs and adds, merges the sums of a CTA's slices and their compensations in a fixed order through shared memory and stores every element of `out` once, with no reduce-add; nothing flushes to zero | | run-to-run bits, K1 | two splits (every `sm_100a` and `sm_103a` row, `sm_107a` `8192 -> 128`): repeat bit for bit (one add of two values); one split (`sm_107a` rows marked `p = 1` below): repeat bit for bit (whole tiles); three or four splits (the other `sm_107a` rows): vary in the last bits | @@ -77,11 +77,14 @@ table below); the tables of (`K1_SPLIT_RULES_BY_SM_TOKENS`) and, for the rows of `K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, at the token ranges of that table (the 212-SM K1 token-range table below), except those token counts; the rows of `K1_SPLIT_RANGE_OVERRIDES_BY_SM` run that table's program at its token ranges (the same table), and those of -`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` at its token ranges where the token count spans an even number of 64-token K blocks -(the same table on 212 SMs, `4096 -> 128` FP32 included; on 152-SM `sm_100a` / `sm_103a` parts the 152-SM two-split -table below). Before all of these, the rows of `K1_CHAIN_RANGES_BY_SM` run the one-split program of several -accumulation chains at that table's token ranges on parts of exactly its SM count, at both block parities (the K1 -accumulation-chain table below). +`K1_TWO_SPLIT_RANGES_BY_SM` at its token ranges, at either number of 64-token K blocks (the same table on 212 SMs, +`4096 -> 128` FP32 included; on 152-SM `sm_100a` / `sm_103a` parts the 152-SM two-split table below). Before all of +these, the rows of `K1_CHAIN_RANGES_BY_SM` run the one-split program of several accumulation chains at that table's +token ranges on parts of exactly its SM count, at both block parities (the K1 accumulation-chain table below). A K1 +program runs only at token counts where its output met the deterministic package's budgets against FP64 (the accuracy +row above) on every measured seed; the ratio of its FP64 relative Frobenius error to the deterministic program's on +the same operands (q) is reported in the tables below as data, not as a rule. Where both programs miss the budget, +the row keeps its route. 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 @@ -115,9 +118,9 @@ it. The measured parts have 148 and 152 SMs (`sm_100a` / | `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` and `K1_SPLIT_OVERRIDES_BY_SM` hold for, `(16172, 16384)` (where they were measured) | | `K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` | `(arch, sm_count)`, then a key of that part's table in `K1_SPLIT_RULES_BY_SM` | `((min_tokens, max_tokens), ...)`: further token counts at which the row takes its verdict from the part's table, except the token counts of `K1_SPLIT_RULE_EXCLUDED_TOKENS`; every other row, part and token count keeps `K1_SPLIT_RULES_BY_SM_TOKENS` alone (`sm_107a`: 212 SMs, the 212-SM K1 token-range table below) | | `K1_SPLIT_RANGE_OVERRIDES_BY_SM` | `(arch, sm_count)`, then a key of that part's table in `K1_SPLIT_RULES_BY_SM` | `((min_tokens, max_tokens, (block_n, num_splits)), ...)`: on a part of exactly `sm_count` SMs below the row's floor, a token count from `min_tokens` to `max_tokens` runs that program, except the token counts of `K1_SPLIT_RULE_EXCLUDED_TOKENS`; a range lies outside `K1_SPLIT_RULES_BY_SM_TOKENS` and the row's `K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` (`sm_107a`: 212 SMs, the 212-SM K1 token-range table below) | -| `K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` | `(arch, sm_count)`, then a key as in `K1_SPLIT_RULES` (below the row's floor: a key of that part's table in `K1_SPLIT_RULES_BY_SM`) | `((min_tokens, max_tokens, (block_n, num_splits)), ...)`: on a part of exactly `sm_count` SMs, below the row's floor or at or above it, a token count from `min_tokens` to `max_tokens` that spans an even number of 64-token K blocks (`ceil(T / 64)` even) runs that program, except the token counts of `K1_SPLIT_RULE_EXCLUDED_TOKENS`; an odd number keeps the routing of the tables above; a range lies outside `K1_SPLIT_RULES_BY_SM_TOKENS` and the row's ranges in `K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` and `K1_SPLIT_RANGE_OVERRIDES_BY_SM` (`sm_107a`: 212 SMs, the two-split rows of the 212-SM K1 token-range table below, `4096 -> 128` FP32 without a verdict in the part's table of `K1_SPLIT_RULES_BY_SM`; `sm_100a` / `sm_103a`: 152 SMs, the 152-SM two-split table below) | +| `K1_TWO_SPLIT_RANGES_BY_SM` | `(arch, sm_count)`, then a key as in `K1_SPLIT_RULES` (below the row's floor: a key of that part's table in `K1_SPLIT_RULES_BY_SM`) | `((min_tokens, max_tokens, (block_n, num_splits)), ...)`: on a part of exactly `sm_count` SMs, below the row's floor or at or above it, a token count from `min_tokens` to `max_tokens` runs that program at either number of 64-token K blocks, except the token counts of `K1_SPLIT_RULE_EXCLUDED_TOKENS`; a range lies outside `K1_SPLIT_RULES_BY_SM_TOKENS` and the row's ranges in `K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` and `K1_SPLIT_RANGE_OVERRIDES_BY_SM` (`sm_107a`: 212 SMs, the two-split rows of the 212-SM K1 token-range table below, `4096 -> 128` FP32 without a verdict in the part's table of `K1_SPLIT_RULES_BY_SM`; `sm_100a` / `sm_103a`: 152 SMs, the 152-SM two-split table below) | | `K1_SPLIT_OVERRIDES_BY_SM` | `(arch, sm_count)` | verdicts (keys and values as in `K1_SPLIT_RULES`) that replace a row's `K1_SPLIT_RULES` verdict on parts of exactly `sm_count` SMs at or above the row's floor, at the token counts of `K1_SPLIT_RULES_BY_SM_TOKENS` only, except the token counts of `K1_SPLIT_RULE_EXCLUDED_TOKENS`; every other token count and SM count keeps the `K1_SPLIT_RULES` verdict (no entry) | -| `K1_CHAIN_RANGES_BY_SM` | `(arch, sm_count)`, then a key of `K1_SPLIT_RULES` | `((min_tokens, max_tokens, (block_n, chains)), ...)`: on a part of exactly `sm_count` SMs, a token count from `min_tokens` to `max_tokens` runs this package's one-split program of that tile width with that many accumulation chains (`k1_chain_plan`; the chain counts this package ships are `K1_CHAINS_CHOICES`), before every table above, `K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` included (`sm_107a`: 212 and 216 SMs, the K1 accumulation-chain table below) | +| `K1_CHAIN_RANGES_BY_SM` | `(arch, sm_count)`, then a key of `K1_SPLIT_RULES` | `((min_tokens, max_tokens, (block_n, chains)), ...)`: on a part of exactly `sm_count` SMs, a token count from `min_tokens` to `max_tokens` runs this package's one-split program of that tile width with that many accumulation chains (`k1_chain_plan`; the chain counts this package ships are `K1_CHAINS_CHOICES`), before every table above, `K1_TWO_SPLIT_RANGES_BY_SM` included (`sm_107a`: 212 and 216 SMs, the K1 accumulation-chain table below) | | `ROUTER_RULES` | `(arch, layout, hidden, experts)`: `layout` `"kk"` (forward, `hidden = K`), `"nn_t"` (weight gradient, `hidden` = the launch's `M`) or `"kn"` (input gradient, `hidden` = the launch's `N`), `experts` = the launch's `N` (`K` of the input gradient) | `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 and of the input gradient; 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"`, `"Hx"`, `"Hxp"`, `"forward"`, `"pairs"`, `"tiles"`, `"S"`, a unit count (this package's stream-K program) or `"simt_nms[
]"` (the exact-FP32 SIMT program of the row's layout, `router_simt`: plain without a form, compensated with `g` or, on a weight-gradient row, `dot2`; the only spec of an input-gradient row) | | `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`; below a SIMT window too) or with a table of its own in `ROUTER_RULES_BY_SM` | | `ROUTER_RULES_BY_SM` | `(arch, sm_count)` | the router verdicts of the parts with exactly `sm_count` SMs (keys and values as in `ROUTER_RULES`; `router_rules` picks the table of a device); such a part takes its own table only: a row without an entry, and every token count outside its windows, runs the deterministic program (`sm_100a` / `sm_103a`: 148 SMs; `sm_107a`: 208 and 212 SMs) | @@ -126,35 +129,35 @@ it. The measured parts have 148 and 152 SMs (`sm_100a` / | architecture | K1 weight gradient `d_in -> d_out` (`block_n`, `num_splits`), BF16 and FP32 out unless marked | K2 forward | K2 weight gradient | |---|---|---|---| -| `sm_100a` (>= 148 SMs) | `4096 -> 32, 48, 64`; `5120 -> 32, 48, 64`; `6144 -> 32`: (64, 2). `4096 -> 128`: deterministic (measured, missed); FP32 on 152 SMs: (64, 2) at an even K-block count from 1024 to 4096 tokens (the 152-SM two-split table) | deterministic (4 splits) | 256 experts on 148 SMs: the compensated SIMT programs from 1 to 160 tokens at hidden 4096, 5120 and 6144 (the compensated SIMT table below). Hidden 5120 from 1 to 128 tokens on every other SM count where it applies: `"waves"`. Every other token count, SM count and hidden size: deterministic (6 splits); the SIMT and head/tail windows of 152-SM parts missed the earlier FP64 error rule (the router FP64 error table below) | -| `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); `4096 -> 128` FP32 on 152 SMs: (64, 2) at an even K-block count from 1024 to 6144 tokens (the 152-SM two-split table) | deterministic (4 splits) | 256 experts on 148 SMs: the compensated SIMT programs from 1 to 160 tokens at hidden 4096, 5120 and 6144 (the compensated SIMT table below). Hidden 5120 from 1 to 160 tokens on every other SM count where it applies: `"waves"`. Every other token count, SM count and hidden size: deterministic (6 splits); the head/tail windows of 152-SM parts, and `"waves"` from 161 tokens, missed the earlier FP64 error rule (the router FP64 error table below) | +| `sm_100a` (>= 148 SMs) | `4096 -> 32, 48, 64`; `5120 -> 32, 48, 64`; `6144 -> 32`: (64, 2). `4096 -> 128`: deterministic (measured, missed); FP32 on 152 SMs: (64, 2) from 1024 to 4096 tokens (the 152-SM two-split table) | deterministic (4 splits) | 256 experts on 148 SMs: the compensated SIMT programs from 1 to 160 tokens at hidden 4096, 5120 and 6144 (the compensated SIMT table below). Hidden 5120 from 1 to 128 tokens on every other SM count where it applies: `"waves"`. Every other token count, SM count and hidden size: deterministic (6 splits); the SIMT and head/tail windows of 152-SM parts missed the earlier FP64 error rule (the router FP64 error table below) | +| `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); `4096 -> 128` FP32 on 152 SMs: (64, 2) from 1024 to 6144 tokens (the 152-SM two-split table) | deterministic (4 splits) | 256 experts on 148 SMs: the compensated SIMT programs from 1 to 160 tokens at hidden 4096, 5120 and 6144 (the compensated SIMT table below). Hidden 5120 from 1 to 160 tokens on every other SM count where it applies: `"waves"`. Every other token count, SM count and hidden size: deterministic (6 splits); the head/tail windows of 152-SM parts, and `"waves"` from 161 tokens, missed the earlier FP64 error rule (the router FP64 error table below) | | `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 (the four FP32 rows with one split: from 961 tokens up). `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), except the token counts of the `sm_107a` K1 token table. `4096 -> 512`, `5120 -> 512`, `4096 -> 1024`, `5120 -> 1024` FP32: deterministic (their one-split verdicts, (128, 1) and (256, 1) `p = 1`, fail the FP64 comparison of the `sm_107a` K1 token table at these token counts); `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 and the ranges of the 212-SM K1 token-range table) | `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. At the two-split ranges of the 212-SM K1 token-range table, where the token count spans an even number of 64-token K blocks (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`): `4096 -> 512`, `5120 -> 512` FP32 (128, 2); `4096 -> 1024` FP32 (192, 2) (row-major-store program); `5120 -> 1024` FP32 (256, 2); `4096 -> 128` FP32 (64, 2) from 1024 to 7680 tokens (no verdict in the 212-SM table: deterministic at every other token count). `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: (256, 1) `p = 1` at its range of the 212-SM K1 token-range table (`K1_SPLIT_RANGE_OVERRIDES_BY_SM`), deterministic at 16172 to 16384 tokens (1.3 to 2.5 % faster there, missed on 208 SMs) | deterministic (4 splits) | 256 experts, hidden 5120 and 6144 at 16172 to 16384 tokens: `"H0.5x4"` (the head/tail program with the BF16 parts of `A` in tensor memory, `ROUTER_HEAD_TAIL_TEMPLATES_BY_SM`). Every other token count, and hidden 4096: deterministic (6 splits): the SIMT windows (hidden 4096 and 5120 from 1 to 256 tokens, hidden 6144 from 1 to 192) missed the FP64 error rule (the router FP64 error table below); next to a resident kernel `"H0.5x4"` missed at hidden 4096 with both head/tail programs, and at hidden 6144 with the three-chunk-buffer program | +| `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 and the ranges of the 212-SM K1 token-range table) | `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. At the two-split ranges of the 212-SM K1 token-range table, at either number of 64-token K blocks (`K1_TWO_SPLIT_RANGES_BY_SM`): `4096 -> 512`, `5120 -> 512` FP32 (128, 2); `4096 -> 1024` FP32 (192, 2) (row-major-store program); `5120 -> 1024` FP32 (256, 2); `4096 -> 128` FP32 (64, 2) from 1024 to 7680 tokens (no verdict in the 212-SM table: deterministic at every other token count). `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: (256, 1) `p = 1` at its range of the 212-SM K1 token-range table (`K1_SPLIT_RANGE_OVERRIDES_BY_SM`), deterministic at 16172 to 16384 tokens (1.3 to 2.5 % faster there, missed on 208 SMs) | deterministic (4 splits) | 256 experts, hidden 5120 and 6144 at 16172 to 16384 tokens: `"H0.5x4"` (the head/tail program with the BF16 parts of `A` in tensor memory, `ROUTER_HEAD_TAIL_TEMPLATES_BY_SM`). Every other token count, and hidden 4096: deterministic (6 splits): the SIMT windows (hidden 4096 and 5120 from 1 to 256 tokens, hidden 6144 from 1 to 192) missed the FP64 error rule (the router FP64 error table below); next to a resident kernel `"H0.5x4"` missed at hidden 4096 with both head/tail programs, and at hidden 6144 with the three-chunk-buffer program | *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 (the FP32 ones up to 960 tokens only: from 961 tokens up their FP64 error is about twice the -deterministic program's, the `sm_107a` K1 token table; on 212-SM and 216-SM parts these FP32 rows run three -accumulation chains at one split at the ranges of the K1 accumulation-chain table). Every K1 choice also needs +of the row is slower (the FP32 ones on 216 SMs and more up to 960 tokens only: from 961 tokens up they are unchanged, +qualification on a 216-SM part pending, the `sm_107a` K1 token table; on 212-SM and 216-SM parts these FP32 rows run +three accumulation chains at one split at the ranges of the K1 accumulation-chain table). Every K1 choice also needs `num_splits * 64 <= T + 63` (shorter token 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, and parts of exactly 148 SMs; `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`), and the 212-SM choices -of the rows of the 212-SM K1 token-range table also at its ranges.* +tables below run the deterministic program (`sm_100a` / `sm_103a`: parts of 152 SMs or more, and parts of exactly 148 +SMs; `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`), and the 212-SM choices of the rows of the 212-SM K1 token-range table also at its +ranges.* | K1 row `d_in -> d_out` | `sm_100a` BF16 | `sm_100a` FP32 | `sm_103a` BF16 | `sm_103a` FP32 | |---|---|---|---|---| | `4096 -> 32` | 512; 65536 and up† | -- | 512; 65536 and up§ | 512; 32768 and up§ | | `4096 -> 48` | 512; 49153 and up† | 512 | 512; 32768 and up§ | 512; 32768 and up§ | | `4096 -> 64` | 512; 65536 and up† | 512 | no verdict | no verdict | -| `5120 -> 32` | 512; 65537 and up† | 512 | 512; 32768 and up‡ | 512 | -| `5120 -> 48` | 512 to 2048; 49153 and up† | 512 to 2048 | 512 to 2048; 32768 and up‡ | 512 to 2048 | -| `5120 -> 64` | 512 to 2048; 8192 | 512 to 2048; 8192 | 512 to 2048; 32768 and up‡ | 512 to 2048 | -| `5120 -> 128` | no verdict | no verdict | 512 to 8192; 32768 and up‡ | 512 to 4096 | +| `5120 -> 32` | 512; 65537 and up† | 512 | 512; 45058 and up‡ | 512 | +| `5120 -> 48` | 512 to 2048; 49153 and up† | 512 to 2048 | 512 to 2048; 40961 and up‡ | 512 to 2048 | +| `5120 -> 64` | 512 to 2048; 8192 | 512 to 2048; 8192 | 512 to 2048; 45058 and up‡ | 512 to 2048 | +| `5120 -> 128` | no verdict | no verdict | 512 to 8192; 32768 and up¶ | 512 to 4096 | | `6144 -> 32` | 1024 | -- | 512 | 512 | | `6144 -> 128` | no verdict | no verdict | 512 to 8192 | 512 to 8192 | @@ -171,33 +174,45 @@ process per row, on two parts per token count (one part for `sm_100a` `6144 -> 3 error of 5 seeds on each of two parts per architecture (prepared launches, CUDA graph replays, the convenience call) next to the deterministic program's on the same operands. A measured token count runs the deterministic program where the split program failed in at least one run (the K1 rule above; the backward pass or the checkpointed step more than -1.3 % slower than with the deterministic program; or the deterministic package's FP64 error budget missed with an error -above the deterministic program's) and the deterministic program leaves the backward pass and the checkpointed step no -more than 1.3 % slower (medians and paired medians, eager launches and graph replays), or where the split program's -error exceeded the deterministic program's by more than 2 % for the same seed on both parts (relative Frobenius error or -BF16 error-budget ratio; marked ‡). A token count between two consecutive such measured token counts runs the -deterministic program too (not measured), and so does every token count above 65536 of a row marked ‡ at 65536 tokens. -Marked §: above 65536 tokens the split program of `sm_103a` `4096 -> 32` and `4096 -> 48` gives the deterministic -program's bits at an even number of 64-token K blocks and fails the FP64 comparison at an odd number, as on 148 SMs -(the 148-SM K1 token table below; the deterministic plan of these rows is the same on 148 and 152 SMs; on a 148-SM -part, 10 seeds per token count: an FP32 relative Frobenius error 1.021 to 1.024 times the deterministic program's at -65537 tokens; where the deterministic program misses the BF16 budget, a budget ratio above 1.30 times the deterministic -program's for some seed: `4096 -> 32` 1.305 at 65600 tokens (the 148-SM table), `4096 -> 48` up to 1.379 at 65537 -tokens), so these rows run the deterministic program from their entry up. Marked †: `sm_100a` BF16 entries from 148 -SMs up (also in the 148-SM table below; `4096 -> 64` from 152 SMs up, its 148-SM entry is below) where the -deterministic program misses the BF16 budget for some seed and the split program fails the rule of the 148-SM table -there (10 seeds per token count on a 148-SM B200 part; the deterministic plan of these rows is the same on 148 and 152 -SMs): a budget ratio above 1.30 times the deterministic program's for some seed at `4096 -> 48` and `5120 -> 48` 49153 -tokens (1.473, 1.320), `5120 -> 32` 65537 tokens (1.328) and `4096 -> 32` 65600 tokens (1.305, up to 1.727 at 106497 -tokens; the entry starts at 65536 as on `sm_103a`), and for `4096 -> 64` a median ratio of 1.041 over its odd token -counts from 65537 tokens; so these rows run the deterministic program from their entry up. -Token counts below 512, above 65536 (other rows), between two measured token counts with different outcomes, and on -parts with fewer than 152 SMs were not measured and keep the verdict. Measured token counts that failed the rule but -keep the verdict, because the deterministic program makes the backward pass or the checkpointed step more than 1.3 % -slower there: BF16 output at 1024 tokens on `sm_103a` `4096 -> 48`, `6144 -> 32`; BF16 output at 2048 tokens on -`sm_100a` `5120 -> 32`, `6144 -> 32`; `sm_103a` `4096 -> 48`; FP32 output at 512 tokens on `sm_100a` `4096 -> 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.* +1.3 % slower than with the deterministic program; or the deterministic package's FP64 error budget missed for some seed) +and the deterministic program leaves the backward pass and the checkpointed step no more than 1.3 % slower (medians and +paired medians, eager launches and graph replays), or where the split program's output missed the deterministic +package's FP64 error budget for some seed (marked ‡, †). The ratio of the split program's error to the deterministic +program's (q) is data, not part of the rule. A token count between two consecutive such measured token counts runs the +deterministic program too (not measured). Marked ‡: the BF16 output of `sm_103a` `5120 -> 32` and `5120 -> 64` met the +BF16 budget on every seed at 32768 tokens (30 seeds on two parts) and at every measured token count up to 45057 tokens +on a 152-SM GB300 part (10 seeds, 20 runs each); at the next measured token count, 47105, both programs missed it for +some seed, so the row keeps the deterministic program from the token count after the last passing one (the split program +misses it on every seed at 65536 tokens, 30 seeds on two parts). `5120 -> 48` met it on that part on every seed up to +49152 tokens (the deterministic program missed it for one seed at 45057 and at 49152 tokens), but on a 148-SM B300 part, +which runs the same program, it missed it for one of 10 seeds at 40961 and at 40962 tokens while the deterministic +program met it (the 148-SM table below), so its entry starts at 40961 tokens; from 49153 tokens both programs miss it on +the 152-SM part. Marked ¶: `sm_103a` `5120 -> 128` BF16 is not restored from 32768 tokens. Its split program met the +BF16 budget up to 45057 tokens as the rows marked ‡ did, but a strict timing re-measurement on that part failed the K1 +rule at 5 of 11 token counts from 32768 to 45057 (32768, 33793, 38913, 43009 and 44033: at most 4 of 8 paired groups +faster, up to 1.5 % slower at the operation) and kept it at 34817, 36865, 39937, 40960, 41985 and 45057, so no measured +window holds. Marked §: the entries of `sm_103a` `4096 -> 32` and `4096 -> 48` reach above 65536 tokens (up to 65536 +tokens they come from the rule above). There neither program meets the BF16 budget on any seed (10 seeds per token +count from 65537 tokens up on a 148-SM part; the deterministic plan of these rows is the same on 148 and 152 SMs), and +the FP32 split program gives the deterministic program's bits at an even number of 64-token K blocks; at an odd number +it met the FP32 budget on every seed at 65537 tokens (q 1.021 to 1.024) and both programs miss it from 98305 tokens up. +It was not measured faster there, so these rows keep the deterministic program from their entry up. Marked †: `sm_100a` +BF16 entries from 148 SMs up (also in the 148-SM table below; `4096 -> 64` from 152 SMs up, its 148-SM entry is below) +where the split program misses the BF16 budget for some seed (10 seeds per token count on a 148-SM B200 part; the +deterministic plan of these rows is the same on 148 and 152 SMs): at 49153 tokens 8 of 10 seeds met it for `4096 -> 48` +and 9 for `5120 -> 48`; at 65536 tokens 3 for `4096 -> 32`; from 65537 tokens up none. The deterministic program misses +it there too (`4096 -> 48` 9 of 10 seeds and `5120 -> 48` 8 of 10 at 49153 tokens, `4096 -> 32` 3 of 10 at 65536 tokens +and none from 65537 tokens up), so these rows keep the deterministic program from their entry up. `5120 -> 32` (10 +seeds, 20 runs each, at 32769 to 65025 tokens, finer from 45569 to 49665 tokens, on a 148-SM B200 part and two 152-SM +GB200 parts): both programs met the budget on every seed up to 46593 tokens, and from 47105 tokens each missed it for +some seed at some token counts, neither more accurate at every one, so the row keeps its entry from 65537 tokens. Token +counts below 512, above 65536 (other rows), between two measured token counts with different outcomes, and on parts with +fewer than 152 SMs were not measured and keep the verdict. Measured token counts that failed the rule but keep the +verdict, because the deterministic program makes the backward pass or the checkpointed step more than 1.3 % slower +there: BF16 output at 1024 tokens on `sm_103a` `4096 -> 48`, `6144 -> 32`; BF16 output at 2048 tokens on `sm_100a` +`5120 -> 32`, `6144 -> 32`; `sm_103a` `4096 -> 48`; FP32 output at 512 tokens on `sm_100a` `4096 -> 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` | `sm_100a` BF16 | `sm_100a` FP32 | `sm_103a` BF16 | `sm_103a` FP32 | |---|---|---|---|---| @@ -205,107 +220,107 @@ slower there: BF16 output at 1024 tokens on `sm_103a` `4096 -> 48`, `6144 -> 32` | `4096 -> 48` | 512; 49153 and up† | 512 | no verdict | no verdict | | `4096 -> 64` | 512; 65536 and up‡ | 512; 65536 and up‡ | no verdict | no verdict | | `5120 -> 32` | 65537 and up† | -- | -- | -- | -| `5120 -> 48` | 512 to 2048; 49153 and up† | 512 | 512 to 2048 | 512 to 2048 | -| `5120 -> 64` | 512 to 2048; 8192 | 512 to 1024; 8192 | 512 to 1024 | 512 to 1024 | +| `5120 -> 48` | 512 to 2048; 49153 and up† | 512 | 512 to 2048; 40961 to 49151 | 512 to 2048 | +| `5120 -> 64` | 512 to 2048; 8192 | 512 to 1024; 8192 | 512 to 1024; 47105 to 49152 | 512 to 1024 | | `5120 -> 128` | no verdict | no verdict | 512 to 4096; 32768 to 65536 | 512 to 4096; 32768 to 65536 | | `6144 -> 32` | -- | -- | -- | -- | | `6144 -> 128` | no verdict | no verdict | 512 to 8192 | 512 to 8192 | -*Table: the 148-SM K1 token table: the K1 token counts that run the deterministic program on `sm_100a` / `sm_103a` -parts of exactly 148 SMs (`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`; †: the entries of the 152-SM table above that hold -from 148 SMs up; `--`: none; `no verdict`: the row runs the deterministic program at every token count, `sm_103a` -`4096 -> 48` below its 152-SM floor). Every other token count -keeps the row's verdict, and parts of any other SM count keep their routing. Measured with the protocol and the rule of -the 152-SM table above at the same token counts, on B300 (`sm_103a`) and B200 (`sm_100a`) parts of 148 SMs: a screen of -every split-routed row and token count on one part (9 paired groups, one long-lived process per GPU), the points that -failed there again on a second part (15 paired groups in a long-lived process and in a fresh process per point; B300: -a part of a second node, which also screened every point; B200: another part of the same node), and the FP64 error of -5 seeds on two parts at 32768 and 65536 tokens and at the first and the last token count of every entry and one token -either side of it (98304 and 131072 tokens too for the entries that reach 65536 tokens). ‡: entries from the FP64 -error. The split program fails the FP64 comparison at a token count where, for some seed, its relative Frobenius -error is more than 2 % above the deterministic program's, or, where the deterministic program meets the error budget -on every seed, its output is outside the budget. Where the deterministic program misses the budget for some seed -(BF16 output at the largest token counts), the BF16 error-budget ratio, a maximum over the elements, moves by up to -about 25 % either way between the two programs at the same relative Frobenius error; there the split program fails -where the median of its ratio to the deterministic program's over a row's measured token counts and seeds (10 seeds -per token count) exceeds 1.02, or where one seed's ratio exceeds 1.30. BF16 output of `5120 -> 32` (23 token counts -from 57336 to 131072 tokens: median 1.000, largest 1.26) and `5120 -> 128` (65536 to 131072 tokens: 1.015, largest -1.15) on `sm_103a` passed it. Above 65536 tokens the split program of `4096 -> 32` (`sm_103a`) and `4096 -> 64` -(`sm_100a`) gives other bits than the deterministic program where the token count spans an odd number of 64-token -blocks (measured at 65537, 65538, 65600 and 65665 tokens; `4096 -> 64` BF16 on `sm_100a` also at 11 more odd counts -from 65728 to 98240 tokens) and the same bits where it spans an even number (65601 to 131072 tokens, 12 counts measured). At -those odd token counts the FP32 output's relative Frobenius error is 2.1 to 2.3 % above the deterministic program's on -every seed, and the BF16 budget ratio exceeds 1.30 times the deterministic program's for some seed (`4096 -> 32` on -`sm_103a`: 1.305 at 65600 tokens; `4096 -> 64` on `sm_100a`: up to 1.65, at 4 of 12 odd counts from 65600 to 98240 -tokens, median 1.000). At the even counts the split program does not make the backward pass or the checkpointed step -1 % faster (`sm_100a`: the backward pass up to 3.5 % slower). So these rows run the deterministic program from 65536 -tokens up. Unmeasured token counts -keep the verdict. Measured token counts that failed the rule but keep the verdict, because the -deterministic program makes the backward pass or the checkpointed step more than 1.3 % slower there on a part: -`sm_100a`: BF16 output at 512 tokens on `4096 -> 32`, `5120 -> 32`, `6144 -> 32`, at 1024 on `5120 -> 32`, at 2048 on -`6144 -> 32`, at 32768 on `5120 -> 48`, at 65536 on `4096 -> 32`, `4096 -> 48`, `5120 -> 48`, `5120 -> 64`, `6144 -> 32`; -FP32 output at 512 tokens on `4096 -> 32`, at 1024 and 2048 on `5120 -> 48`, at 2048 on `5120 -> 64`, at 65536 on -`5120 -> 48`; `sm_103a`: BF16 output at 512 tokens on `4096 -> 32`, `5120 -> 32`, `6144 -> 32`, at 2048 on -`5120 -> 64`, at 8192 on `5120 -> 128`, at 65536 on `5120 -> 48`, `5120 -> 64`, `6144 -> 32`, `6144 -> 128`; FP32 -output at 512 tokens on `6144 -> 32`, at 2048 on `5120 -> 64`, at 8192 on `5120 -> 128`, at 32768 on `4096 -> 32`.* +*Table: the 148-SM K1 token table: the K1 token counts that run the deterministic program on `sm_100a` / `sm_103a` parts +of exactly 148 SMs (`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`; †: the entries of the 152-SM table above that hold from 148 +SMs up; `--`: none; `no verdict`: the row runs the deterministic program at every token count, `sm_103a` `4096 -> 48` +below its 152-SM floor). Every other token count keeps the row's verdict, and parts of any other SM count keep their +routing. Measured with the protocol and the rule of the 152-SM table above at the same token counts, on B300 +(`sm_103a`) and B200 (`sm_100a`) parts of 148 SMs: a screen of every split-routed row and token count on one part (9 +paired groups, one long-lived process per GPU), the points that failed there again on a second part (15 paired groups in +a long-lived process and in a fresh process per point; B300: a part of a second node, which also screened every point; +B200: another part of the same node), and the FP64 error of 5 seeds on two parts at 32768 and 65536 tokens and at the +first and the last token count of every entry and one token either side of it (98304 and 131072 tokens too for the +entries that reach 65536 tokens). ‡: entries from the FP64 error and the time. `sm_103a` `5120 -> 48` BF16 at 40961 to +49151 tokens and `5120 -> 64` BF16 at 47105 to 49152 tokens: from the FP64 error alone, measured with 10 seeds (20 runs +each) at 36865 to 61441 tokens: the split program missed the BF16 budget for one of 10 seeds at 40961 and at 40962 +tokens (`5120 -> 64`: 47105) while the deterministic program met it on every seed (`5120 -> 48` met it at 41985 to 47105 +tokens inside its range); from 49152 tokens (`5120 -> 64`: 49153) both programs miss it for some seed, so the split +program stays there. Above 65536 tokens the split program of `4096 -> 32` (`sm_103a`) and `4096 -> 64` (`sm_100a`) +gives other bits than the deterministic program where the token count spans an odd number of 64-token blocks (measured +at 65537, 65538, 65600 and 65665 tokens; `4096 -> 64` BF16 on `sm_100a` also at 11 more odd counts from 65728 to 98240 +tokens) and the same bits where it spans an even number (65601 to 131072 tokens, 12 counts measured). There the BF16 +output of both programs misses the BF16 budget for some seed, and the FP32 output of both misses the FP32 budget at an +odd count from 98305 tokens up (q 1.021 to 1.023 at the measured odd counts). At the even counts the split program does +not make the backward pass or the checkpointed step 1 % faster (`sm_100a`: the backward pass up to 3.5 % slower). So +these rows run the deterministic program from 65536 tokens up. The BF16 output of `5120 -> 32` (23 token counts from +57336 to 131072 tokens) and `5120 -> 128` (65536 to 131072 tokens) on `sm_103a` keeps the split program on 148 SMs (q +median 1.000 and 1.015, largest 1.26 and 1.15); at 65537 tokens `5120 -> 128` missed the BF16 budget on every seed with +both programs (10 seeds), and `5120 -> 64` missed it for some seed with both at 47104 tokens (split 8 of 10 seeds, +deterministic 9). Unmeasured token counts keep the verdict. Measured token counts that failed the rule but keep the +verdict, because the deterministic program makes the backward pass or the checkpointed step more than 1.3 % slower there +on a part: `sm_100a`: BF16 output at 512 tokens on `4096 -> 32`, `5120 -> 32`, `6144 -> 32`, at 1024 on `5120 -> 32`, at +2048 on `6144 -> 32`, at 32768 on `5120 -> 48`, at 65536 on `4096 -> 32`, `4096 -> 48`, `5120 -> 48`, `5120 -> 64`, +`6144 -> 32`; FP32 output at 512 tokens on `4096 -> 32`, at 1024 and 2048 on `5120 -> 48`, at 2048 on `5120 -> 64`, at +65536 on `5120 -> 48`; `sm_103a`: BF16 output at 512 tokens on `4096 -> 32`, `5120 -> 32`, `6144 -> 32`, at 2048 on +`5120 -> 64`, at 8192 on `5120 -> 128`, at 65536 on `5120 -> 48`, `5120 -> 64`, `6144 -> 32`, `6144 -> 128`; FP32 output +at 512 tokens on `6144 -> 32`, at 2048 on `5120 -> 64`, at 8192 on `5120 -> 128`, at 32768 on `4096 -> 32`.* | K1 row `d_in -> d_out` | parts | token counts with an odd number of 64-token K blocks that run the deterministic program | |---|---|---| -| `sm_100a` `4096 -> 32` FP32 | 148 SMs and more | 1025 and up | -| `sm_100a` `4096 -> 48` FP32 | 148 SMs and more | 1025 and up | -| `sm_100a` `4096 -> 64` FP32 | 148 SMs and more | 1025 and up | -| `sm_100a` `5120 -> 32` FP32 | 148 SMs and more | 12289 and up | -| `sm_100a` `5120 -> 48` FP32 | 148 SMs and more | 12289 and up | -| `sm_103a` `4096 -> 32` FP32 | 148 SMs and more | 1025 and up | -| `sm_103a` `4096 -> 48` FP32 | 152 SMs and more | 1025 and up | -| `sm_107a` `8192 -> 128` FP32 | 212 SMs; 216 SMs and more | 1025 to 8192 | +| `sm_100a` `4096 -> 32` FP32 | 148 SMs and more | 65538 and up | +| `sm_100a` `4096 -> 48` FP32 | 148 SMs and more | 65538 and up | +| `sm_100a` `4096 -> 64` FP32 | 148 SMs and more | 65538 and up | +| `sm_100a` `5120 -> 32` FP32 | 148 SMs and more | 65538 and up | +| `sm_100a` `5120 -> 48` FP32 | 148 SMs and more | 65538 and up | +| `sm_103a` `4096 -> 32` FP32 | 148 SMs and more | 65538 and up | +| `sm_103a` `4096 -> 48` FP32 | 152 SMs and more | 14338 to 32767‡; 65538 and up | +| `sm_107a` `8192 -> 128` FP32 | 216 SMs and more | 1025 to 8192§ | *Table: the odd-block K1 table (`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS`): these rows run the deterministic program at a token count of the third column that spans an odd number of 64-token K blocks (`ceil(T / 64)` odd); an even number, and every other row, keeps the routing of the tables above. In these ranges the deterministic program of the row cuts the token reduction of its tiles, taken in order, into units of half the tile's K blocks, rounded up (at an odd block count -a unit runs on into the next tile; its automatic two-way split), and adds the partials in a fixed order: the `4096` rows -from 961 tokens (16 K blocks) up; `sm_100a` `5120 -> 32` and `5120 -> 48` from where one wave's operands outgrow the L2 -(from 12289 tokens with 126 MiB of L2, between 12292 and 12337 tokens on a B200; below, whole tiles); `sm_107a` -`8192 -> 128` while they fit the L2 (to 7808 tokens with 126 MiB of L2, to between 4097 and 8192 tokens on a 212-SM -R200; above, a synchronised two-way split that the two-split program matches). At an even block count the two-split -program's output is the deterministic program's bit for bit; at an odd count its FP32 output's FP64 relative Frobenius -error was more than 2 % above the deterministic program's on every seed at every measured odd token count of the -ranges from 1025 to 65537 tokens (10 seeds each): the `4096` rows (`sm_100a` on a 148-SM B200 part, `sm_103a` on a -148-SM B300 part; the deterministic plan of these rows is the same on 148 and 152 SMs) 1.03 to 1.15 times at 1025, -1088, 2049, 3000, 4097, 8193, 12289, 16173, 24577, 32769 and 49153 tokens and 1.021 to 1.024 times at 65537; -`sm_100a` `5120 -> 32` and `5120 -> 48` 1.025 to 1.075 times at 12337 to 32769 tokens; `sm_107a` `8192 -> 128` 1.04 to -1.14 times at 1025, 1088, 2049 and 4097 tokens (212-SM R200 parts, the 212-SM and the 216-SM plan; the deterministic -plan of the row is the same on 212 and 216 SMs). The excess falls with the token count (`5120 -> 32` and `5120 -> 48`: -1.7 to 1.96 % at 49153 tokens, 0.7 to 1.5 % at 65537 to 131073 tokens; the `4096` rows: 1.2 to 1.7 % at 98305 and -131073 tokens); their ranges stay open because the two-split program is not more accurate than the deterministic -program at any odd block count in them. The BF16 outputs of these rows had the deterministic program's relative -Frobenius error and met the error budget on every seed where the deterministic program met it.* - -| K1 row `d_in -> d_out` (FP32 out) | 152 SMs: token ranges, even K-block counts only (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`) | measured token counts outside them | +a unit runs on into the next tile; its automatic two-way split), and adds the partials in a fixed order. At an even +block count the two-split program's output is the deterministic program's bit for bit. At an odd count its FP32 output +met the deterministic package's FP32 budget on every seed at every measured odd token count from 1025 to 65537 tokens +(10 seeds each: the `4096` rows at 1025, 1088, 2049, 3000, 4097, 8193, 12289, 16173, 24577, 32769, 49153 and 65537 +tokens, `sm_100a` on a 148-SM B200 part and `sm_103a` on a 148-SM B300 part, the deterministic plan of these rows being +the same on 148 and 152 SMs; `sm_100a` `5120 -> 32` and `5120 -> 48` at 12337 to 65537 tokens), so those token counts +run the two-split program; q was 1.03 to 1.15, 1.021 to 1.024 at 65537 tokens and 1.007 to 1.075 for the `5120` rows +(data). From 98305 tokens up both programs miss the FP32 budget at an odd count; no odd token count from 65538 to 98304 +was measured, and the rows run the deterministic program at an odd count from 65538 tokens up. Marked ‡: a range from +the time, not the budget: on a 152-SM GB300 part, at odd block counts the two-split program of `sm_103a` `4096 -> 48` +left the checkpointed step 2.0 to 2.3 % slower in graph replays at 16173 tokens (two measurements), the backward pass +1.6 and 2.0 % slower in eager launches at 20481 and 28673 tokens, and failed the K1 rule in one of 8 groups at 18433 +tokens, while it kept the rule and the 1.3 % bound at 1025, 4097, 6145, 8193, 10241, 12289 and 14337 tokens. Marked §: +`sm_107a` `8192 -> 128` FP32 keeps its range on 216 SMs and more only, where its deterministic program's two-way split +runs while one wave's operands fit the L2 (to 7808 tokens with 126 MiB of L2): at odd block counts from 1025 to 8191 +tokens its two-split program met the FP32 budget on every seed (10 seeds on 212-SM R200 parts with the 212-SM and the +216-SM plan; q 1.04 to 1.14 at 1025, 1088, 2049 and 4097 tokens), and its time there was measured on 212-SM parts only, +so the range is unchanged (qualification on a 216-SM part pending). The BF16 outputs of these rows had the +deterministic program's relative Frobenius error and met the error budget on every seed where the deterministic program +met it.* + +| K1 row `d_in -> d_out` (FP32 out) | 152 SMs: token ranges, every K-block count (`K1_TWO_SPLIT_RANGES_BY_SM`) | measured token counts outside them | |---|---|---| | `sm_100a` `4096 -> 128`: (64, 2) | 1024 to 4096 | 512; 6144; 8192 | | `sm_103a` `4096 -> 128`: (64, 2) | 1024 to 6144 | 512; 8192 | -*Table: the 152-SM two-split table (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` on `sm_100a` / `sm_103a` parts of exactly 152 -SMs, above the SM-count floor): the FP32 row `4096 -> 128`, which has no verdict in `K1_SPLIT_RULES`, runs two splits of -64-wide tiles, each adding half of an output's K blocks, at a token count of its range that spans an even number of -64-token K blocks (`ceil(T / 64)` even). At an odd number, outside the range, with BF16 output and on parts of any -other SM count (148 SMs included) it runs the deterministic program. At an even block count its output was the -deterministic program's bit for bit (5 seeds at 1024, 2048, 4096, 6144 and 8192 tokens); at 1088 and 2112 tokens (17 -and 33 blocks) it was not, with an FP64 relative Frobenius error 1.107 and 1.142 times the deterministic program's (5 -seeds on three parts). Measured at 512, 1024, 2048, 4096 and 8192 tokens on two 152-SM parts of each GPU (two GB200 -nodes, two GB300 nodes) and at 6144 tokens on one GB200 part and two GB300 parts: 15 paired groups of 20 launches in -four runs per part (two launch orders, a long-lived process and a fresh process per row), 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. A range runs between measured token counts at which the two-split program passed the K1 rule in every run -on every part, left the backward pass and the checkpointed step no more than 1.3 % slower than the deterministic -program, and met the FP64 rule of the 148-SM table above (an error no more than 2 % above the deterministic program's -for every seed); the token counts between two consecutive measured ones inside it take the verdict too (not -measured). From 1024 to 4096 tokens its time was 0.70 to 0.79 times the deterministic program's (graph replays, both -GPUs; 8.65 against 12.14 us on a GB200 part at 1024 tokens). At 512 tokens it is slower than the deterministic -program, and at 8192 tokens slower next to a kernel holding 32 SMs; at 6144 tokens on the GB200 part its time next to -a kernel holding SMs on another stream was 1.099 times the deterministic program's.* +*Table: the 152-SM two-split table (`K1_TWO_SPLIT_RANGES_BY_SM` on `sm_100a` / `sm_103a` parts of exactly 152 SMs, above +the SM-count floor): the FP32 row `4096 -> 128`, which has no verdict in `K1_SPLIT_RULES`, runs two splits of 64-wide +tiles, each adding half of an output's K blocks, at every token count of its range, at either number of 64-token K +blocks. Outside the range, with BF16 output and on parts of any other SM count (148 SMs included) it runs the +deterministic program. At an even block count its output was the deterministic program's bit for bit (5 seeds at 1024, +2048, 4096, 6144 and 8192 tokens); at 1088 and 2112 tokens (17 and 33 blocks) it was not, with q 1.107 and 1.142 (5 +seeds on three parts; data: the deterministic program's FP64 relative Frobenius error there is about 1 % of the FP32 +budget). Measured at 512, 1024, 2048, 4096 and 8192 tokens on two 152-SM parts of each GPU (two GB200 nodes, two GB300 +nodes) and at 6144 tokens on one GB200 part and two GB300 parts: 15 paired groups of 20 launches in four runs per part +(two launch orders, a long-lived process and a fresh process per row), 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. A range +runs between measured token counts at which the two-split program passed the K1 rule in every run on every part, left +the backward pass and the checkpointed step no more than 1.3 % slower than the deterministic program, and met the +deterministic package's FP64 error budget for every seed; the token counts between two consecutive measured ones inside +it, odd block counts included, take the verdict too (not measured). From 1024 to 4096 tokens its time was 0.70 to 0.79 +times the deterministic program's (graph replays, both GPUs; 8.65 against 12.14 us on a GB200 part at 1024 tokens). At +512 tokens it is slower than the deterministic program, and at 8192 tokens slower next to a kernel holding 32 SMs; at +6144 tokens on the GB200 part its time next to a kernel holding SMs on another stream was 1.099 times the deterministic +program's.* | K1 row `d_in -> d_out` (FP32 out unless marked BF16) | 208 SMs | 212 SMs | 216 SMs | |---|---|---|---| @@ -313,58 +328,65 @@ a kernel holding SMs on another stream was 1.099 times the deterministic program | `4096 -> 32` BF16 | -- | -- | 512¶ | | `4096 -> 48` | 16231; 16384 | -- | 512¶; 16231; 16384 | | `4096 -> 64` | 16231; 16384 | -- | 512¶; 16231; 16384 | -| `4096 -> 512` | 16172 to 16384§ | 16172 to 16384§ | 961 and up§ | -| `4096 -> 1024` | 16172 to 16384§ | 16172 to 16384§ | 961 and up§ | +| `4096 -> 512` | 16231; 16384 | 16231; 16384 | 961 and up§ | +| `4096 -> 1024` | 16231; 16384‡ | 16231 | 961 and up§ | | `5120 -> 32` | -- | -- | 512¶ | | `5120 -> 32` BF16 | -- | -- | 512¶ | | `5120 -> 48` | 16231; 16384 | 16231 | 512¶; 16231; 16384 | | `5120 -> 64` | 16231 | 16231 | 512¶; 16231; 16384 | -| `5120 -> 512` | 16172 to 16384§ | 16172 to 16384§ | 961 and up§ | -| `5120 -> 1024` | 16172 to 16384§ | 16172 to 16384§ | 961 and up§ | +| `5120 -> 512` | -- | 16231; 16384 | 961 and up§ | +| `5120 -> 1024` | -- | -- | 961 and up§ | | `5120 -> 1024` BF16 | -- | -- | 16231 | | `6144 -> 32` BF16 | -- | -- | 512¶ | | `8192 -> 128`, BF16 and FP32 | -- | -- | 512¶; 65536¶ | *Table: the K1 token counts that run the deterministic program on `sm_107a` parts (`K1_SPLIT_RULE_EXCLUDED_TOKENS`; -`--`: none; `8192 -> 128` FP32 also at odd block counts from 1025 to 8192 tokens on 212 and 216 SMs: the odd-block K1 -table above). Unmarked ones 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, 212 -and 216 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), 216 SMs: two parts. -A measured token count runs the deterministic program where the row's program failed the K1 rule above in some run, or +`--`: none; `8192 -> 128` FP32 also at odd block counts from 1025 to 8192 tokens on 216 SMs: the odd-block K1 table +above). Unmarked ones 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, 212 and 216 +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), 216 SMs: two parts. 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 entry of `5120 -> 1024` BF16 (16231 tokens) is not a token count of those rules: the deterministic plan there is -the row's plan at every token count, and the row's program left the backward pass 1.5 % slower in eager launches (ratio -of the medians over the runs). Marked §: entries of their own, not from the rules above. These FP32 rows run the -deterministic program from 961 tokens up on parts of at least 216 SMs, and at 16172 to 16384 tokens on 208 and 212 SMs, -because their one-split program fails the FP64 comparison there: it adds the whole token reduction of an output in one +SMs and `5120 -> 1024` on 208 SMs were measured twice; the rule takes the runs of both measurements together. On 216 +SMs the 192-wide row-major program of `4096 -> 1024` FP32 was measured in the same runs against the 256-wide tiles and +the deterministic plan: faster than both at 16172 and 16384 tokens in every run, not faster than the deterministic +package's 16231-token plan; it is not shipped, these token counts lying in the row's § range. The 216-SM entry of +`5120 -> 1024` BF16 (16231 tokens) is not a token count of those rules: the deterministic plan there is the row's plan +at every token count, and the row's program left the backward pass 1.5 % slower in eager launches (ratio of the medians +over the runs). Marked ‡: an entry of a later strict re-measurement of the restored window on a 208-SM R200 part: at +16384 tokens the deterministic program's plan (160-wide tiles) was faster; the row's program took 61.27 against 53.49 us +in graph replays (8 of 8 groups failed the K1 rule) and left the checkpointed step 6.3 % slower in graph replays. +Marked §: entries of their own, not from the rules above. On parts of at least 216 SMs these FP32 rows run the +deterministic program from 961 tokens up. Their one-split program adds the whole token reduction of an output in one chain, while from 961 tokens (16 K blocks) up the deterministic program of these rows cuts that reduction into units of -half an output's K blocks, rounded up (its automatic two-way split), and adds the partials in a fixed order. The -one-split program's FP64 relative Frobenius error is 2.0 to 2.4 times the deterministic program's on every seed from 961 -tokens up, above the 2 % the rule allows (and from 45056 tokens up above the deterministic package's FP32 error budget, -5e-5); up to 960 tokens the deterministic program adds one chain as well, and the two errors are equal. Measured with -the 216-SM plan on 212-SM R200 parts (the same program and accumulation order as on 216 SMs; the deterministic plan of -these rows is the same on 208, 212 and 216 SMs; the error does not depend on the SM count) against an FP64 reference, -each row on two parts, 5 seeds, at 512 to 65536 tokens, and in a later run with 10 seeds at 704, 896, 960, 961 and 1024 -tokens: the same bits as the deterministic program up to 960 tokens, 2.09 to 2.11 times its error at 961 and 1024 -tokens. The 212-SM two-split ranges of these rows (the 212-SM K1 token-range table below) lie outside 16172 to 16384 -tokens and keep their program. Marked ¶: single token counts from a re-measurement of chosen points at 512, 8192, 32768 -and 65536 tokens on two 216-SM R200 parts per point (the protocol above: 15 paired groups, eager launches and graph -replays, next to a kernel holding 8, 16 or 32 SMs in both launch orders, the backward pass and the checkpointed step -timed). A point is an entry where the row's program failed the K1 rule in some run and the deterministic program leaves -the backward pass and the checkpointed step no more than 1.3 % slower on both parts (the worst of the ratio of the -medians and the median of the paired ratios, eager launches and graph replays: the rule of the `sm_100a` / `sm_103a` -entries); the neighbouring token counts were not measured and keep the verdict. Measured points that failed but keep -the verdict, because the deterministic program makes the backward pass or the checkpointed step more than 1.3 % slower -there on a part: BF16 output at 8192 tokens on `4096 -> 32`, at 32768 tokens on `4096 -> 32` and `5120 -> 1024`, at -65536 tokens on `5120 -> 1024`; FP32 output at 512 tokens on `4096 -> 32` and `6144 -> 32`.* +half an output's K blocks, rounded up (its automatic two-way split), and adds the partials in a fixed order; up to 960 +tokens both programs add one chain and give the same bits. Measured with the 216-SM plan on 212-SM R200 parts (the same +program and accumulation order as on 216 SMs; the error does not depend on the SM count) against an FP64 reference, each +row on two parts, 5 seeds, at 32768, 36864, 40960, 45056, 49152, 57344, 65536 and 131072 tokens, the one-split program's +output met the deterministic package's FP32 budget on every seed up to 40960 tokens (1.86e-5 at 16384 tokens, 3.77e-5 at +32768; q 2.0 to 2.4, data) and missed it on every seed from 45056 tokens up (7.6e-5 at 65536 tokens and 1.5e-4 at +131072; the deterministic program: 3.8e-5 and 7.6e-5). Its time from 961 to 40960 tokens on 216-SM parts has not been +measured under the rule above, so these entries are unchanged (qualification on a 216-SM part pending). At 65536 tokens +the deterministic program is slower: on a 212-SM part it takes 1.81 (`4096 -> 512`), 1.05 (`4096 -> 1024`), 1.72 +(`5120 -> 512`) and 1.02 (`5120 -> 1024`) times the one-split program's time (CUDA graph replays) and leaves the +backward pass or the checkpointed step up to 51.7, 4.8, 41.5 and 3.9 % slower. Marked ¶: single token counts from a +re-measurement of chosen points at 512, 8192, 32768 and 65536 tokens on two 216-SM R200 parts per point (the protocol +above: 15 paired groups, eager launches and graph replays, next to a kernel holding 8, 16 or 32 SMs in both launch +orders, the backward pass and the checkpointed step timed). A point is an entry where the row's program failed the K1 +rule in some run and the deterministic program leaves the backward pass and the checkpointed step no more than 1.3 % +slower on both parts (the worst of the ratio of the medians and the median of the paired ratios, eager launches and +graph replays: the rule of the `sm_100a` / `sm_103a` entries); the neighbouring token counts were not measured and keep +the verdict. Measured points that failed but keep the verdict, because the deterministic program makes the backward +pass or the checkpointed step more than 1.3 % slower there on a part: BF16 output at 8192 tokens on `4096 -> 32`, at +32768 tokens on `4096 -> 32` and `5120 -> 1024`, at 65536 tokens on `5120 -> 1024`; FP32 output at 512 tokens on +`4096 -> 32` and `6144 -> 32`.* | K1 row `d_in -> d_out` (FP32 out unless marked BF16) | 212 SMs: token ranges besides 16172 to 16384 | measured token counts outside them | |---|---|---| @@ -373,54 +395,54 @@ there on a part: BF16 output at 8192 tokens on `4096 -> 32`, at 32768 tokens on | `8192 -> 128`, BF16 and FP32 | 1024 to 16172; 16384 to 32768 | 512; 65536 | | `4096 -> 512`, `4096 -> 1024`, `5120 -> 512`, `5120 -> 1024` | 512 | 1024 to 8192; 32768; 65536 | | `5120 -> 1024` BF16: (256, 1) `p = 1` (`K1_SPLIT_RANGE_OVERRIDES_BY_SM`) | 512 to 12288 | 32768 | -| `4096 -> 512`, `5120 -> 512`: (128, 2), even K-block counts only (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`) | 1024 to 8192; 32768 to 65536 | 512 | -| `4096 -> 1024`: (192, 2), even K-block counts only (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`) | 2048 to 8192; 32768 to 65536 | 512; 1024 | -| `5120 -> 1024`: (256, 2), even K-block counts only (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`) | 1024 to 8192 | 512; 32768; 65536 | -| `4096 -> 128`: (64, 2), even K-block counts only (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`; no verdict in the 212-SM table) | 1024 to 7680 | 512; 640; 768; 896; 8000; 8192; 12288; 16384 | +| `4096 -> 512`, `5120 -> 512`: (128, 2), every K-block count (`K1_TWO_SPLIT_RANGES_BY_SM`) | 1024 to 8192; 32768 to 65536 | 512 | +| `4096 -> 1024`: (192, 2), every K-block count (`K1_TWO_SPLIT_RANGES_BY_SM`) | 2048 to 8192; 32768 to 65536 | 512; 1024 | +| `5120 -> 1024`: (256, 2), every K-block count (`K1_TWO_SPLIT_RANGES_BY_SM`) | 1024 to 8192 | 512; 32768; 65536 | +| `4096 -> 128`: (64, 2), every K-block count (`K1_TWO_SPLIT_RANGES_BY_SM`; no verdict in the 212-SM table) | 1024 to 7680 | 512; 640; 768; 896; 8000; 8192; 12288; 16384 | *Table: the 212-SM K1 token-range table (`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`): further token counts at which these rows of the 212-SM table of `K1_SPLIT_RULES_BY_SM` run their program (the third column runs the deterministic program), -except the token counts of the `sm_107a` K1 token table above. Measured at 512, 1024, 2048, 4096, 8192, 32768 and -65536 tokens on two 212-SM R200 nodes of four parts each, with the protocol of the window above (one part per launch -layout and row order), and the FP64 error of 5 seeds on two parts. A measured token count joins a range where the -row's program passed the K1 rule in every run, left the backward pass and the checkpointed step no more than 1.3 % -slower than the deterministic program, and met the deterministic package's FP64 error budget with an error no more than -2 % above the deterministic program's for every seed on both parts. A range ends at a measured token count that passed -(or at the window) and holds no measured token count that failed; the token counts between two consecutive measured -ones inside it take the verdict too (not measured). Token counts below 512, above 65536 and between a measured token -count that passed and one that failed run the deterministic program. The four one-split FP32 rows pass at 512 tokens -only: from 1024 tokens up their FP64 error is about twice the deterministic program's (the one-split program adds an -output's whole token reduction in one chain). The other rows are not faster than the deterministic program in every -run at 512 tokens, nor `8192 -> 128` at 65536. The two-split rows (`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`) run two splits -of the row's 212-SM tile width, each adding half of an output's K blocks, at a token count of their ranges that spans an -even number of 64-token K blocks (`ceil(T / 64)` even); at an odd number they keep the routing above (the deterministic -program). So the token counts in the third column of the one-split FP32 rows run the two-split program where a -two-split range holds them. At an even block count from 1024 tokens up their output was the deterministic program's bit -for bit (5 seeds on two parts at 1024, 2048, 4096, 8192, 16231, 16384, 32768 and 65536 tokens, on one part at 6000, -12000 and 48000); at an odd count from 17 blocks up it was not, with an FP64 relative Frobenius error 1.01 to 1.17 times -the deterministic program's (5 seeds, 1025 to 65537 tokens). They were measured with the protocol above at 512, 1024, -2048, 4096, 8192, 32768 and 65536 tokens (`4096 -> 512` also at 40961, which runs the deterministic program: an odd -block count) on two 212-SM R200 nodes of four parts each; a range runs between measured token counts at which the -two-split program passed the K1 rule in every run and left the backward pass and the checkpointed step no more than 1.3 -% slower than the deterministic program, and holds none that failed. At 512 tokens the two-split program is slower than -the deterministic plan; `4096 -> 1024` is not faster in every run at 1024 tokens; `5120 -> 1024` is not faster in every -run at 32768 tokens and slower at 65536. `4096 -> 128` FP32 has no program in the 212-SM table either; its two-split -range was measured with the same protocol at 512, 640, 768, 896, 1024, 2048, 4096, 6144, 6656, 7168, 7680, 8000, 8192, -12288 and 16384 tokens on two nodes of each of two R200 systems (6144, 8000, 8192, 12288 and 16384 tokens and 640 on -one system). From 1024 to 7680 tokens it passed in every run (0.70 to 0.76 times the deterministic program's time) and -its output was the deterministic program's bit for bit at an even block count (5 seeds); at an odd count (1088 to -5000 tokens) its FP64 error was 1.11 to 1.16 times the deterministic program's. It failed the rule at 512, 640, 768 -and 896 tokens; from 8000 tokens the deterministic plan changes and the two-split error was 1.54 times its error, and -at 12288 and 16384 tokens it was 5 to 7 % slower. `5120 -> 1024` BF16 has no program in the 212-SM table (at 16172 to -16384 tokens the deterministic plan is faster); its range runs the 256-wide one-split program of `K1_SPLIT_RULES` (its -216-SM verdict), measured at 512, 1024, 2048, 4096, 8192, 12288 and 32768 tokens with the same protocol on a node of -four parts, twice with the parts' launch layouts and orders rotated: faster than the deterministic plan in every run up to 12288 tokens, with the FP64 error of the -deterministic program on every seed; at 32768 tokens it failed the K1 rule in every run (0.1 to 0.6 % faster than the -deterministic plan alone), and its checkpointed step measured 2.5 to 4.3 % slower, as did a second deterministic launch -against the first there (1.9 to 4.4 %). On 212 SMs -`4096 -> 1024` FP32 keeps the 192-wide program at 512 tokens: the 256-wide program is 15 % faster alone and leaves the -backward pass 4 to 6 % faster, but the checkpointed step 0.2 to 1.2 % slower in graph replays (median of the paired -ratios, ratio of the medians).* +except the token counts of the `sm_107a` K1 token table above. Measured at 512, 1024, 2048, 4096, 8192, 32768 and 65536 +tokens on two 212-SM R200 nodes of four parts each, with the protocol of the window above (one part per launch layout +and row order), and the FP64 error of 5 seeds on two parts. A measured token count joins a range where the row's +program passed the K1 rule in every run, left the backward pass and the checkpointed step no more than 1.3 % slower than +the deterministic program, and met the deterministic package's FP64 error budget for every seed on both parts. A range +ends at a measured token count that passed (or at the window) and holds no measured token count that failed; the token +counts between two consecutive measured ones inside it take the verdict too (not measured). Token counts below 512, +above 65536 and between a measured token count that passed and one that failed run the deterministic program. The four +one-split FP32 rows take this table at 512 tokens only (their 212-SM program in the window as well); from 960 tokens the +K1 accumulation-chain table below runs three chains first, and the two-split ranges hold more token counts (q of the +one-split program about 2 from 1024 tokens up, data: it adds an output's whole token reduction in one chain). The other +rows are not faster than the deterministic program in every run at 512 tokens, nor `8192 -> 128` at 65536. The +two-split rows (`K1_TWO_SPLIT_RANGES_BY_SM`) run two splits of the row's 212-SM tile width, each adding half of an +output's K blocks, at every token count of their ranges, at either number of 64-token K blocks. So the token counts in +the third column of the one-split FP32 rows run the two-split program where a two-split range holds them. At an even +block count from 1024 tokens up their output was the deterministic program's bit for bit (5 seeds on two parts at 1024, +2048, 4096, 8192, 16231, 16384, 32768 and 65536 tokens, on one part at 6000, 12000 and 48000); at an odd count from 17 +blocks up it was not, with q 1.01 to 1.17 (data) and within the FP32 budget on every seed (5 seeds, 1025 to 65537 +tokens; relative Frobenius error at most 3.8e-5, at 65537 tokens). They were measured with the protocol above at 512, +1024, 2048, 4096, 8192, 32768 and 65536 tokens (`4096 -> 512` also at 40961, an odd block count) on two 212-SM R200 +nodes of four parts each; a range runs between measured token counts at which the two-split program passed the K1 rule +in every run and left the backward pass and the checkpointed step no more than 1.3 % slower than the deterministic +program, and holds none that failed; the token counts between, odd block counts included, take the verdict too (not +measured). At 512 tokens the two-split program is slower than the deterministic plan; `4096 -> 1024` is not faster in +every run at 1024 tokens; `5120 -> 1024` is not faster in every run at 32768 tokens and slower at 65536. `4096 -> 128` +FP32 has no program in the 212-SM table either; its two-split range was measured with the same protocol at 512, 640, +768, 896, 1024, 2048, 4096, 6144, 6656, 7168, 7680, 8000, 8192, 12288 and 16384 tokens on two nodes of each of two R200 +systems (6144, 8000, 8192, 12288 and 16384 tokens and 640 on one system). From 1024 to 7680 tokens it passed in every +run (0.70 to 0.76 times the deterministic program's time) and its output was the deterministic program's bit for bit at +an even block count (5 seeds); at an odd count (1088 to 5000 tokens) q was 1.11 to 1.16 (data), so these token counts +run it too. It failed the rule at 512, 640, 768 and 896 tokens; from 8000 tokens the deterministic plan changes and the +two-split error was 1.54 times its error, and at 12288 and 16384 tokens it was 5 to 7 % slower. `5120 -> 1024` BF16 has +no program in the 212-SM table (at 16172 to 16384 tokens the deterministic plan is faster); its range runs the 256-wide +one-split program of `K1_SPLIT_RULES` (its 216-SM verdict), measured at 512, 1024, 2048, 4096, 8192, 12288 and 32768 +tokens with the same protocol on a node of four parts, twice with the parts' launch layouts and orders rotated: faster +than the deterministic plan in every run up to 12288 tokens, with the FP64 error of the deterministic program on every +seed; at 32768 tokens it failed the K1 rule in every run (0.1 to 0.6 % faster than the deterministic plan alone), and +its checkpointed step measured 2.5 to 4.3 % slower, as did a second deterministic launch against the first there (1.9 to +4.4 %). On 212 SMs `4096 -> 1024` FP32 keeps the 192-wide program at 512 tokens: the 256-wide program is 15 % faster +alone and leaves the backward pass 4 to 6 % faster, but the checkpointed step 0.2 to 1.2 % slower in graph replays +(median of the paired ratios, ratio of the medians).* | K1 row `d_in -> d_out`, FP32 out | 212 SMs: `(block_n, chains)`, token ranges | 216 SMs: `(block_n, chains)`, token ranges | |---|---|---| @@ -437,17 +459,18 @@ programs were timed at (1024 to 8192 and 32768 to 40960 tokens) three chains ran checkpointed step 2 to 21 % faster than two splits (medians, eager launches and graph replays). The program cuts each tile's K blocks into three contiguous chains, each summed from an `init` MMA in an accumulator of its own, and its epilogue adds them once, in order: no workspace and the same bits on every run. Its FP64 relative Frobenius error was -0.31 to 0.73 times the deterministic program's on every seed (5 seeds on a 212-SM R200 part at 960, 961, 1024, 2048, -4096, 8192, 16172, 16231, 16384, 32768 and 40961 tokens, even and odd block counts), where the one-chain program's is -2.0 to 2.3 times. The ranges follow the rule and the construction of the 212-SM K1 token-range table (the K1 rule in +0.31 to 0.73 times the deterministic program's on every seed (q, data) (5 seeds on a 212-SM R200 part at 960, 961, 1024, +2048, 4096, 8192, 16172, 16231, 16384, 32768 and 40961 tokens, even and odd block counts), where the one-chain program's +is 2.0 to 2.3 times. The ranges follow the rule and the construction of the 212-SM K1 token-range table (the K1 rule in every run, the backward pass and the checkpointed step no more than 1.3 % slower than the deterministic program by the ratio of the medians over the runs and the median of the paired ratios, eager launches and graph replays). 212 SMs: measured in two jobs on four parts of one node each (one launch layout and row order per part), at 960, 4096, 16172, 16384, 32768 and 40961 tokens and at 960, 1024, 2048, 4096, 4097, 8192, 16172, 16384, 32768, 32769, 40960 and 40961 tokens; `4096 -> 1024` left the checkpointed step 1.7 % slower at 16384 tokens in one of them. 216 SMs: measured at 960, 961, 1024, 2048, 4096, 8192, 16172, 16231, 16384, 32768 and 40961 tokens on one part (two row orders); up to 960 -tokens these rows keep their one-chain program, whose error equals the deterministic program's there. Token counts -outside the ranges keep the routing of the tables above.* +tokens these rows keep their one-chain program, which gives the deterministic program's bits there. Token counts +outside the ranges keep the routing of the tables above (216 SMs: the one-chain program up to 960 tokens, the +deterministic program from 961 tokens up).* `"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 diff --git a/flashinfer/experimental/dense_projection_gemm/cake_backend.py b/flashinfer/experimental/dense_projection_gemm/cake_backend.py index c2b4b1ead81..9217a50abcb 100644 --- a/flashinfer/experimental/dense_projection_gemm/cake_backend.py +++ b/flashinfer/experimental/dense_projection_gemm/cake_backend.py @@ -290,16 +290,22 @@ def _ceil_div(a: int, b: int) -> int: # the inner (unit-stride) extent of the stored output and ``N`` its row count. A weight gradient # ``dW[d_out, d_in] = G[T, d_out].T @ X[T, d_in]`` has ``A = X``, ``B = G``, ``M = d_in``, ``N = d_out``, ``K = T``. -#: Per-row verdicts of the K1 weight gradient, ``(arch, L, M, N, out_dtype) -> (block_n, num_splits)`` (this -#: package's program with that tile width and split count; a width of :data:`K1_ROW_MAJOR_BLOCK_N` runs the row-major -#: program with the operands exchanged) or ``-> None`` (the deterministic program); ``M`` / ``N`` as above, -#: ``out_dtype`` ``"bf16"`` or ``"f32"``. A row carries a program only where that program was measured -#: faster than the deterministic package's shipped plan of the row in every measurement run: a one-sided paired sign -#: test (p < 0.05) with a median time ratio below 0.99, in eager launches and in CUDA graph replays, at token counts -#: 16172 to 16384, and no more than 3 % slower than that plan next to a kernel holding 8, 16 or 32 SMs on another -#: stream. A ``None`` entry is a measured row that missed; a row without an entry was not measured. Both run the -#: deterministic program. A verdict does not hold at the token counts of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` and, -#: where the token count spans an odd number of K blocks, of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS`. +#: Per-row verdicts of the K1 weight gradient, ``(arch, L, M, N, out_dtype) -> (block_n, num_splits)`` (this package's +#: program with that tile width and split count; a width of :data:`K1_ROW_MAJOR_BLOCK_N` runs the row-major program with +#: the operands exchanged) or ``-> None`` (the deterministic program); ``M`` / ``N`` as above, ``out_dtype`` ``"bf16"`` +#: or ``"f32"``. A row carries a program only where that program was measured faster than the deterministic package's +#: shipped plan of the row in every measurement run: a one-sided paired sign test (p < 0.05) with a median time ratio +#: below 0.99, in eager launches and in CUDA graph replays, at token counts 16172 to 16384, and no more than 3 % slower +#: than that plan next to a kernel holding 8, 16 or 32 SMs on another stream. A ``None`` entry is a measured row that +#: missed; a row without an entry was not measured. Both run the deterministic program. A verdict does not hold at the +#: token counts of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` and, where the token count spans an odd number of K blocks, of +#: :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS`. At every token count a verdict holds, its output met the deterministic +#: package's error budgets against an FP64 reference of the exact operands (BF16 output ``atol = rtol = 1e-2`` +#: elementwise with no element outside; FP32 output a relative Frobenius error ``<= 5e-5`` and a max error +#: ``<= 1e-3 * max|ref|``) on every measured seed, the bound the deterministic package states for its own programs; the +#: split program's FP64 error relative to the deterministic program's is reported data, not a rule. Where both programs +#: miss that budget for some seed (BF16 output at the largest token counts, FP32 output from about 98305 tokens), the +#: budget does not choose between them, and the tables keep the routing they give there. # fmt: off K1_SPLIT_RULES:dict[tuple[str, int, int, int, str], Optional[tuple[int, int]]] = { # sm_100a: two splits of 64-wide tiles (one wave of the CTA pairs; two partials, bitwise reproducible). d_out 128 @@ -360,8 +366,8 @@ def _ceil_div(a: int, b: int) -> int: ("sm_103a", 1, 4096, 48, "f32"): 152, } #: Token counts at which a row of :data:`K1_SPLIT_RULES` (or its entry in :data:`K1_SPLIT_OVERRIDES_BY_SM`) or of a -#: table of :data:`K1_SPLIT_RULES_BY_SM` does not take its verdict, ``key -> ((sm_count, min_tokens, max_tokens), -#: ...)``: a token count (``K`` of the weight gradient) from +#: table of :data:`K1_SPLIT_RULES_BY_SM` does not take its verdict, +#: ``key -> ((sm_count, min_tokens, max_tokens), ...)``: a token count (``K`` of the weight gradient) from #: ``min_tokens`` to ``max_tokens`` (``None``: no upper bound) runs the deterministic program on a part of at least #: ``sm_count`` SMs where ``sm_count`` is at or above the row's SM-count floor (:data:`K1_MIN_SM_COUNT`, #: :data:`K1_RULE_MIN_SM_COUNT`), and on a part of exactly ``sm_count`` SMs below it (the parts of @@ -369,54 +375,68 @@ def _ceil_div(a: int, b: int) -> int: #: at 16172 to 16384 tokens. The sm_100a / sm_103a entries come from a re-measurement at 512, 1024, 2048, 4096, 8192, #: 16231, 32768 and 65536 tokens on two 152-SM parts per row and token count (GB200, GB300; the strict runs of 8 GB300 #: points ran on two other parts, four GB300 parts in all; the protocol is in the K1 token table of the README). A -#: measured token count is an entry where the split program failed in at least one run (the rule above; the backward -#: pass or the checkpointed step more than 1.3 % slower than with the deterministic program; or the deterministic -#: package's FP64 error budget missed with an error above the deterministic program's) and the deterministic program -#: leaves the backward pass and the checkpointed step no more than 1.3 % slower, or where the split program's FP64 error -#: (relative Frobenius error or BF16 error-budget ratio) exceeded the deterministic program's by more than 2 % for some -#: seed on both parts. A token count between two consecutive measured entries is an entry too (not measured), and so is -#: every token count above 65536 of a row whose split program failed that FP64 comparison at 65536 tokens, and of -#: sm_103a ``4096 -> 32`` and ``4096 -> 48`` (BF16 and FP32): above 65536 tokens their two-split program's output is the -#: deterministic program's bits at an even number of 64-token K blocks and fails that comparison at an odd number (the -#: reason of their 148-SM entries in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`; the deterministic plan of these rows -#: is the same on 148 and 152 SMs; on a 148-SM part, 10 seeds: an FP32 relative Frobenius error 1.021 to 1.024 times -#: the deterministic program's at 65537 tokens; where the deterministic program misses the BF16 budget, a budget ratio -#: above 1.30 times the deterministic program's for some seed: ``4096 -> 32`` 1.305 at 65600 tokens, ``4096 -> 48`` up -#: to 1.379 at 65537 tokens). The open-ended sm_100a BF16 entries come from that rule where the deterministic program -#: misses the BF16 budget for some seed (the rule of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`; 10 seeds per token -#: count on a 148-SM B200 part, and the deterministic plan of these rows is the same on 148 and 152 SMs): a budget ratio -#: above 1.30 times the deterministic program's for some seed from the first token count listed -- ``4096 -> 48`` and -#: ``5120 -> 48`` 1.473 and 1.320 at 49153 tokens, ``5120 -> 32`` 1.328 at 65537 tokens, ``4096 -> 32`` 1.305 at -#: 65600 tokens (up to 1.727 at 106497; the entry starts at 65536 as on sm_103a) -- and ``4096 -> 64`` a median ratio -#: of 1.041 over its odd token counts from 65537 tokens (152 SMs; the 148-SM entry is in -#: :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`). Unmeasured token counts keep the verdict: below 512, above 65536 -#: (except the rows with open-ended entries), between two measured token counts with different outcomes, and on parts -#: with fewer than 152 SMs (148-SM parts: their own entries in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`). The -#: sm_107a entries are single token counts of the deterministic package's sm_107a rules for one token count (16231 and -#: 16384 tokens; 16231 only for ``4096 -> 1024``), which give the row a faster deterministic plan there, and 16231 -#: tokens of ``5120 -> 1024`` BF16 on 216 SMs (measured in the same runs; the deterministic plan there is the row's -#: plan at every token count). The other single 216-SM -#: token counts (512 and 65536) come from a re-measurement of chosen points at 512, 8192, 32768 and 65536 tokens -#: on two 216-SM parts per point under the rule of the sm_100a / sm_103a entries (the protocol is in the sm_107a K1 -#: token table of the README); their neighbouring token counts were not measured. On 208, 212 and 216 SMs the window -#: entries come from a re-measurement at 16172, 16231 and 16384 tokens against those plans (the protocol is in the -#: sm_107a K1 token table of the README): a measured token count is an entry where the row's program failed the 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). The open-ended sm_107a entries, and the 208-SM and 212-SM entries from 16172 to 16384 tokens, of the FP32 -#: rows ``4096 -> 512``, ``4096 -> 1024``, ``5120 -> 512`` and ``5120 -> 1024`` come from neither rule: their one-split -#: program adds the whole token reduction of an output in one chain, while from 961 tokens (16 K blocks) up the -#: deterministic program of these rows cuts it into units of ``ceil(blocks / 2)`` K blocks (the two halves of each -#: output at an even block count) and adds the partials in a fixed order (its automatic two-way split). From 961 -#: tokens up the one-split program's FP64 relative Frobenius error is 2.0 to 2.4 times the deterministic program's on -#: every seed, above the 2 % the rule allows; up to 960 tokens the deterministic program adds one chain as well and the -#: two errors are equal. Measured with the 216-SM plan on 212-SM R200 parts (the same program and accumulation order as -#: on 216 SMs, and the deterministic plan of these rows is the same on 208, 212 and 216 SMs; the error does not depend -#: on the SM count) against an FP64 reference: each row on two parts, 5 seeds, at 512 to 65536 tokens, and in a later -#: run with 10 seeds at 704 to 1024 tokens (the deterministic program's bits up to 960 tokens, 2.09 to 2.11 times its -#: error at 961 and 1024 tokens). So these rows run the deterministic program from 961 tokens up on parts of at least -#: 216 SMs, and from 16172 to 16384 tokens on 208 and 212 SMs (the 212-SM two-split ranges of -#: :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` lie outside those tokens). +#: measured token count is an entry where the split program failed in at least one run (the rule above, or the backward +#: pass or the checkpointed step more than 1.3 % slower than with the deterministic program) and the deterministic +#: program leaves the backward pass and the checkpointed step no more than 1.3 % slower, or where the split program's +#: output missed the deterministic package's error budgets against an FP64 reference of the exact operands (BF16 output +#: ``atol = rtol = 1e-2`` elementwise with no element outside; FP32 output a relative Frobenius error ``<= 5e-5`` and a +#: max error ``<= 1e-3 * max|ref|``) for some seed while the deterministic program's met it. A token count between two +#: consecutive measured entries is an entry too (not measured). The split program's FP64 error relative to the +#: deterministic program's is reported, not a rule. Where the deterministic program misses the budget for some seed +#: too, the budget does not choose a program, and these open-ended entries keep the deterministic program: sm_103a +#: ``4096 -> 32`` and ``4096 -> 48`` above 65536 tokens (BF16: both programs outside the budget on every seed from 65537 +#: tokens; FP32: the two-split program's output is the deterministic program's bits at an even number of 64-token K +#: blocks, and at an odd number it met the budget at 65537 tokens and missed it on every seed from 98305 tokens, as the +#: deterministic program's did; the deterministic plan of these rows is the same on 148 and 152 SMs, and their entries +#: start at the token counts at which the split program failed the timing rule above), and the open-ended sm_100a BF16 +#: entries of ``4096 -> 48`` and ``5120 -> 48`` from 49153 tokens, ``4096 -> 32`` from 65536 tokens and ``4096 -> 64`` +#: from 65536 tokens on 152 SMs (the 148-SM entry is in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`), where the split +#: program's output missed the BF16 budget for some seed and the deterministic program's did too (10 seeds per token +#: count on a 148-SM B200 part; the deterministic plan of these rows is the same on 148 and 152 SMs). sm_100a +#: ``5120 -> 32`` BF16 keeps its entry from 65537 tokens: measured with 10 seeds (20 runs each) at 32769 to 65025 tokens +#: on a 148-SM B200 part and two 152-SM GB200 parts (finer from 45569 to 49665 tokens), both programs met the BF16 +#: budget on every seed up to 46593 tokens, and from 47105 tokens each missed it for some seed at some token counts, +#: neither more accurate at every one; the budget does not choose a program there. The sm_103a BF16 entries of +#: ``5120 -> 32`` and ``5120 -> 64`` start at 45058 tokens: their split program's output met the BF16 budget on every +#: seed at 32768 tokens (30 seeds on two parts) and at every measured token count up to 45057 tokens on a 152-SM GB300 +#: part (10 seeds, 20 runs each), and missed it for some seed at the next measured token count, 47105, where the +#: deterministic program's output missed it too; the row keeps the deterministic program from there up (the split +#: program misses it on every seed at 65536 tokens, 30 seeds on two parts). That of ``5120 -> 128`` starts at 32768 +#: tokens: its split program met the budget up to 45057 tokens as well, but in a strict timing re-measurement on that +#: part it failed the rule above at 5 of 11 token counts from 32768 to 45057 (32768, 33793, 38913, 43009 and 44033) and +#: kept it at the other 6, so no measured window holds. That of ``5120 -> 48`` starts at 40961 tokens: on the 152-SM +#: part its split program's output met the budget on every seed up to 49152 tokens (the deterministic program's missed +#: it for one seed at 45057 and at 49152 tokens), but on a 148-SM B300 part, which runs the same program, it missed it +#: for one of 10 seeds at 40961 and at 40962 tokens while the deterministic program's met it on every seed; from 49153 +#: tokens both miss it on the 152-SM part. Unmeasured token counts keep the verdict: below 512, above 65536 (except the +#: rows with open-ended entries), between two measured token counts with different outcomes, and on parts with fewer +#: than 152 SMs (148-SM parts: their own entries in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM`). The sm_107a entries +#: are single token counts of the deterministic package's sm_107a rules for one token count (16231 and 16384 tokens; +#: 16231 only for ``4096 -> 1024``), which give the row a faster deterministic plan there, and 16231 tokens of +#: ``5120 -> 1024`` BF16 on 216 SMs (measured in the same runs; the deterministic plan there is the row's plan at every +#: token count). The other single 216-SM token counts (512 and 65536) come from a re-measurement of chosen points at +#: 512, 8192, 32768 and 65536 tokens on two 216-SM parts per point under the rule of the sm_100a / sm_103a entries (the +#: protocol is in the sm_107a K1 token table of the README); their neighbouring token counts were not measured. On 208, +#: 212 and 216 SMs the window entries come from a re-measurement at 16172, 16231 and 16384 tokens against those plans +#: (the protocol is in the sm_107a K1 token table of the README): a measured token count is an entry where the row's +#: program failed the 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). The 208-SM entry of ``4096 -> 1024`` FP32 at 16384 tokens comes from a later strict +#: re-measurement of the restored window on a 208-SM R200 part: the deterministic program's plan there (160-wide tiles) +#: was faster, and the row's program was 14.5 % slower at the operation and left the checkpointed step 6.3 % slower in +#: graph replays. The open-ended sm_107a entries come from neither rule: on parts of at least 216 SMs the FP32 rows +#: ``4096 -> 512``, ``4096 -> 1024``, ``5120 -> 512`` and ``5120 -> 1024`` run the deterministic program from 961 tokens +#: (16 K blocks) up, outside the ranges of :data:`K1_CHAIN_RANGES_BY_SM` (which come first). From there the +#: deterministic program of these rows cuts the token reduction into units of ``ceil(blocks / 2)`` K blocks and adds the +#: partials in a fixed order (its automatic two-way split), while their one-split program adds the whole token reduction +#: of an output in one chain; up to 960 tokens the two give the same bits. Measured with the 216-SM plan on 212-SM R200 +#: parts (the same program and accumulation order as on 216 SMs; the error does not depend on the SM count) against an +#: FP64 reference, each row on two parts, 5 seeds, at 512 to 65536 tokens, the one-split program's output met the +#: deterministic package's FP32 error budget (5e-5) on every seed up to 40960 tokens (1.86e-5 at 16384 tokens, 3.77e-5 +#: at 32768; 2.0 to 2.4 times the deterministic program's error from 961 tokens: reported) and missed it on every seed +#: from 45056 tokens up (7.6e-5 at 65536 tokens, where the deterministic program's is 3.8e-5). Its time from 961 to +#: 40960 tokens on 216-SM parts has not been measured under the rule above, so these entries are unchanged +#: (qualification on a 216-SM part pending). # fmt: off K1_SPLIT_RULE_EXCLUDED_TOKENS: dict[ tuple[str, int, int, int, str], tuple[tuple[int, int, Optional[int]], ...] @@ -437,11 +457,11 @@ def _ceil_div(a: int, b: int) -> int: ("sm_103a", 1, 4096, 32, "f32"): ((152, 512, 512), (152, 32768, None)), ("sm_103a", 1, 4096, 48, "bf16"): ((152, 512, 512), (152, 32768, None)), ("sm_103a", 1, 4096, 48, "f32"): ((152, 512, 512), (152, 32768, None)), - ("sm_103a", 1, 5120, 32, "bf16"): ((152, 512, 512), (152, 32768, None)), + ("sm_103a", 1, 5120, 32, "bf16"): ((152, 512, 512), (152, 45058, None)), ("sm_103a", 1, 5120, 32, "f32"): ((152, 512, 512),), - ("sm_103a", 1, 5120, 48, "bf16"): ((152, 512, 2048), (152, 32768, None)), + ("sm_103a", 1, 5120, 48, "bf16"): ((152, 512, 2048), (152, 40961, None)), ("sm_103a", 1, 5120, 48, "f32"): ((152, 512, 2048),), - ("sm_103a", 1, 5120, 64, "bf16"): ((152, 512, 2048), (152, 32768, None)), + ("sm_103a", 1, 5120, 64, "bf16"): ((152, 512, 2048), (152, 45058, None)), ("sm_103a", 1, 5120, 64, "f32"): ((152, 512, 2048),), ("sm_103a", 1, 5120, 128, "bf16"): ((152, 512, 8192), (152, 32768, None)), ("sm_103a", 1, 5120, 128, "f32"): ((152, 512, 4096),), @@ -459,8 +479,12 @@ def _ceil_div(a: int, b: int) -> int: (208, 16231, 16231), (208, 16384, 16384), (216, 512, 512), (216, 16231, 16231), (216, 16384, 16384), ), - ("sm_107a", 1, 4096, 512, "f32"): ((208, 16172, 16384), (212, 16172, 16384), (216, 961, None)), - ("sm_107a", 1, 4096, 1024, "f32"): ((208, 16172, 16384), (212, 16172, 16384), (216, 961, None)), + ("sm_107a", 1, 4096, 512, "f32"): ( + (208, 16231, 16231), (208, 16384, 16384), (212, 16231, 16231), (212, 16384, 16384), (216, 961, None), + ), + ("sm_107a", 1, 4096, 1024, "f32"): ( + (208, 16231, 16231), (208, 16384, 16384), (212, 16231, 16231), (216, 961, None), + ), ("sm_107a", 1, 5120, 32, "bf16"): ((216, 512, 512),), ("sm_107a", 1, 5120, 32, "f32"): ((216, 512, 512),), ("sm_107a", 1, 5120, 48, "f32"): ( @@ -471,9 +495,9 @@ def _ceil_div(a: int, b: int) -> int: (208, 16231, 16231), (212, 16231, 16231), (216, 512, 512), (216, 16231, 16231), (216, 16384, 16384), ), - ("sm_107a", 1, 5120, 512, "f32"): ((208, 16172, 16384), (212, 16172, 16384), (216, 961, None)), + ("sm_107a", 1, 5120, 512, "f32"): ((212, 16231, 16231), (212, 16384, 16384), (216, 961, None)), ("sm_107a", 1, 5120, 1024, "bf16"): ((216, 16231, 16231),), - ("sm_107a", 1, 5120, 1024, "f32"): ((208, 16172, 16384), (212, 16172, 16384), (216, 961, None)), + ("sm_107a", 1, 5120, 1024, "f32"): ((216, 961, None),), ("sm_107a", 1, 6144, 32, "bf16"): ((216, 512, 512),), ("sm_107a", 1, 8192, 128, "bf16"): ((216, 512, 512), (216, 65536, 65536)), ("sm_107a", 1, 8192, 128, "f32"): ((216, 512, 512), (216, 65536, 65536)), @@ -482,29 +506,30 @@ def _ceil_div(a: int, b: int) -> int: #: Token counts at which a row of :data:`K1_SPLIT_RULES` runs the deterministic program on the parts of exactly one SM #: count, ``(arch, sm_count) -> {key: ((min_tokens, max_tokens), ...)}`` (``max_tokens`` ``None``: no upper bound). An #: entry of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` at or above a row's SM-count floor holds from its SM count up; these -#: hold on exactly their SM count, and every other part keeps its routing. The sm_100a / sm_103a 148-SM entries (B200 -#: / B300 parts) follow the rule of the 152-SM entries of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`, measured on 148-SM -#: parts at 512, 1024, 2048, 4096, 8192, 16231, 32768 and 65536 tokens (the protocol is in the 148-SM K1 token table of -#: the README): every point in a screen on one part, the points that failed there again on a second part (B300: every -#: point on a part of a second node too), and the FP64 error of 5 seeds on two parts at 32768 and 65536 tokens and at -#: the first and the last token count of every entry and one token either side of it (and at 98304 and 131072 tokens -#: for the entries that reach 65536 tokens). A measured token count is an entry where the split program failed in some -#: run and the deterministic program leaves the backward pass and the checkpointed step no more than 1.3 % slower on -#: both parts (the worst of the ratio of the medians and the median of the paired ratios, eager launches and graph -#: replays), or where the split program failed the FP64 comparison: for some seed, a relative Frobenius error more than -#: 2 % above the deterministic program's, or, where the deterministic program meets the deterministic package's error -#: budget on every seed, an output outside it. Where the deterministic program misses the budget for some seed (BF16 -#: output at the largest token counts), the BF16 error-budget ratio, a maximum over the elements, moves by up to about -#: 25 % either way between the two programs at the same relative Frobenius error; there the split program fails where -#: the median over a row's measured token counts and seeds (10 seeds per token count) of its ratio to the deterministic -#: program's exceeds 1.02, or where one seed's ratio exceeds 1.30. A token count between two consecutive measured -#: entries is an entry too. Entries without an upper bound: sm_103a ``4096 -> 32`` and sm_100a ``4096 -> 64`` (BF16 -#: and FP32) from 65536 tokens up: above 65536 tokens the split program's output differs from the deterministic -#: program's where the token count spans an odd number of 64-token blocks (65537, 65538, 65600 and 65665 tokens: an FP32 -#: relative Frobenius error 2.1 to 2.3 % above on every seed, a BF16 budget ratio 1.305 times the deterministic -#: program's for one seed; sm_100a BF16 at 12 odd counts from 65600 to 98240 tokens, 10 seeds each: up to 1.65 times) -#: and is the deterministic program's bits where it spans an even number (65601 to 131072 tokens), where it does not -#: make the backward pass or the checkpointed step 1 % faster (sm_100a: the backward pass up to 3.5 % slower). +#: hold on exactly their SM count, and every other part keeps its routing. The sm_100a / sm_103a 148-SM entries (B200 / +#: B300 parts) follow the rule of the 152-SM entries of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`, measured on 148-SM parts +#: at 512, 1024, 2048, 4096, 8192, 16231, 32768 and 65536 tokens (the protocol is in the 148-SM K1 token table of the +#: README): every point in a screen on one part, the points that failed there again on a second part (B300: every point +#: on a part of a second node too), and the FP64 error of 5 seeds on two parts at 32768 and 65536 tokens and at the +#: first and the last token count of every entry and one token either side of it (and at 98304 and 131072 tokens for the +#: entries that reach 65536 tokens). A measured token count is an entry where the split program failed in some run and +#: the deterministic program leaves the backward pass and the checkpointed step no more than 1.3 % slower on both parts +#: (the worst of the ratio of the medians and the median of the paired ratios, eager launches and graph replays), or +#: where the split program's output missed the deterministic package's error budget against FP64 for some seed while the +#: deterministic program's met it. Where both miss it for some seed (BF16 output at the largest token counts), the +#: budget does not choose a program and the entry keeps the deterministic program. A token count between two +#: consecutive measured entries is an entry too. Entries without an upper bound: sm_103a ``4096 -> 32`` and sm_100a +#: ``4096 -> 64`` (BF16 and FP32) from 65536 tokens up: there the split program's output is the deterministic program's +#: bits where the token count spans an even number of 64-token blocks (65601 to 131072 tokens), where it does not make +#: the backward pass or the checkpointed step 1 % faster (sm_100a: the backward pass up to 3.5 % slower); where it spans +#: an odd number (65537, 65538, 65600 and 65665 tokens) its BF16 output and the deterministic program's both miss the +#: budget for some seed, and its FP32 output (a relative Frobenius error 2.1 to 2.3 % above the deterministic program's: +#: reported) misses the FP32 budget on every seed from 98305 tokens, as the deterministic program's does. sm_103a +#: ``5120 -> 48`` BF16 runs the deterministic program at 40961 to 49151 tokens and ``5120 -> 64`` BF16 at 47105 to 49152 +#: tokens on 148 SMs: on a 148-SM B300 part (10 seeds, 20 runs each, at 36865 to 61441 tokens) their split program's +#: output missed the BF16 budget for one of 10 seeds at 40961 and at 40962 tokens (``5120 -> 64``: 47105) where the +#: deterministic program's met it on every seed (``5120 -> 48`` met it at 41985 to 47105 tokens inside its range), and +#: from 49152 tokens (``5120 -> 64``: 49153) both programs missed it for some seed, so the split program stays there. #: Unmeasured token counts keep the verdict. # fmt: off K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM: dict[ @@ -524,9 +549,9 @@ def _ceil_div(a: int, b: int) -> int: ("sm_103a", 148): { ("sm_103a", 1, 4096, 32, "bf16"): ((65536, None),), ("sm_103a", 1, 4096, 32, "f32"): ((65536, None),), - ("sm_103a", 1, 5120, 48, "bf16"): ((512, 2048),), + ("sm_103a", 1, 5120, 48, "bf16"): ((512, 2048), (40961, 49151)), ("sm_103a", 1, 5120, 48, "f32"): ((512, 2048),), - ("sm_103a", 1, 5120, 64, "bf16"): ((512, 1024),), + ("sm_103a", 1, 5120, 64, "bf16"): ((512, 1024), (47105, 49152)), ("sm_103a", 1, 5120, 64, "f32"): ((512, 1024),), ("sm_103a", 1, 5120, 128, "bf16"): ((512, 4096), (32768, 65536)), ("sm_103a", 1, 5120, 128, "f32"): ((512, 4096), (32768, 65536)), @@ -539,38 +564,43 @@ def _ceil_div(a: int, b: int) -> int: #: deterministic program where the token count spans an odd number of 64-element K blocks (``ceil(K / 64)`` odd), #: ``key -> ((sm_count, min_tokens, max_tokens), ...)`` (parts and ``max_tokens`` as in #: :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`). A token count with an even block count, and every token count outside the -#: ranges, keeps the routing of the tables above. In these ranges the deterministic program of the row cuts the token -#: reduction of its tiles, taken in order, into units of ``ceil(blocks / 2)`` K blocks (at an odd block count a unit -#: runs on into the next tile) and adds the partials in a fixed order: from 961 tokens (16 K blocks) up for -#: ``4096 -> 32``, ``4096 -> 48`` and ``4096 -> 64``; for sm_100a ``5120 -> 32`` and ``5120 -> 48`` from where one -#: wave's operands outgrow the L2 (from 12289 tokens with 126 MiB of L2, between 12292 and 12337 tokens on a B200; -#: below, whole tiles); for sm_107a ``8192 -> 128`` while they fit the L2 (to 7808 tokens with 126 MiB of L2, to -#: between 4097 and 8192 tokens on a 212-SM R200; above, a synchronised two-way split that the two-split program -#: matches). At an even block count the two-split program's output is the deterministic program's bit for bit; at an -#: odd block count it is not, and its FP32 output's FP64 relative Frobenius error was more than the 2 % of the rule -#: above over the deterministic program's on every seed at every measured odd token count in the ranges from 1025 to -#: 65537 tokens (10 seeds each): sm_103a ``4096 -> 32`` and ``4096 -> 48`` and sm_100a ``4096 -> 32``, ``4096 -> 48`` -#: and ``4096 -> 64`` 1.03 to 1.15 times at 1025 to 49153 tokens and 1.021 to 1.024 times at 65537 tokens (148-SM B300 -#: and B200 parts; the deterministic plan of these rows is the same on 148 and 152 SMs); sm_100a ``5120 -> 32`` and -#: ``5120 -> 48`` 1.025 to 1.075 times at 12337 to 32769 tokens (B200); sm_107a ``8192 -> 128`` 1.04 to 1.14 times at -#: 1025 to 4097 tokens (212-SM R200, the 212-SM and the 216-SM plan; the deterministic plan of the row is the same on -#: 212 and 216 SMs). The excess falls with the token count (``5120 -> 32`` and ``5120 -> 48``: 1.7 to 1.96 % at 49153 -#: tokens, 0.7 to 1.5 % at 65537 to 131073 tokens; the ``4096`` rows: 1.2 to 1.7 % at 98305 and 131073 tokens); their -#: ranges stay open because the two-split program is not more accurate than the deterministic program at any odd block -#: count in them. The BF16 outputs of these rows had the deterministic program's relative Frobenius error and met the -#: error budget on every seed where the deterministic program met it. +#: ranges, keeps the routing of the tables above. In these rows the deterministic program cuts the token reduction of +#: its tiles, taken in order, into units of ``ceil(blocks / 2)`` K blocks (at an odd block count a unit runs on into the +#: next tile) and adds the partials in a fixed order: from 961 tokens (16 K blocks) up for ``4096 -> 32``, +#: ``4096 -> 48`` and ``4096 -> 64``; for sm_100a ``5120 -> 32`` and ``5120 -> 48`` from where one wave's operands +#: outgrow the L2 (from 12289 tokens with 126 MiB of L2, between 12292 and 12337 tokens on a B200; below, whole tiles). +#: At an even block count the two-split program's output is the deterministic program's bit for bit; at an odd block +#: count it is not, and its FP32 output met the deterministic package's FP32 error budget on every seed at every +#: measured odd token count from 1025 to 65537 tokens (10 seeds each; a relative Frobenius error 1.03 to 1.15 times the +#: deterministic program's at 1025 to 49153 tokens and 1.021 to 1.024 times at 65537 for the ``4096`` rows on 148-SM +#: B200 and B300 parts, whose deterministic plan is the same on 148 and 152 SMs, and 1.025 to 1.075 times at 12337 to +#: 32769 tokens for sm_100a ``5120 -> 32`` and ``5120 -> 48`` on a B200: reported, not a rule). From 98305 tokens up +#: neither program's FP32 output met the budget on every seed (98305 and 131073 tokens), so the ranges start above the +#: last measured odd token count at which the two-split program met it (65537; the token counts up to 98305 were not +#: measured). The BF16 outputs of these rows had the deterministic program's relative Frobenius error and met the error +#: budget on every seed where the deterministic program met it. The sm_103a ``4096 -> 48`` range from 14338 to 32767 +#: tokens comes from the time instead: on a 152-SM GB300 part, at odd block counts the two-split program left the +#: checkpointed step 2.0 to 2.3 % slower in graph replays at 16173 tokens (two measurements), the backward pass 1.6 and +#: 2.0 % slower in eager launches at 20481 and 28673 tokens, and failed the K1 rule in one of 8 groups at 18433 tokens, +#: while it kept the rule and the 1.3 % bound at 1025, 4097, 6145, 8193, 10241, 12289 and 14337 tokens. sm_107a +#: ``8192 -> 128`` FP32 (two splits on 212 and 216 SMs) has an entry from 216 SMs up only: at an odd block count from +#: 1025 to 8192 tokens, where its deterministic program's two-way split runs while one wave's operands fit the L2 (to +#: 7808 tokens with 126 MiB of L2, to between 4097 and 8192 tokens on a 212-SM R200), its output differs from the +#: deterministic program's and met the FP32 budget on every seed (10 seeds on 212-SM parts with the 212-SM and the +#: 216-SM plan; an error 1.04 to 1.14 times the deterministic program's at 1025 to 4097 tokens: reported), and its time +#: there was measured on 212-SM parts only. # fmt: off K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS: dict[ tuple[str, int, int, int, str], tuple[tuple[int, int, Optional[int]], ...] ] = { - ("sm_100a", 1, 4096, 32, "f32"): ((148, 1025, None),), - ("sm_100a", 1, 4096, 48, "f32"): ((148, 1025, None),), - ("sm_100a", 1, 4096, 64, "f32"): ((148, 1025, None),), - ("sm_100a", 1, 5120, 32, "f32"): ((148, 12289, None),), - ("sm_100a", 1, 5120, 48, "f32"): ((148, 12289, None),), - ("sm_103a", 1, 4096, 32, "f32"): ((148, 1025, None),), - ("sm_103a", 1, 4096, 48, "f32"): ((152, 1025, None),), - ("sm_107a", 1, 8192, 128, "f32"): ((212, 1025, 8192), (216, 1025, 8192)), + ("sm_100a", 1, 4096, 32, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 4096, 48, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 4096, 64, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 5120, 32, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 5120, 48, "f32"): ((148, 65538, None),), + ("sm_103a", 1, 4096, 32, "f32"): ((148, 65538, None),), + ("sm_103a", 1, 4096, 48, "f32"): ((152, 14338, 32767), (152, 65538, None)), + ("sm_107a", 1, 8192, 128, "f32"): ((216, 1025, 8192),), } # fmt: on #: Per-SM-count verdicts of the parts of an architecture with fewer SMs than its :data:`K1_MIN_SM_COUNT`, @@ -642,17 +672,16 @@ def _ceil_div(a: int, b: int) -> int: #: seeds on two parts. A measured token count joins a range where the row's program passed the K1 rule above in every #: run, left the backward pass and the checkpointed step no more than 1.3 % slower than the deterministic program (ratio #: of the medians and median of the paired ratios, eager launches and graph replays), and met the deterministic -#: package's FP64 error budget with an error (relative Frobenius error or BF16 error-budget ratio) no more than 2 % -#: above the deterministic program's for every seed on both parts. A range runs from such a measured token count to +#: package's FP64 error budget for every seed on both parts. A range runs from such a measured token count to #: the window's first token count, or from the window's last one to such a measured token count, and holds no measured #: token count that failed; the token counts between two consecutive measured ones inside it take the verdict too (not #: measured). The token counts below a range's first measured token count, above its last, and those between a #: measured token count that passed and one that failed run the deterministic program. On 212 SMs: the FP32 rows #: ``4096 -> 32, 48, 64`` and ``5120 -> 32, 48, 64`` from 1024 to 65536 tokens (at 512 not faster in every run), #: ``8192 -> 128`` from 1024 to 32768 tokens (at 512 slower, at 65536 not faster in every run), and -#: the FP32 rows ``4096 / 5120 -> 512 / 1024`` at 512 tokens only: from 1024 tokens up their one-split program, which -#: adds an output's whole token reduction in one chain, has about twice the deterministic program's FP64 error (their -#: two-split program runs at the ranges of :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`). +#: the FP32 rows ``4096 / 5120 -> 512 / 1024`` at 512 tokens only (from 1024 tokens up their one-split program was not +#: taken for its FP64 error, about twice the deterministic program's, which is reported data and not a rule; there +#: these rows run the programs of :data:`K1_CHAIN_RANGES_BY_SM` and :data:`K1_TWO_SPLIT_RANGES_BY_SM`). # fmt: off K1_SPLIT_RULES_BY_SM_TOKEN_RANGES: dict[ tuple[str, int], dict[tuple[str, int, int, int, str], tuple[tuple[int, int], ...]] @@ -695,42 +724,45 @@ def _ceil_div(a: int, b: int) -> int: }, } # fmt: on -#: Verdicts of rows on a part of exactly one SM count at further token ranges that hold only where the token count spans -#: an even number of 64-element K blocks (``ceil(K / 64)`` even), +#: Verdicts of rows on a part of exactly one SM count at further token ranges, #: ``(arch, sm_count) -> {key: ((min_tokens, max_tokens, (block_n, num_splits)), ...)}`` (keys and values as in #: :data:`K1_SPLIT_RULES`): on a part of exactly ``sm_count`` SMs, below the row's SM-count floor or at or above it, a -#: token count from ``min_tokens`` to ``max_tokens`` with an even block count runs that range's program, except the -#: token counts of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for the part and the token counts of the row's ranges in -#: :data:`K1_CHAIN_RANGES_BY_SM` for the part, which come first. A token count with an odd block count, and every -#: other token count, part and row, keeps the routing of the tables above. A range lies outside -#: :data:`K1_SPLIT_RULES_BY_SM_TOKENS` and outside the row's ranges in :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` and -#: :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM`. On 152-SM sm_100a and sm_103a parts the FP32 row ``4096 -> 128`` (no verdict -#: in :data:`K1_SPLIT_RULES`) runs two splits of 64-wide tiles, each adding half of an output's K blocks: at an even -#: block count its output was the deterministic program's bit for bit (5 seeds at 1024 to 8192 tokens), at an odd count -#: (1088 and 2112 tokens) it was not, with an FP64 relative Frobenius error 1.107 and 1.142 times the deterministic -#: program's. It passed the K1 rule (the README) at 1024, 2048 and 4096 tokens on two parts of each GPU (0.70 to 0.79 -#: times the deterministic program's time in CUDA graph replays), and at 6144 tokens on two sm_103a parts; on the -#: sm_100a part measured at 6144 tokens its time next to a kernel holding SMs on another stream was 1.099 times the -#: deterministic program's, and at 512 and 8192 tokens it failed on every part. 148-SM parts keep the deterministic -#: program. On 212-SM sm_107a parts the same row (no verdict in the part's table of :data:`K1_SPLIT_RULES_BY_SM` -#: either) runs that program from 1024 to 7680 tokens at an even block count: its output was the deterministic program's -#: bit for bit (5 seeds at 1024, 2048, 4096, 6144, 6656, 7168 and 7680 tokens), at an odd count (1088 to 5000 tokens) -#: its FP64 error was 1.11 to 1.16 times the deterministic program's; it passed the K1 rule there on two nodes of each -#: of two R200 systems (0.70 to 0.76 times the deterministic program's time), and failed at 512, 768 and 896 tokens; -#: from 8000 tokens the deterministic plan changes and the two-split error was 1.54 times its error, and at 12288 and -#: 16384 tokens it was 5 to 7 % slower. On 212 SMs the FP32 rows ``4096 / 5120 -> 512 / 1024`` run two splits of their -#: 212-SM tile width (128, 192 for ``4096 -> 1024``, 256 for ``5120 -> 1024``), each adding half of an output's K -#: blocks. From 1024 tokens up, at an even block count, their output was the deterministic program's bit for bit (5 -#: seeds on two parts at 1024, 2048, 4096, 8192, 16231, 16384, 32768 and 65536 tokens, on one part at 6000, 12000 and -#: 48000); at an odd count from 17 blocks up it was not, with an FP64 relative Frobenius error 1.01 to 1.17 times the -#: deterministic program's (5 seeds, 1025 to 65537 tokens). Measured as the ranges of -#: :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, at 512, 1024, 2048, 4096, 8192, 16231, 16384, 32768 and 65536 tokens (the -#: protocol and the rule are in the 212-SM K1 token-range table of the README): a range runs between measured token -#: counts at which the two-split program passed, and holds none that failed. At 512 tokens two splits are not faster -#: than the deterministic plan, ``4096 -> 1024`` is not faster in every run at 1024 tokens, and ``5120 -> 1024`` failed -#: the rule at every measured token count from 16231 up. +#: token count from ``min_tokens`` to ``max_tokens`` runs that range's program, at an even and at an odd number of +#: 64-element K blocks, except the token counts of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for the part and the token +#: counts of the row's ranges in :data:`K1_CHAIN_RANGES_BY_SM` for the part, which come first. Every other token count, +#: part and row keeps the routing of the tables above. A range lies outside :data:`K1_SPLIT_RULES_BY_SM_TOKENS` and +#: outside the row's ranges in :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` and :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM`. +#: Every program here is two splits, each adding half of an output's K blocks: at an even block count its output is the +#: deterministic program's bit for bit; at an odd count it is not, and its FP32 output met the deterministic package's +#: FP32 error budget on every measured seed (its FP64 error relative to the deterministic program's is reported, not a +#: rule). On 152-SM sm_100a and sm_103a parts the FP32 row ``4096 -> 128`` (no verdict in :data:`K1_SPLIT_RULES`) runs +#: two splits of 64-wide tiles: its output was the deterministic program's bit for bit at an even block count (5 seeds +#: at 1024 to 8192 tokens); at an odd count (1088 and 2112 tokens) it was not, with an FP64 relative Frobenius error +#: 1.107 and 1.142 times the deterministic program's. It passed the K1 rule (the README) at 1024, 2048 and 4096 tokens +#: on two parts of each GPU (0.70 to 0.79 times the deterministic program's time in CUDA graph replays), and at 6144 +#: tokens on two sm_103a parts; on the sm_100a part measured at 6144 tokens its time next to a kernel holding SMs on +#: another stream was 1.099 times the deterministic program's, and at 512 and 8192 tokens it failed on every part. +#: 148-SM parts keep the deterministic program. On 212-SM sm_107a parts the same row (no verdict in the part's table of +#: :data:`K1_SPLIT_RULES_BY_SM` either) runs that program from 1024 to 7680 tokens: its output was the deterministic +#: program's bit for bit at an even block count (5 seeds at 1024, 2048, 4096, 6144, 6656, 7168 and 7680 tokens), at an +#: odd count (1088 to 5000 tokens) its FP64 error was 1.11 to 1.16 times the deterministic program's; it passed the K1 +#: rule there on two nodes of each of two R200 systems (0.70 to 0.76 times the deterministic program's time; 0.48 to +#: 0.51 times at 1088 tokens), and failed at 512, 768 and 896 tokens; from 8000 tokens the deterministic plan changes +#: and the two-split error was 1.54 times its error, and at 12288 and 16384 tokens it was 5 to 7 % slower. On 212 SMs +#: the FP32 rows ``4096 / 5120 -> 512 / 1024`` run two splits of their 212-SM tile width (128, 192 for ``4096 -> 1024``, +#: 256 for ``5120 -> 1024``). From 1024 tokens up, at an even block count, their output was the deterministic program's +#: bit for bit (5 seeds on two parts at 1024, 2048, 4096, 8192, 16231, 16384, 32768 and 65536 tokens, on one part at +#: 6000, 12000 and 48000); at an odd count from 17 blocks up it was not, with an FP64 relative Frobenius error 1.01 to +#: 1.17 times the deterministic program's (5 seeds, 1025 to 65537 tokens; at most 3.8e-5 at 65537 tokens). Measured as +#: the ranges of :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, at 512, 1024, 2048, 4096, 8192, 16231, 16384, 32768 and +#: 65536 tokens (the protocol and the rule are in the 212-SM K1 token-range table of the README): a range runs between +#: measured token counts at which the two-split program passed, and holds none that failed; the token counts between two +#: measured ones inside it take its verdict at either block parity (not measured). At 512 tokens two splits are not +#: faster than the deterministic plan, ``4096 -> 1024`` is not faster in every run at 1024 tokens, and ``5120 -> 1024`` +#: failed the rule at every measured token count from 16231 up. The ranges of these four rows below 16172 tokens and +#: from 32768 to 40961 tokens run the three-chain program of :data:`K1_CHAIN_RANGES_BY_SM` instead. # fmt: off -K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM: dict[ +K1_TWO_SPLIT_RANGES_BY_SM: dict[ tuple[str, int], dict[tuple[str, int, int, int, str], tuple[tuple[int, int, tuple[int, int]], ...]], ] = { @@ -755,35 +787,37 @@ def _ceil_div(a: int, b: int) -> int: #: parts), at the token counts of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` only, except the token counts of #: :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for the part. Every other token count, every other SM count and every row #: without an entry keep the verdict of :data:`K1_SPLIT_RULES`. No part has an entry: the one measured verdict, the -#: 192-wide one-split program of ``4096 -> 1024`` FP32 on 216 SMs, fails the FP64 comparison of the sm_107a FP32 -#: one-split rows from 961 tokens up (:data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`), and the row runs the deterministic program -#: at these token counts. +#: 192-wide one-split program of ``4096 -> 1024`` FP32 on 216 SMs, lies inside the row's 216-SM entry of +#: :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` (from 961 tokens up), and the row runs the deterministic program at these token +#: counts. K1_SPLIT_OVERRIDES_BY_SM: dict[ tuple[str, int], dict[tuple[str, int, int, int, str], Optional[tuple[int, int]]] ] = {} #: Token ranges at which a row runs this package's one-split program with several accumulation chains on the parts of #: exactly one SM count, ``(arch, sm_count) -> {key: ((min_tokens, max_tokens, (block_n, chains)), ...)}`` (keys as in #: :data:`K1_SPLIT_RULES`): a token count from ``min_tokens`` to ``max_tokens`` runs the program of that tile width at -#: one split with that many chains (:func:`k1_chain_plan`) in place of the routing of the tables above, whose entries -#: for these rows keep their one-chain program off these token counts on accuracy, and in place of the two splits of -#: :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` at an even number of 64-element K blocks: at the even block counts of those -#: ranges both programs were timed at on 212 SMs (1024 to 8192 and 32768 to 40960 tokens), three chains ran the weight -#: gradient 7 to 52 % and the checkpointed step 2 to 21 % faster than two splits (medians, eager and graph), with an -#: FP64 error below the deterministic program's, which two splits match bit for bit. The FP32 rows -#: ``4096 / 5120 -> 512 / 1024``: from 961 tokens up their one-chain program adds an output's whole token reduction in -#: one chain, about twice the deterministic program's FP64 error; three chains of a third of the K blocks each had an -#: FP64 relative Frobenius error 0.31 to 0.73 times the deterministic program's on every seed (5 seeds on a 212-SM R200 -#: part at 960, 961, 1024, 2048, 4096, 8192, 16172, 16231, 16384, 32768 and 40961 tokens, even and odd block counts; the -#: error does not depend on the SM count) and ran within 5 % of the one-chain program. Measured as the ranges of -#: :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` (the rule, the bound on the backward pass and the checkpointed step, and -#: the construction of a range are that table's): on 212 SMs in two jobs on four parts of one node each, at 960, 4096, -#: 16172, 16384, 32768 and 40961 tokens and at 960, 1024, 2048, 4096, 4097, 8192, 16172, 16384, 32768, 32769, 40960 and -#: 40961 tokens; on 216 SMs at 960, 961, 1024, 2048, 4096, 8192, 16172, 16231, 16384, 32768 and 40961 tokens on one -#: part. Up to 960 tokens on 216 SMs the one-chain program keeps its :data:`K1_SPLIT_RULES` verdict (the deterministic -#: program adds one chain there too). ``4096 -> 512`` and ``5120 -> 512`` are not faster than the deterministic plan in -#: every run at 16384 tokens on 212 SMs, nor are ``4096 -> 512`` and ``4096 -> 1024`` at 16231 and 16384 tokens on 216 -#: SMs; ``4096 -> 1024`` at 16384 tokens on 212 SMs (in one of the two jobs) and ``5120 -> 1024`` at 32768 tokens on 216 -#: SMs left the checkpointed step more than 1.3 % slower than the deterministic program. +#: one split with that many chains (:func:`k1_chain_plan`) in place of the routing of the tables above (the one-chain +#: program of :data:`K1_SPLIT_RULES` and of the 212-SM table of :data:`K1_SPLIT_RULES_BY_SM`, and the two splits of +#: :data:`K1_TWO_SPLIT_RANGES_BY_SM`): at the even block counts of the two-split ranges both programs were timed at on +#: 212 SMs (1024 to 8192 and 32768 to 40960 tokens), three chains ran the weight gradient 7 to 52 % and the checkpointed +#: step 2 to 21 % faster than two splits (medians, eager and graph), with an FP64 error below the deterministic +#: program's, which two splits match bit for bit. The FP32 rows ``4096 / 5120 -> 512 / 1024``: from 961 tokens up their +#: one-chain program adds an output's whole token reduction in one chain, about twice the deterministic program's FP64 +#: error (reported); three chains of a third of the K blocks each had an FP64 relative Frobenius error 0.31 to 0.73 +#: times the deterministic program's on every seed (5 seeds on a 212-SM R200 part at 960, 961, 1024, 2048, 4096, 8192, +#: 16172, 16231, 16384, 32768 and 40961 tokens, even and odd block counts; the error does not depend on the SM count) +#: and ran within 5 % of the one-chain program. Measured as the ranges of :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` +#: (the rule, the bound on the backward pass and the checkpointed step, and the construction of a range are that +#: table's): on 212 SMs in two jobs on four parts of one node each, at 960, 4096, 16172, 16384, 32768 and 40961 tokens +#: and at 960, 1024, 2048, 4096, 4097, 8192, 16172, 16384, 32768, 32769, 40960 and 40961 tokens; on 216 SMs at 960, 961, +#: 1024, 2048, 4096, 8192, 16172, 16231, 16384, 32768 and 40961 tokens on one part. Up to 960 tokens on 216 SMs the +#: one-chain program keeps its :data:`K1_SPLIT_RULES` verdict (the deterministic program adds one chain there too). +#: ``4096 -> 512`` and ``5120 -> 512`` are not faster than the deterministic plan in every run at 16384 tokens on 212 +#: SMs, nor are ``4096 -> 512`` and ``4096 -> 1024`` at 16231 and 16384 tokens on 216 SMs; ``4096 -> 1024`` at 16384 +#: tokens on 212 SMs (in one of the two jobs) and ``5120 -> 1024`` at 32768 tokens on 216 SMs left the checkpointed step +#: more than 1.3 % slower than the deterministic program. Outside these ranges the rows keep the routing of the tables +#: above (on 216 SMs and up the one-chain program up to 960 tokens and the deterministic program from 961 tokens +#: up: :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`). # fmt: off K1_CHAIN_RANGES_BY_SM: dict[ tuple[str, int], @@ -849,20 +883,19 @@ def k1_plan( out_f32: bool, sm_count: int, ) -> Optional[tuple[int, int]]: - """``(block_n, num_splits)`` of this package's K1 program for the transposed-store launch ``out[l, n, m]`` - (``M``, ``N`` as in the comment above; ``K`` the token count) on a device of ``sm_count`` SMs, or ``None``: the - row runs the deterministic program. The row's verdict in :data:`K1_SPLIT_RULES` from the SM-count floor of the - row (:data:`K1_MIN_SM_COUNT`, :data:`K1_RULE_MIN_SM_COUNT`) up, replaced at a ``K`` of + """``(block_n, num_splits)`` of this package's K1 program for the transposed-store launch ``out[l, n, m]`` (``M``, + ``N`` as in the comment above; ``K`` the token count) on a device of ``sm_count`` SMs, or ``None``: the row runs the + deterministic program. The row's verdict in :data:`K1_SPLIT_RULES` from the SM-count floor of the row + (:data:`K1_MIN_SM_COUNT`, :data:`K1_RULE_MIN_SM_COUNT`) up, replaced at a ``K`` of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` by the row's entry in the table of :data:`K1_SPLIT_OVERRIDES_BY_SM` for exactly - ``sm_count``, if any, and at a ``K`` of the row's range in :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` for exactly - ``sm_count`` that spans an even number of 64-element K blocks by that range's verdict; below the floor the row's - verdict in the table of :data:`K1_SPLIT_RULES_BY_SM` for exactly ``sm_count``, only at a ``K`` of - :data:`K1_SPLIT_RULES_BY_SM_TOKENS` or of the row's ranges in :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` for that - SM count, and the verdict of the row's range in :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` for that SM count at a - ``K`` of that range, else the verdict of the row's range in :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` for that SM - count at a ``K`` of that range that spans an even number of 64-element K blocks; either except at a ``K`` of the - row's entries in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for this part (from the entry's SM count up where that count - is at or above the floor, exactly that SM count below it), in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM` for + ``sm_count``, if any, and at a ``K`` of the row's range in :data:`K1_TWO_SPLIT_RANGES_BY_SM` for exactly + ``sm_count`` by that range's verdict; below the floor the row's verdict in the table of :data:`K1_SPLIT_RULES_BY_SM` + for exactly ``sm_count``, only at a ``K`` of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` or of the row's ranges in + :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` for that SM count, and the verdict of the row's range in + :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` for that SM count at a ``K`` of that range, else the verdict of the row's + range in :data:`K1_TWO_SPLIT_RANGES_BY_SM` for that SM count at a ``K`` of that range; either except at a ``K`` of + the row's entries in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for this part (from the entry's SM count up where that + count is at or above the floor, exactly that SM count below it), in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM` for exactly ``sm_count``, or, at a ``K`` that spans an odd number of 64-element K blocks, in :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS` for this part (as the first table). A row without a verdict, a ``K`` of those entries, every other ``K`` below the floor and a ``K`` too short for the split count (:func:`k1_splits_fit`) @@ -883,11 +916,10 @@ def k1_plan( if in_window: overrides = K1_SPLIT_OVERRIDES_BY_SM.get((arch, int(sm_count)), {}) choice = overrides.get(key, choice) - if _ceil_div(K, K1_BLOCK_K) % 2 == 0: - even_part = K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM.get((arch, int(sm_count)), {}) - own = [v for lo, hi, v in even_part.get(key, ()) if lo <= int(K) <= hi] - if own: - choice = own[0] + two_part = K1_TWO_SPLIT_RANGES_BY_SM.get((arch, int(sm_count)), {}) + own = [v for lo, hi, v in two_part.get(key, ()) if lo <= int(K) <= hi] + if own: + choice = own[0] else: part = (arch, int(sm_count)) ranges = K1_SPLIT_RULES_BY_SM_TOKEN_RANGES.get(part, {}).get(key, ()) @@ -896,9 +928,9 @@ def k1_plan( for lo, hi, v in K1_SPLIT_RANGE_OVERRIDES_BY_SM.get(part, {}).get(key, ()) if lo <= int(K) <= hi ] - if not own and _ceil_div(K, K1_BLOCK_K) % 2 == 0: - even = K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM.get(part, {}).get(key, ()) - own = [v for lo, hi, v in even if lo <= int(K) <= hi] + if not own: + two = K1_TWO_SPLIT_RANGES_BY_SM.get(part, {}).get(key, ()) + own = [v for lo, hi, v in two if lo <= int(K) <= hi] if own: choice = own[0] elif in_window or any(lo <= int(K) <= hi for lo, hi in ranges): diff --git a/tests/experimental/test_cake_dense_projection_gemm.py b/tests/experimental/test_cake_dense_projection_gemm.py index 5d914a67075..4965a49c639 100644 --- a/tests/experimental/test_cake_dense_projection_gemm.py +++ b/tests/experimental/test_cake_dense_projection_gemm.py @@ -683,27 +683,28 @@ def no_k1_chains(monkeypatch): @pytest.mark.parametrize( "arch, sm_count, rows", [ - ("sm_100a", 148, 23), - ("sm_103a", 152, 30), - ("sm_103a", 148, 27), + ("sm_100a", 148, 28), + ("sm_103a", 152, 31), + ("sm_103a", 148, 28), ("sm_107a", 216, 26), ("sm_107a", 212, 19), - ("sm_107a", 208, 8), + ("sm_107a", 208, 14), ("sm_107a", 210, 0), ], ) def test_k1_planners_take_the_measured_rows(arch, sm_count, rows): - """The verdicts of the three architectures (on sm_107a with the chain ranges of 216 SMs) take 23 / 30 / 26 of the + """The verdicts of the three architectures (on sm_107a with the chain ranges of 216 SMs) take 28 / 31 / 26 of the 124 weight-gradient rows (the MLA head rows, ``L = 64``, run the deterministic program), at 16172 and 16231 tokens; - on 148 SMs sm_103a leaves ``4096 -> 48`` to the deterministic program (27 rows). 16172 tokens span an odd number of - 64-token K blocks (253): sm_100a ``4096 -> 32``, ``4096 -> 48``, ``4096 -> 64``, ``5120 -> 32`` and ``5120 -> 48`` - FP32 and sm_103a ``4096 -> 32`` and ``4096 -> 48`` FP32 run the deterministic program there - (:data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS`). On sm_107a the token counts of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` - run the deterministic program (seven FP32 rows and one BF16 row at 16231 tokens on 216 SMs; the four one-split FP32 - rows at both token counts on 208 SMs, and on 212 and 216 SMs outside their ranges in :data:`K1_CHAIN_RANGES_BY_SM`: - 16231 tokens for ``4096 -> 512``, ``4096 -> 1024`` and ``5120 -> 512`` on 212 SMs and for ``4096 -> 512`` and - ``4096 -> 1024`` on 216 SMs, where the other five and six (row, token count) pairs run three accumulation chains). - The per-SM-count sm_107a tables and the chain ranges take 19 rows on 212 SMs, the tables 8 on 208 SMs; an SM count + on 148 SMs sm_103a leaves ``4096 -> 48`` to the deterministic program (28 rows). 16172 tokens span an odd number of + 64-token K blocks (253), below the first token count of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS` (65538), so the + FP32 rows of that table run their two-split program there, except sm_103a ``4096 -> 48`` on 152 SMs (its range from + 14338 to 32767 tokens). On sm_107a the token counts of + :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` run the deterministic program (seven FP32 rows and one BF16 row at 16231 + tokens on 216 SMs; on 212 and 216 SMs outside their ranges in :data:`K1_CHAIN_RANGES_BY_SM`: 16231 tokens for + ``4096 -> 512``, ``4096 -> 1024`` and ``5120 -> 512`` on 212 SMs and for ``4096 -> 512`` and ``4096 -> 1024`` on 216 + SMs, where the other five and six (row, token count) pairs run three accumulation chains; on 208 SMs ``4096 -> 512`` + and ``4096 -> 1024`` FP32 at 16231 tokens, where the other six pairs of the four one-split FP32 rows run one split). + The per-SM-count sm_107a tables and the chain ranges take 19 rows on 212 SMs, the tables 14 on 208 SMs; an SM count below 216 without a table takes none.""" taken = 0 for d_in, d_out in WGRAD_SHAPES: @@ -790,7 +791,7 @@ def test_k1_planner_choices(): cake_backend.k1_plan(arch="sm_100a", M=4096, N=64, sm_count=120, **kw) is None ) # d_out = 128 at d_in 4096 was measured and missed on every architecture (152 SMs: FP32 two-split ranges of - # K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM below the window, test_k1_152_sm_even_block_ranges) + # K1_TWO_SPLIT_RANGES_BY_SM below the window, test_k1_152_sm_two_split_ranges) for arch, sms in SM_COUNTS.items(): assert ( cake_backend.k1_plan(arch=arch, M=4096, N=128, sm_count=sms, **kw) is None @@ -848,11 +849,12 @@ def p107(M, N, sms, out_f32=False): w512 = dict(arch="sm_107a", M=4096, N=512, sm_count=216, L=1, out_f32=True) assert cake_backend.k1_plan(K=960, **w512) == (128, 1) assert cake_backend.k1_chain_plan(K=960, **w512) is None - for T in (961, 16172): + for T in (961, 16172, 32768, 40961): assert cake_backend.k1_plan(K=T, **w512) == (128, 1), T assert cake_backend.k1_chain_plan(K=T, **w512) == (128, 3), T - for T in (16231, 16384): + for T in (16173, 16231, 16384, 32767, 40962, 65536): assert cake_backend.k1_plan(K=T, **w512) is None, T + assert cake_backend.k1_chain_plan(K=T, **w512) is None, T # measured and missed on sm_107a: 4096 -> 64 BF16, 4096 -> 512 BF16, 4096 -> 4096 FP32 for N, out_f32 in ((64, False), (512, False), (4096, True)): key = ("sm_107a", 1, 4096, N, "f32" if out_f32 else "bf16") @@ -965,19 +967,20 @@ def test_k1_small_token_counts_run_the_deterministic_program(): def test_k1_excluded_token_counts(no_k1_chains): """A row of :data:`K1_SPLIT_RULES` or of a table of :data:`K1_SPLIT_RULES_BY_SM` runs the deterministic program at - the token counts of its :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` entries: from the entry's SM count up where that - count is at or above the row's SM-count floor, on exactly that SM count below it. Every entry belongs to a row - with a verdict on its part; per SM count its ranges are closed, disjoint and not adjacent, start at a measured - token count (512 to 65537; the sm_107a FP32 one-split ranges at 961, the first token count at which the - deterministic program of these rows splits K) and end at one or have no upper - bound; the row runs the deterministic program at both ends and keeps its verdict (inside the token window, the - override of :data:`K1_SPLIT_OVERRIDES_BY_SM` for the part's SM count, if any) one token outside a range (inside the - token window of a per-SM-count table), on a part with fewer SMs than an entry at or above the floor, on the floor's - part for an entry below it, and above 65536 tokens unless a range has no upper bound. A token count with an odd - number of 64-token blocks inside an entry of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS` for the part runs the - deterministic program too.""" + the token counts of its :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` entries: from the entry's SM count up where that count + is at or above the row's SM-count floor, on exactly that SM count below it. Every entry belongs to a row with a + verdict on its part; per SM count its ranges are closed, disjoint and not adjacent, start at a measured token count + or one token above one (512 to 65537; the sm_107a FP32 one-split ranges at 961, the first token count at which the + deterministic program of these rows splits K) and end at one or have no upper bound; the row runs the + deterministic program at both ends and keeps its verdict (inside the token window, the override of + :data:`K1_SPLIT_OVERRIDES_BY_SM` for the part's SM count, if any) one token outside a range (inside the token window + of a per-SM-count table), on a part with fewer SMs than an entry at or above the floor, on the floor's part for an + entry below it, and above 65536 tokens unless a range has no upper bound. A token count with an odd number of + 64-token blocks inside an entry of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS` for the part, or inside an entry of + :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM` for the part's exact SM count, runs the deterministic program too.""" table = cake_backend.K1_SPLIT_RULE_EXCLUDED_TOKENS odd_table = cake_backend.K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS + exact_table = cake_backend.K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM by_sm = cake_backend.K1_SPLIT_RULES_BY_SM overrides = cake_backend.K1_SPLIT_OVERRIDES_BY_SM routed = {k for k, v in cake_backend.K1_SPLIT_RULES.items() if v is not None} @@ -1007,7 +1010,8 @@ def cut(sm_count, T, key=key, floor=floor): on_part = sm_count >= sms if sms >= floor else sm_count == sms if on_part and a <= T and (b is None or b >= T): return True - return False + exact = exact_table.get((key[0], sm_count), {}).get(key, ()) + return any(a <= T and (b is None or b >= T) for a, b in exact) def kept(sm_count, T, want, key=key, floor=floor): override = overrides.get((key[0], sm_count), {}).get(key) @@ -1071,7 +1075,9 @@ def kept(sm_count, T, want, key=key, floor=floor): assert cake_backend.k1_plan(K=512, sm_count=152, **kw) is None assert cake_backend.k1_plan(K=511, sm_count=152, **kw) == (64, 2) assert cake_backend.k1_plan(K=512, sm_count=151, **kw) == (64, 2) - assert cake_backend.k1_plan(K=32768, sm_count=152, **kw) is None + for T in (32768, 45057): + assert cake_backend.k1_plan(K=T, sm_count=152, **kw) == (64, 2), T + assert cake_backend.k1_plan(K=45058, sm_count=152, **kw) is None assert cake_backend.k1_plan(K=131072, sm_count=152, **kw) is None assert cake_backend.k1_plan(K=131072, sm_count=149, **kw) == (64, 2) # a row without an entry keeps its verdict at every token count its split count fits @@ -1216,10 +1222,9 @@ def test_k1_212_sm_token_ranges(no_k1_chains): every token count of their ranges (both ends and the measured token counts inside), except a token count of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for the part, and run the deterministic program one token outside a range and at every measured token count outside the ranges and the window (except the two-split ranges of - :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` at an even number of K blocks: ``test_k1_212_sm_even_block_ranges``). - Every range starts and ends at a measured token count, lies inside 512 to 65536 tokens and belongs to a row with - a program on that part. Another SM count - (208, 210, 216, 217) and every other row keep their routing: 208 SMs and the other 212-SM rows run the + :data:`K1_TWO_SPLIT_RANGES_BY_SM`: ``test_k1_212_sm_two_split_ranges``). Every range starts and ends at a + measured token count, lies inside 512 to 65536 tokens and belongs to a row with a program on that part. Another + SM count (208, 210, 216, 217) and every other row keep their routing: 208 SMs and the other 212-SM rows run the deterministic program outside the window, 216 SMs and up decide by :data:`K1_SPLIT_RULES`.""" table = cake_backend.K1_SPLIT_RULES_BY_SM_TOKEN_RANGES assert set(table) == {("sm_107a", 212)} @@ -1265,9 +1270,9 @@ def p(key, sm_count, T): if first <= T <= last or any(a <= T <= b for a, b in spans): continue assert p(key, 212, T) is None, (key, T) - even = K1_212_EVEN_BLOCK_RANGES.get((key[2], key[3], key[4]), ()) + two = K1_212_TWO_SPLIT_RANGES.get((key[2], key[3], key[4]), ()) for T in K1_212_MEASURED_TOKENS: - two_split = _k1_blocks(T) % 2 == 0 and any(a <= T <= b for a, b, _ in even) + two_split = any(a <= T <= b for a, b, _ in two) if T not in held and not first <= T <= last and not two_split: assert p(key, 212, T) is None, (key, T) for T in (131072, 511): @@ -1384,17 +1389,17 @@ def p(key, sm_count, T): assert p(key, 212, T) is None, T -#: The ranges of :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` on 212-SM sm_107a parts, ``(d_in, d_out, out_dtype) -> -#: ((min_tokens, max_tokens, (block_n, num_splits)), ...)``: two splits of the row's 212-SM tile width at a token count -#: that spans an even number of 64-token K blocks. -K1_212_EVEN_BLOCK_RANGES = { +#: The ranges of :data:`K1_TWO_SPLIT_RANGES_BY_SM` on 212-SM sm_107a parts, ``(d_in, d_out, out_dtype) -> +#: ((min_tokens, max_tokens, (block_n, num_splits)), ...)``: two splits of the row's 212-SM tile width at every token +#: count of the range (either number of 64-token K blocks). +K1_212_TWO_SPLIT_RANGES = { (4096, 512, "f32"): ((1024, 8192, (128, 2)), (32768, 65536, (128, 2))), (4096, 1024, "f32"): ((2048, 8192, (192, 2)), (32768, 65536, (192, 2))), (5120, 512, "f32"): ((1024, 8192, (128, 2)), (32768, 65536, (128, 2))), (5120, 1024, "f32"): ((1024, 8192, (256, 2)),), } #: The measured token counts of those rows: 512 to 65536 tokens and the window's 16231 and 16384. -K1_212_EVEN_BLOCK_MEASURED_TOKENS = ( +K1_212_TWO_SPLIT_MEASURED_TOKENS = ( 512, 1024, 2048, @@ -1412,21 +1417,21 @@ def _k1_blocks(T): return -(-int(T) // 64) -def test_k1_212_sm_even_block_ranges(monkeypatch, no_k1_chains): - """On 212-SM sm_107a parts a token count of a range of :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` that spans an even - number of 64-token K blocks runs that range's program (two splits of the row's 212-SM tile width), except a token - count of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for the part; one that spans an odd number (1025, 8193, 40961), and - every token count outside the ranges, keeps the plan the part has without the table. A range starts and ends at a - measured token count with an even block count, lies inside 1024 to 65536 tokens, outside the window and outside the - row's ranges in :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` and :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM`. 208, 210, - 211, 213, 216 and 217 SMs plan every row and token count as without the table.""" - table = cake_backend.K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM - assert set(table) == {("sm_107a", 212)} | set(K1_152_EVEN_BLOCK_RANGES) - # 4096 -> 128: test_k1_212_sm_4096_to_128_even_block_range +def test_k1_212_sm_two_split_ranges(monkeypatch, no_k1_chains): + """On 212-SM sm_107a parts every token count of a range of :data:`K1_TWO_SPLIT_RANGES_BY_SM` (an even or an odd + number of 64-token K blocks: 1025, 2049, 6000, 40961) runs that range's program (two splits of the row's 212-SM tile + width), except a token count of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` for the part; every token count outside the + ranges keeps the plan the part has without the table. A range starts and ends at a measured token count, lies + inside 1024 to 65536 tokens, outside the window and outside the row's ranges in + :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES` and :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM`. 208, 210, 211, 213, 216 and + 217 SMs plan every row and token count as without the table.""" + table = cake_backend.K1_TWO_SPLIT_RANGES_BY_SM + assert set(table) == {("sm_107a", 212)} | set(K1_152_TWO_SPLIT_RANGES) + # 4096 -> 128: test_k1_212_sm_4096_to_128_two_split_range narrow = ("sm_107a", 1, 4096, 128, "f32") rows = {k: v for k, v in table[("sm_107a", 212)].items() if k != narrow} assert set(table[("sm_107a", 212)]) == set(rows) | {narrow} - assert {(k[2], k[3], k[4]) for k in rows} == set(K1_212_EVEN_BLOCK_RANGES) + assert {(k[2], k[3], k[4]) for k in rows} == set(K1_212_TWO_SPLIT_RANGES) by_sm = cake_backend.K1_SPLIT_RULES_BY_SM[("sm_107a", 212)] token_ranges = cake_backend.K1_SPLIT_RULES_BY_SM_TOKEN_RANGES[("sm_107a", 212)] overrides = cake_backend.K1_SPLIT_RANGE_OVERRIDES_BY_SM.get(("sm_107a", 212), {}) @@ -1449,55 +1454,54 @@ def p(M, N, T, sm_count=212): for key, spans in rows.items(): assert key[0] == "sm_107a" and key[1] == 1 and key[4] == "f32", key - want = K1_212_EVEN_BLOCK_RANGES[(key[2], key[3], key[4])] + want = K1_212_TWO_SPLIT_RANGES[(key[2], key[3], key[4])] assert tuple(spans) == want, key verdict = by_sm[key] assert verdict is not None, key for (_a, b, _v), (c, _d, _w) in zip(spans, spans[1:], strict=False): assert b < c, key for lo, hi, choice in spans: - assert lo in K1_212_EVEN_BLOCK_MEASURED_TOKENS, key - assert hi in K1_212_EVEN_BLOCK_MEASURED_TOKENS, key + assert lo in K1_212_TWO_SPLIT_MEASURED_TOKENS, key + assert hi in K1_212_TWO_SPLIT_MEASURED_TOKENS, key assert 1024 <= lo <= hi <= 65536, key - assert _k1_blocks(lo) % 2 == 0 and _k1_blocks(hi) % 2 == 0, key assert hi < first or lo > last, key own = token_ranges.get(key, ()) assert not any(a <= hi and lo <= b for a, b in own), key over = overrides.get(key, ()) assert not any(a <= hi and lo <= b for a, b, _ in over), key assert tuple(choice) == (verdict[0], 2), key - # an odd number of K blocks inside a range (17, 129 and 641 blocks) runs the deterministic program, an even number - # (16, 94 and 1024 blocks) the two-split program + # inside a range the two-split program at an odd number of K blocks (17, 33, 641 blocks) as at an even number + # (16, 94, 128, 512, 1024 blocks); one token outside a range the deterministic program for M, N in ((4096, 512), (5120, 512)): - for T in (1025, 8193, 40961): - assert p(M, N, T) is None, (M, N, T) - for T in (1024, 6000, 65536): + for T in (1024, 1025, 6000, 8192, 32768, 40961, 65536): assert p(M, N, T) == (128, 2), (M, N, T) - for T in (1024, 1025, 2049, 8193, 40961): + for T in (8193, 32767, 65537): + assert p(M, N, T) is None, (M, N, T) + for T in (1024, 8193, 32767, 65537): assert p(4096, 1024, T) is None, T - for T in (2048, 6000, 65536): + for T in (2048, 2049, 6000, 40961, 65536): assert p(4096, 1024, T) == (192, 2), T - for T in (1025, 8193, 32768, 65536): + for T in (1023, 8193, 32768, 65536): assert p(5120, 1024, T) is None, T - for T in (1024, 6000, 8192): + for T in (1024, 1025, 6000, 8192): assert p(5120, 1024, T) == (256, 2), T - # outside the ranges (an even number of blocks each) the rows keep their 212-SM routing: the one-split program at - # 512 tokens (K1_SPLIT_RULES_BY_SM_TOKEN_RANGES), the deterministic program elsewhere outside the window - for M, N, _ in K1_212_EVEN_BLOCK_RANGES: + # outside the ranges the rows keep their 212-SM routing: the one-split program at 512 tokens + # (K1_SPLIT_RULES_BY_SM_TOKEN_RANGES), the deterministic program elsewhere outside the window + for M, N, _ in K1_212_TWO_SPLIT_RANGES: assert p(M, N, 512) == tuple(by_sm[("sm_107a", 1, M, N, "f32")]), (M, N) for T in (8320, 12000, 32640, 65664, 131072): assert p(M, N, T) is None, (M, N, T) # 208 and 216 SMs: the deterministic program (216 SMs: the one-split verdict up to 960 tokens only) rule = tuple(cake_backend.K1_SPLIT_RULES[("sm_107a", 1, M, N, "f32")]) assert p(M, N, 960, sm_count=216) == rule, (M, N) - for T in (1024, 6000, 8192, 32768): + for T in (1024, 6000, 8192, 32768, 40960, 40961, 65536): assert p(M, N, T, sm_count=208) is None, (M, N, T) assert p(M, N, T, sm_count=216) is None, (M, N, T) sms_list = (208, 210, 211, 212, 213, 216, 217) Ts = { T + d - for T in K1_212_EVEN_BLOCK_MEASURED_TOKENS + (6000, 12000, 40961, 131072) + for T in K1_212_TWO_SPLIT_MEASURED_TOKENS + (6000, 12000, 40961, 131072) for d in (-1, 0, 1) } for spans in rows.values(): @@ -1522,12 +1526,12 @@ def plans(): return out with_table = plans() - monkeypatch.setattr(cake_backend, "K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM", {}) + monkeypatch.setattr(cake_backend, "K1_TWO_SPLIT_RANGES_BY_SM", {}) without = plans() routed = 0 for (key, sm_count, T), plan in with_table.items(): held = [tuple(v) for lo, hi, v in rows[key] if lo <= T <= hi] - two_split = sm_count == 212 and bool(held) and _k1_blocks(T) % 2 == 0 + two_split = sm_count == 212 and bool(held) if two_split and not cut(key, sm_count, T): assert plan == held[0], (key, T) assert without[(key, sm_count, T)] is None, (key, T) @@ -1537,26 +1541,26 @@ def plans(): assert routed -#: The ranges of :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` on 152-SM sm_100a and sm_103a parts, ``(arch, sm_count) -> +#: The ranges of :data:`K1_TWO_SPLIT_RANGES_BY_SM` on 152-SM sm_100a and sm_103a parts, ``(arch, sm_count) -> #: ((min_tokens, max_tokens, (block_n, num_splits)), ...)`` of the FP32 row ``4096 -> 128``: two splits of 64-wide -#: tiles at a token count that spans an even number of 64-token K blocks. -K1_152_EVEN_BLOCK_RANGES = { +#: tiles at every token count of the range (either number of 64-token K blocks). +K1_152_TWO_SPLIT_RANGES = { ("sm_100a", 152): ((1024, 4096, (64, 2)),), ("sm_103a", 152): ((1024, 6144, (64, 2)),), } #: The measured token counts of that row on 152 SMs. -K1_152_EVEN_BLOCK_MEASURED_TOKENS = (512, 1024, 1088, 2048, 2112, 4096, 6144, 8192) +K1_152_TWO_SPLIT_MEASURED_TOKENS = (512, 1024, 1088, 2048, 2112, 4096, 6144, 8192) -def test_k1_152_sm_even_block_ranges(monkeypatch): +def test_k1_152_sm_two_split_ranges(monkeypatch): """On 152-SM sm_100a and sm_103a parts (at or above the SM-count floor) the FP32 row ``4096 -> 128``, which has no - verdict in :data:`K1_SPLIT_RULES`, runs two splits of 64-wide tiles at a token count of its range in - :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` that spans an even number of 64-token K blocks: 1024 to 4096 tokens on - sm_100a, 1024 to 6144 on sm_103a. An odd number (1025, 1088, 2112, 4097), a token count outside the range (512, - 1023, 4161 / 6209, 8192), BF16, and every other SM count (148, 150, 151, 153, 160) run the deterministic program - as without the table. A range starts and ends at a measured token count with an even block count and lies outside - the window of :data:`K1_SPLIT_RULES_BY_SM_TOKENS`.""" - table = cake_backend.K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM + verdict in :data:`K1_SPLIT_RULES`, runs two splits of 64-wide tiles at every token count of its range in + :data:`K1_TWO_SPLIT_RANGES_BY_SM`, at an even or an odd number of 64-token K blocks (1025, 1088, 2112, 4033): 1024 + to 4096 tokens on sm_100a, 1024 to 6144 on sm_103a. A token count outside the range (512, 960, 1023, 4097 / 6145, + 8192), BF16, and every other SM count (148, 150, 151, 153, 160) run the deterministic program as without the table. + A range starts and ends at a measured token count and lies outside the window of + :data:`K1_SPLIT_RULES_BY_SM_TOKENS`.""" + table = cake_backend.K1_TWO_SPLIT_RANGES_BY_SM first, last = cake_backend.K1_SPLIT_RULES_BY_SM_TOKENS def p(arch, T, sm_count=152, out_f32=True): @@ -1564,7 +1568,7 @@ def p(arch, T, sm_count=152, out_f32=True): arch=arch, L=1, M=4096, N=128, K=T, out_f32=out_f32, sm_count=sm_count ) - for part, want in K1_152_EVEN_BLOCK_RANGES.items(): + for part, want in K1_152_TWO_SPLIT_RANGES.items(): arch, sms = part key = (arch, 1, 4096, 128, "f32") assert table[part] == {key: want}, part @@ -1575,16 +1579,15 @@ def p(arch, T, sm_count=152, out_f32=True): ) assert floor <= sms, key for lo, hi, choice in want: - assert lo in K1_152_EVEN_BLOCK_MEASURED_TOKENS, key - assert hi in K1_152_EVEN_BLOCK_MEASURED_TOKENS, key - assert _k1_blocks(lo) % 2 == 0 and _k1_blocks(hi) % 2 == 0, key + assert lo in K1_152_TWO_SPLIT_MEASURED_TOKENS, key + assert hi in K1_152_TWO_SPLIT_MEASURED_TOKENS, key assert hi < first or lo > last, key assert tuple(choice) == (64, 2), key - # the measured token counts: two splits at 1024, 2048 and 4096 tokens (and 6144 on sm_103a), the deterministic - # program at 512 and 8192 tokens and at an odd number of K blocks (1088, 2112) - det_ts = (512, 960, 1023, 1025, 1088, 2112, 4097, 8192, 16231) + # the measured token counts: two splits at 1024, 1088, 2048, 2112 and 4096 tokens (and 6144 on sm_103a), the + # deterministic program at 512 and 8192 tokens + det_ts = (512, 960, 1023, 8192, 16231) for arch, top in (("sm_100a", 4096), ("sm_103a", 6144)): - for T in (1024, 1100, 2048, 3072, 4033, 4096, top): + for T in (1024, 1025, 1088, 1100, 2048, 2112, 3072, 4033, 4096, top): assert p(arch, T) == (64, 2), (arch, T) for T in det_ts + (top + 1, top + 64, top + 65): assert p(arch, T) is None, (arch, T) @@ -1595,30 +1598,30 @@ def p(arch, T, sm_count=152, out_f32=True): sms_list = (148, 150, 151, 152, 153, 160) near = (-64, -1, 0, 1, 64) - Ts = {T + d for T in K1_152_EVEN_BLOCK_MEASURED_TOKENS for d in near} + Ts = {T + d for T in K1_152_TWO_SPLIT_MEASURED_TOKENS for d in near} Ts |= {16172, 16231, 16384, 65536} def plans(): return { (arch, sms, T, f32): p(arch, T, sm_count=sms, out_f32=f32) - for arch, _ in K1_152_EVEN_BLOCK_RANGES + for arch, _ in K1_152_TWO_SPLIT_RANGES for sms in sms_list for T in Ts for f32 in (False, True) } with_table = plans() - monkeypatch.setattr(cake_backend, "K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM", {}) + monkeypatch.setattr(cake_backend, "K1_TWO_SPLIT_RANGES_BY_SM", {}) without = plans() routed = 0 for (arch, sms, T, f32), plan in with_table.items(): assert without[(arch, sms, T, f32)] is None, (arch, sms, T, f32) held = [ tuple(v) - for lo, hi, v in K1_152_EVEN_BLOCK_RANGES.get((arch, sms), ()) + for lo, hi, v in K1_152_TWO_SPLIT_RANGES.get((arch, sms), ()) if lo <= T <= hi ] - if f32 and held and _k1_blocks(T) % 2 == 0: + if f32 and held: assert plan == held[0], (arch, sms, T) routed += 1 else: @@ -1626,24 +1629,24 @@ def plans(): assert routed -#: The range of :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` of the FP32 row ``4096 -> 128`` on 212-SM sm_107a parts (no +#: The range of :data:`K1_TWO_SPLIT_RANGES_BY_SM` of the FP32 row ``4096 -> 128`` on 212-SM sm_107a parts (no #: verdict in :data:`K1_SPLIT_RULES` or in the part's table of :data:`K1_SPLIT_RULES_BY_SM`): two splits of 64-wide -#: tiles at a token count that spans an even number of 64-token K blocks. +#: tiles at every token count of the range (either number of 64-token K blocks). K1_212_4096_128_RANGE = ((1024, 7680, (64, 2)),) #: The measured token counts of that row on 212 SMs. K1_212_N128_TOKENS = (512, 768, 896, 1024, 2048, 4096, 6144, 6656, 7168, 7680) -def test_k1_212_sm_4096_to_128_even_block_range(monkeypatch): +def test_k1_212_sm_4096_to_128_two_split_range(monkeypatch): """On 212-SM sm_107a parts (below the SM-count floor) the FP32 row ``4096 -> 128``, which has no verdict in - :data:`K1_SPLIT_RULES` or in the part's table of :data:`K1_SPLIT_RULES_BY_SM`, runs two splits of 64-wide tiles at a - token count from 1024 to 7680 that spans an even number of 64-token K blocks. An odd number (1025, 1088, 3000, - 4160, 5000, 7615), a token count outside the range (512, 896, 1023, 7745, 8192, 16384), BF16, and every other SM - count (208, 210, 211, 213, 216, 217) run the deterministic program as without the table.""" + :data:`K1_SPLIT_RULES` or in the part's table of :data:`K1_SPLIT_RULES_BY_SM`, runs two splits of 64-wide tiles at + every token count from 1024 to 7680, at an even or an odd number of 64-token K blocks (1025, 1088, 3000, 4160, 5000, + 7615). A token count outside the range (512, 896, 1023, 7681, 7745, 8192, 16384), BF16, and every other SM count + (208, 210, 211, 213, 216, 217) run the deterministic program as without the table.""" key = ("sm_107a", 1, 4096, 128, "f32") part = ("sm_107a", 212) - even = cake_backend.K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM[part] - assert even[key] == K1_212_4096_128_RANGE + two = cake_backend.K1_TWO_SPLIT_RANGES_BY_SM[part] + assert two[key] == K1_212_4096_128_RANGE assert cake_backend.K1_SPLIT_RULES[key] is None assert key not in cake_backend.K1_SPLIT_RULES_BY_SM[part] assert key not in cake_backend.K1_SPLIT_RULES_BY_SM_TOKEN_RANGES.get(part, {}) @@ -1652,7 +1655,6 @@ def test_k1_212_sm_4096_to_128_even_block_range(monkeypatch): for lo, hi, choice in K1_212_4096_128_RANGE: assert lo in K1_212_N128_TOKENS, key assert hi in K1_212_N128_TOKENS, key - assert _k1_blocks(lo) % 2 == 0 and _k1_blocks(hi) % 2 == 0, key assert hi < first or lo > last, key assert tuple(choice) == (64, 2), key @@ -1661,9 +1663,11 @@ def p(T, sm_count=212, out_f32=True): arch="sm_107a", L=1, M=4096, N=128, K=T, out_f32=out_f32, sm_count=sm_count ) - for T in (1024, 1100, 2048, 3072, 4096, 6144, 6656, 7168, 7680): + for T in (1024, 1025, 1088, 1100, 2048, 3000, 3072, 4096, 4160, 5000, 6144, 7615): + assert p(T) == (64, 2), T + for T in (6656, 7168, 7680): assert p(T) == (64, 2), T - for T in (512, 896, 1023, 1025, 1088, 3000, 4160, 5000, 7615, 7745, 8192, 16384): + for T in (512, 896, 1023, 7681, 7745, 8192, 16384): assert p(T) is None, T for T in (1024, 4096, 7680): assert p(T, out_f32=False) is None, T @@ -1682,13 +1686,13 @@ def plans(): } with_table = plans() - monkeypatch.setattr(cake_backend, "K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM", {}) + monkeypatch.setattr(cake_backend, "K1_TWO_SPLIT_RANGES_BY_SM", {}) without = plans() routed = 0 for (sms, T, f32), plan in with_table.items(): assert without[(sms, T, f32)] is None, (sms, T, f32) held = [tuple(v) for lo, hi, v in K1_212_4096_128_RANGE if lo <= T <= hi] - if sms == 212 and f32 and held and _k1_blocks(T) % 2 == 0: + if sms == 212 and f32 and held: assert plan == held[0], T routed += 1 else: @@ -1708,17 +1712,30 @@ def plans(): ("sm_100a", 5120, 64, "f32"): ((512, 1024), (8192, 8192)), ("sm_103a", 4096, 32, "bf16"): ((65536, None),), ("sm_103a", 4096, 32, "f32"): ((65536, None),), - ("sm_103a", 5120, 48, "bf16"): ((512, 2048),), + ("sm_103a", 5120, 48, "bf16"): ((512, 2048), (40961, 49151)), ("sm_103a", 5120, 48, "f32"): ((512, 2048),), - ("sm_103a", 5120, 64, "bf16"): ((512, 1024),), + ("sm_103a", 5120, 64, "bf16"): ((512, 1024), (47105, 49152)), ("sm_103a", 5120, 64, "f32"): ((512, 1024),), ("sm_103a", 5120, 128, "bf16"): ((512, 4096), (32768, 65536)), ("sm_103a", 5120, 128, "f32"): ((512, 4096), (32768, 65536)), ("sm_103a", 6144, 128, "bf16"): ((512, 8192),), ("sm_103a", 6144, 128, "f32"): ((512, 8192),), } -#: The measured token counts of the 148-SM audit. -K1_148_MEASURED_TOKENS = (512, 1024, 2048, 4096, 8192, 16231, 32768, 65536) +#: The measured token counts of the 148-SM audit (40961, 47105 and 49152: points of the 10-seed FP64 error grid of +#: sm_103a ``5120 -> 48`` / ``5120 -> 64`` BF16). +K1_148_MEASURED_TOKENS = ( + 512, + 1024, + 2048, + 4096, + 8192, + 16231, + 32768, + 40961, + 47105, + 49152, + 65536, +) def test_k1_exact_sm_count_exclusions(monkeypatch): @@ -1726,8 +1743,8 @@ def test_k1_exact_sm_count_exclusions(monkeypatch): :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM` on a part of exactly that SM count (148-SM sm_100a / sm_103a parts), and every part of any other SM count (147, 149, 152, 160, 208, 212, 216, 220) plans every row and token count as without the table. Every entry belongs to a row of :data:`K1_SPLIT_RULES` with a program on that part; its ranges - are closed, disjoint and not adjacent, start at a measured token count and end at one or have no upper bound; one - token outside a range the row keeps the plan it has without the table.""" + are closed, disjoint and not adjacent, start at a measured token count and end at one, one token below one, or + have no upper bound; one token outside a range the row keeps the plan it has without the table.""" table = cake_backend.K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM got_rows = { (arch, k[2], k[3], k[4]): tuple(v) @@ -1786,7 +1803,7 @@ def plans(sms_list, Ts): assert any(b < T < c for T in K1_148_MEASURED_TOKENS), key for lo, hi in spans: assert lo in K1_148_MEASURED_TOKENS, key - assert hi is None or hi in K1_148_MEASURED_TOKENS, key + assert hi is None or {hi, hi + 1} & set(K1_148_MEASURED_TOKENS), key assert without[(key, 148, lo)] is not None, key assert with_table[(key, 148, lo)] is None, key @@ -1796,9 +1813,10 @@ def test_k1_148_comb_rows_run_det_from_65536(): deterministic program from 65536 tokens up, at odd (65537, 65600, 67520) and even (65601, 67584, 98304, 131072) numbers of 64-token blocks, the split program one token below. On 152 SMs ``sm_100a`` ``4096 -> 64`` runs the deterministic program above 65536 tokens in BF16 (its 152-SM entry of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`, - open-ended) and in FP32 at odd numbers of 64-token blocks (:data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS`), and keeps - the split program in FP32 at even numbers; ``sm_103a`` ``4096 -> 32`` runs the deterministic program there too - (its 152-SM entries of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`, open-ended).""" + open-ended) and in FP32 at odd numbers of 64-token blocks from 65538 tokens up + (:data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS`), and keeps the split program in FP32 at 65537 tokens and at even + numbers; ``sm_103a`` ``4096 -> 32`` runs the deterministic program there too (its 152-SM entries of + :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`, open-ended).""" for arch, N in (("sm_100a", 64), ("sm_103a", 32)): for f32 in (False, True): kw = dict(arch=arch, L=1, M=4096, N=N, out_f32=f32) @@ -1810,7 +1828,7 @@ def test_k1_148_comb_rows_run_det_from_65536(): assert cake_backend.k1_plan(K=65535, sm_count=152, **kw) == (64, 2), kw for T in (65536, 65537, 67520, 67584, 131072): odd = -(-T // 64) % 2 == 1 - want = (64, 2) if f32 and not odd else None + want = (64, 2) if f32 and (not odd or T < 65538) else None assert cake_backend.k1_plan(K=T, sm_count=152, **kw) == want, (kw, T) kw = dict(arch="sm_103a", L=1, M=4096, N=32, out_f32=f32) for T in (65536, 65537, 67520, 67584, 131072): @@ -1915,8 +1933,8 @@ def test_k1_sm_count_overrides_hold_only_at_their_sm_count_and_token_counts( """:data:`K1_SPLIT_OVERRIDES_BY_SM` has no entry: ``4096 -> 1024`` FP32 keeps its :data:`K1_SPLIT_RULES` verdict (256-wide tiles at one split) up to 960 tokens from 216 SMs up and runs the deterministic program from 961 tokens up, the window of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` included; on 215 SMs (no table) it runs the deterministic - program, and on 212 SMs its own tables decide (the 192-wide row-major program at one split at 512 tokens, the - two-split program of :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` at an even K-block count inside its ranges, else + program, and on 212 SMs its own tables decide (the 192-wide row-major program at one split in the window, except at + 16231 tokens, and at 512 tokens, the two-split program of :data:`K1_TWO_SPLIT_RANGES_BY_SM` inside its ranges, else the deterministic program). On 212 SMs at 512 tokens the launch is the row-major-store program's: 88 tiles, 176 CTAs. A table entry acts on exactly its SM count, at the window's token counts only: ``None`` keeps a row on the deterministic program there, a verdict replaces the row's :data:`K1_SPLIT_RULES` verdict there.""" @@ -1933,15 +1951,17 @@ def plan(sms, T, M=d_in, N=d_out): for sms in (216, 217, 280): for T in (512, 704, 960): assert plan(sms, T) == (256, 1), (sms, T) - for T in (961, 1024, 8192, 16172, 16231, 16384, 32768, 131072): + for T in (961, 1024, 8192, 16172, 16231, 16384, 32768, 40960, 131072): assert plan(sms, T) is None, (sms, T) for T in (512, 16172, 16231, 16300, 16384): assert plan(215, T) is None, T - assert plan(208, T) is None, T assert plan(212, 512) == (192, 1) - even_212 = K1_212_EVEN_BLOCK_RANGES[(d_in, d_out, "f32")] - for T in (1024, 2048, 8192, 16171, 16172, 16231, 16384, 16385, 32768): - want = [v for lo, hi, v in even_212 if lo <= T <= hi and _k1_blocks(T) % 2 == 0] + for T in (16172, 16300, 16384): + assert plan(212, T) == (192, 1), T + assert plan(212, 16231) is None + two_212 = K1_212_TWO_SPLIT_RANGES[(d_in, d_out, "f32")] + for T in (1024, 2048, 2049, 8192, 16171, 16385, 32768, 40960): + want = [v for lo, hi, v in two_212 if lo <= T <= hi] assert plan(212, T) == (want[0] if want else None), T # every other row and token count on 216 SMs: the record verdict; BF16 4096 -> 1024 has none assert ( @@ -1992,24 +2012,37 @@ def plan(sms, T, M=d_in, N=d_out): def test_k1_sm107a_token_counts_of_the_deterministic_rules(no_k1_chains): - """The sm_107a entries of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` at 16172 to 16384 tokens are the single token - counts 16231 and 16384 per SM count: an entry of 208 or 212 SMs holds on exactly that SM count (not on the other - table's part, nor from 216 SMs up), an entry of 216 SMs from 216 SMs up; the neighbouring token counts keep the - row's verdict. The four one-split FP32 rows run the deterministic program at every one of these token counts on - 208, 212 and 216 SMs (their ranges: the next test), where :data:`K1_CHAIN_RANGES_BY_SM` (taken out here) does not - run three accumulation chains.""" + """The sm_107a entries of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS` at 16172 to 16384 tokens are the single token counts + 16231 and 16384 per SM count: an entry of 208 or 212 SMs holds on exactly that SM count (not on the other table's + part, nor from 216 SMs up), an entry of 216 SMs from 216 SMs up; the neighbouring token counts keep the row's + verdict. The four one-split FP32 rows run the deterministic program at every one of these token counts on 216 SMs + (their 216-SM ranges: the next test). These are the tables under :data:`K1_CHAIN_RANGES_BY_SM` (taken out here), + whose ranges run three accumulation chains on 212 and 216 SMs. + """ def p(M, N, sms, T, out_f32=True): return cake_backend.k1_plan( arch="sm_107a", L=1, M=M, N=N, K=T, out_f32=out_f32, sm_count=sms ) - # 4096 / 5120 -> 512 / 1024 FP32: the deterministic program at 16172 to 16384 tokens on 208, 212 and 216 SMs + # 4096 -> 512: the deterministic program at 16231 and 16384 tokens on 208 and 212 SMs + for sms in (208, 212): + assert p(4096, 512, sms, 16231) is None and p(4096, 512, sms, 16384) is None + for T in (16172, 16230, 16232, 16383): + assert p(4096, 512, sms, T) == (128, 1), (sms, T) + # 4096 -> 1024: at 16231 tokens on 208 and 212 SMs, and at 16384 on 208 SMs (a strict re-measurement of the + # restored window: the deterministic plan there is faster); 212 SMs (their table) keep the 192-wide row-major + # program at 16384, and 208 SMs the 256-wide tiles below it + for sms in (208, 212): + assert p(4096, 1024, sms, 16231) is None, sms + assert p(4096, 1024, 208, 16384) is None + assert p(4096, 1024, 208, 16383) == (256, 1) + assert p(4096, 1024, 212, 16384) == (192, 1) + # the four one-split FP32 rows on 216 SMs: the deterministic program at every one of these token counts for M in (4096, 5120): for N in (512, 1024): - for sms in (208, 212, 216): - for T in (16172, 16230, 16231, 16232, 16383, 16384): - assert p(M, N, sms, T) is None, (M, N, sms, T) + for T in (16172, 16230, 16231, 16232, 16383, 16384): + assert p(M, N, 216, T) is None, (M, N, T) # 4096 -> 48 and 4096 -> 64: entries on 208 and 216 SMs, none on 212 for N in (48, 64): assert p(4096, N, 212, 16231) == (64, 3) and p(4096, N, 212, 16384) == (64, 3) @@ -2023,6 +2056,14 @@ def p(M, N, sms, T, out_f32=True): assert p(5120, N, 212, 16384) == (64, 3), N assert p(5120, 48, 208, 16384) is None assert p(5120, 64, 208, 16384) == (64, 3) + # 5120 -> 512: 212 SMs only (216 SMs: above) + assert p(5120, 512, 212, 16231) is None and p(5120, 512, 212, 16384) is None + assert p(5120, 512, 208, 16231) == (128, 1) + # 5120 -> 1024 FP32: no entry at these token counts on 208 and 212 SMs (216 SMs: above); both keep the 256-wide + # one-split program at both + for sms in (208, 212): + for T in (16231, 16384): + assert p(5120, 1024, sms, T) == (256, 1), (sms, T) # 5120 -> 1024 BF16: 16231 tokens from 216 SMs up (the neighbouring token counts and 16384 keep the 256-wide # one-split program); 208 and 212 SMs run the deterministic program at every token count (their tables) for sms in (216, 217, 280): @@ -2046,8 +2087,9 @@ def p(M, N, sms, T, out_f32=True): #: The first token count at which the sm_107a FP32 rows ``4096 / 5120 -> 512 / 1024`` run the deterministic program on #: parts of 216 SMs or more: from 961 tokens (16 K blocks) up the deterministic program of these rows splits K in two -#: (adding the partials in a fixed order), and their one-split program's FP64 relative Frobenius error is about twice -#: its error; up to 960 tokens both add one chain. +#: (adding the partials in a fixed order) while their one-split program adds one chain; up to 960 tokens both add one +#: chain. The one-split program met the FP32 accuracy budget up to 40960 tokens (measured with the 216-SM plan on +#: 212-SM parts); its time from 961 tokens up on 216-SM parts has not been measured. K1_SM107A_FP32_ONE_SPLIT_TOKENS = 961 #: Those rows ``(d_in, d_out)`` and their one-split programs ``(block_n, splits)`` below that token count. K1_SM107A_WIDE_FP32_ROWS = { @@ -2059,15 +2101,15 @@ def p(M, N, sms, T, out_f32=True): def test_k1_sm107a_wide_fp32_one_split_rows(no_k1_chains): - """The one-split programs of the sm_107a FP32 rows ``4096 -> 512``, ``4096 -> 1024``, ``5120 -> 512`` and - ``5120 -> 1024`` fail the FP64 comparison wherever the deterministic program of these rows splits K: from - ``K1_SM107A_FP32_ONE_SPLIT_TOKENS`` up, with no upper bound, these rows run the deterministic program on parts of - 216 SMs or more, and at 16172 to 16384 tokens on 208 and 212 SMs (their entries in - :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`); they keep their one-split program up to 960 tokens on 216 SMs and up, and - at 512 tokens on 212 SMs (:data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`). 212 SMs keep the two-split ranges at an - even number of 64-token K blocks (``K1_212_EVEN_BLOCK_RANGES``). No other sm_107a row has a range without an - upper bound, and sm_100a / sm_103a have no verdict for these rows. These are the tables under - :data:`K1_CHAIN_RANGES_BY_SM` (taken out here), whose ranges run three accumulation chains on 212 and 216 SMs.""" + """The sm_107a FP32 rows ``4096 -> 512``, ``4096 -> 1024``, ``5120 -> 512`` and ``5120 -> 1024`` run the + deterministic program from ``K1_SM107A_FP32_ONE_SPLIT_TOKENS`` up, with no upper bound, on parts of 216 SMs or more + (their entries in :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`), and keep their one-split program up to 960 tokens there. + The 208-SM and 212-SM tables hold only at 16172 to 16384 tokens and at 512 tokens on 212 SMs + (:data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`), so those parts run the deterministic program at these token counts, + except the 212-SM two-split ranges at either number of 64-token K blocks (``K1_212_TWO_SPLIT_RANGES``); the other + entries of these rows lie in that window. No other sm_107a row has a range without an upper bound, and sm_100a / + sm_103a have no verdict for these rows. These are the tables under :data:`K1_CHAIN_RANGES_BY_SM` (taken out here), + whose ranges run three accumulation chains on 212 and 216 SMs.""" T0 = K1_SM107A_FP32_ONE_SPLIT_TOKENS rows = K1_SM107A_WIDE_FP32_ROWS excluded = cake_backend.K1_SPLIT_RULE_EXCLUDED_TOKENS @@ -2087,36 +2129,19 @@ def p(M, N, sms, T, arch="sm_107a"): for (M, N), choice in rows.items(): key = ("sm_107a", 1, M, N, "f32") assert cake_backend.K1_SPLIT_RULES[key] == choice - assert sorted(excluded[key]) == [ - (208, 16172, 16384), - (212, 16172, 16384), - (216, T0, None), - ], key + assert [e for e in excluded[key] if e[2] is None] == [(216, T0, None)] + for _, lo, hi in excluded[key]: + assert hi is None or 16172 <= lo <= hi <= 16384, key for sms in (216, 217, 280): for T in (512, 704, T0 - 1): assert p(M, N, sms, T) == choice, (M, N, sms, T) - for T in (T0, 1024, 16172, 16231, 32768, 45056, 65536, 65537, 131072): + for T in (T0, 1024, 16172, 16231, 32768, 40960, 45056, 65536, 131072): assert p(M, N, sms, T) is None, (M, N, sms, T) assert p(M, N, 212, 512) == tuple(by_sm[("sm_107a", 212)][key]), (M, N) - even = K1_212_EVEN_BLOCK_RANGES[(M, N, "f32")] + two = K1_212_TWO_SPLIT_RANGES[(M, N, "f32")] for sms in (208, 212, 215): - for T in ( - T0 - 1, - T0, - 1024, - 16172, - 16231, - 16384, - 32768, - 40960, - 65536, - 131072, - ): - held = [ - tuple(v) - for lo, hi, v in even - if sms == 212 and lo <= T <= hi and _k1_blocks(T) % 2 == 0 - ] + for T in (T0 - 1, T0, 1024, 1025, 32768, 40960, 40961, 65536, 131072): + held = [tuple(v) for lo, hi, v in two if sms == 212 and lo <= T <= hi] want = held[0] if held else None assert p(M, N, sms, T) == want, (M, N, sms, T) for arch, sms in (("sm_100a", 148), ("sm_100a", 152), ("sm_103a", 152)): @@ -2127,24 +2152,34 @@ def p(M, N, sms, T, arch="sm_107a"): #: The entries of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS`, ``key -> ((sm_count, min_tokens, max_tokens), ...)``. K1_ODD_BLOCK_EXCLUDED_TOKENS = { - ("sm_100a", 1, 4096, 32, "f32"): ((148, 1025, None),), - ("sm_100a", 1, 4096, 48, "f32"): ((148, 1025, None),), - ("sm_100a", 1, 4096, 64, "f32"): ((148, 1025, None),), - ("sm_100a", 1, 5120, 32, "f32"): ((148, 12289, None),), - ("sm_100a", 1, 5120, 48, "f32"): ((148, 12289, None),), - ("sm_103a", 1, 4096, 32, "f32"): ((148, 1025, None),), - ("sm_103a", 1, 4096, 48, "f32"): ((152, 1025, None),), - ("sm_107a", 1, 8192, 128, "f32"): ((212, 1025, 8192), (216, 1025, 8192)), + ("sm_100a", 1, 4096, 32, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 4096, 48, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 4096, 64, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 5120, 32, "f32"): ((148, 65538, None),), + ("sm_100a", 1, 5120, 48, "f32"): ((148, 65538, None),), + ("sm_103a", 1, 4096, 32, "f32"): ((148, 65538, None),), + ("sm_103a", 1, 4096, 48, "f32"): ((152, 14338, 32767), (152, 65538, None)), + ("sm_107a", 1, 8192, 128, "f32"): ((216, 1025, 8192),), } +#: The ranges of ``K1_ODD_BLOCK_EXCLUDED_TOKENS`` that come from the time (the backward pass or the checkpointed +#: step more than 1.3 % slower at an odd block count), not from the FP32 budget. +K1_ODD_BLOCK_TIME_RANGES = {("sm_103a", 1, 4096, 48, "f32"): ((152, 14338, 32767),)} +#: The ranges of ``K1_ODD_BLOCK_EXCLUDED_TOKENS`` on parts whose time at an odd block count was not measured (sm_107a +#: ``8192 -> 128`` FP32 from 216 SMs up: measured on 212-SM parts only). +K1_ODD_BLOCK_UNTIMED_RANGES = {("sm_107a", 1, 8192, 128, "f32"): ((216, 1025, 8192),)} def test_k1_odd_block_exclusions(monkeypatch): - """A row of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS` runs the deterministic program at a token count of an - entry's range that spans an odd number of 64-token K blocks, on a part of the entry's SM count or more where the - entry is at or above the row's SM-count floor and on exactly that SM count below it; at an even number, below the - range and above a range's upper bound the row keeps the route it takes without the table. Every entry names an - FP32 two-split row with a verdict on its part, starts at an odd block count (the first token count of a block) and - ends at an even one or has no upper bound; BF16 rows have no entry.""" + """A row of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS` runs the deterministic program at a token count of an entry's + range that spans an odd number of 64-token K blocks, on a part of the entry's SM count or more where the entry is at + or above the row's SM-count floor and on exactly that SM count below it; at an even number and below the range the + row keeps the route it takes without the table. Every entry names an FP32 two-split row with a verdict on its part + and starts at 65538 tokens (above the largest measured token count at which the two-split program met the FP32 + accuracy budget, 65537) with no upper bound, except the ranges of ``K1_ODD_BLOCK_TIME_RANGES`` (sm_103a + ``4096 -> 48`` from 14338 to 32767 tokens on 152 SMs: the time) and of ``K1_ODD_BLOCK_UNTIMED_RANGES`` (sm_107a + ``8192 -> 128`` from 1025 to 8192 tokens from 216 SMs up: the time not measured there); BF16 rows have no entry. + Below 65538 tokens and outside those ranges these rows run their two-split program at either number of K + blocks.""" table = cake_backend.K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS assert table == K1_ODD_BLOCK_EXCLUDED_TOKENS @@ -2164,8 +2199,32 @@ def p(key, sms, T): cake_backend.K1_RULE_MIN_SM_COUNT.get(key, 0), ) for sms, lo, hi in entries: - assert lo % 64 == 1 and _k1_blocks(lo) % 2 == 1, key - assert hi is None or (hi % 64 == 0 and _k1_blocks(hi) % 2 == 0), key + if (sms, lo, hi) in K1_ODD_BLOCK_TIME_RANGES.get(key, ()): + assert sms >= floor and hi is not None and _k1_blocks(lo) % 2 == 1, key + odd = (lo, 64 * _k1_blocks(lo), 16173, 20481, 28673, hi - 64) + even = (64 * _k1_blocks(lo) + 1, 20480, hi) + outside = (lo - 1, 1025) # odd block counts below the range + assert all(_k1_blocks(T) % 2 == 1 for T in odd + outside), key + assert all(_k1_blocks(T) % 2 == 0 for T in even), key + for part in (sms, sms + 1, sms + 64): + cases += [(key, part, T, True) for T in odd] + cases += [(key, part, T, False) for T in even + outside] + continue + if (sms, lo, hi) in K1_ODD_BLOCK_UNTIMED_RANGES.get(key, ()): + assert sms >= floor and lo % 64 == 1 and _k1_blocks(lo) % 2 == 1, key + assert hi % 64 == 0 and _k1_blocks(hi) % 2 == 0, key + assert cake_backend.K1_SPLIT_RULES[key][1] == 2, key + odd = (lo, lo + 63, lo + 3072, hi - 64) + even = (lo - 1, lo + 64, lo + 3071, hi) + # odd block counts below and above the range + outside = (lo - 128, hi + 1) + assert all(_k1_blocks(T) % 2 == 1 for T in odd + outside), key + assert all(_k1_blocks(T) % 2 == 0 for T in even), key + for part in (sms, sms + 1, sms + 64): + cases += [(key, part, T, True) for T in odd] + cases += [(key, part, T, False) for T in even + outside] + continue + assert lo == 65538 and _k1_blocks(lo) % 2 == 1 and hi is None, key if sms >= floor: rule = cake_backend.K1_SPLIT_RULES[key] parts = (sms, sms + 1, sms + 64) @@ -2173,10 +2232,10 @@ def p(key, sms, T): rule = cake_backend.K1_SPLIT_RULES_BY_SM[(arch, sms)][key] parts = (sms,) assert rule is not None and rule[1] == 2, (key, sms) - top = (16173, 65537, 131073) if hi is None else (hi - 64,) - odd = (lo, lo + 63, lo + 3072) + top - even = (lo - 1, lo + 64, lo + 3071) + ((65536,) if hi is None else (hi,)) - outside = (lo - 128,) + (() if hi is None else (hi + 1,)) + last = 64 * _k1_blocks(lo) # the last token count of lo's block + odd = (lo, last, last + 65, lo + 3072, 98305, 131073) + even = (last + 1, 65664, 98304, 131072) + outside = (lo - 128, lo - 1) # odd block counts below the range assert all(_k1_blocks(T) % 2 == 1 for T in odd + outside), key assert all(_k1_blocks(T) % 2 == 0 for T in even), key for part in parts: @@ -2190,17 +2249,33 @@ def p(key, sms, T): assert got[(key, part, T)] is None, (key, part, T) else: assert got[(key, part, T)] == want, (key, part, T) - # the entries change the route of measured odd token counts: without the table these take the two-split verdict + # the entries change the route of odd token counts from 65538 up: without the table these take the two-split + # verdict for key, part, T in ( - (("sm_100a", 1, 4096, 32, "f32"), 148, 1025), - (("sm_100a", 1, 5120, 32, "f32"), 148, 16173), - (("sm_103a", 1, 4096, 32, "f32"), 148, 4097), - (("sm_103a", 1, 4096, 48, "f32"), 152, 16173), - (("sm_107a", 1, 8192, 128, "f32"), 212, 4097), - (("sm_107a", 1, 8192, 128, "f32"), 216, 1025), + (("sm_100a", 1, 4096, 32, "f32"), 148, 65538), + (("sm_100a", 1, 5120, 48, "f32"), 148, 98305), + (("sm_103a", 1, 4096, 32, "f32"), 149, 65665), ): assert got[(key, part, T)] is None, (key, part, T) - assert p(key, part, T) is not None, (key, part, T) + assert p(key, part, T) == (64, 2), (key, part, T) + # below 65538 tokens odd block counts take the two-split program with the table in place (1025, 4097, 14337, 16173 + # and 65537 tokens; sm_103a 4096 -> 48 on 152 SMs not at 14338 to 32767); sm_107a ``8192 -> 128`` FP32 has an entry + # from 216 SMs up only + monkeypatch.undo() + table = cake_backend.K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS + assert table == K1_ODD_BLOCK_EXCLUDED_TOKENS + for key, part, T, want in ( + (("sm_100a", 1, 4096, 32, "f32"), 148, 1025, (64, 2)), + (("sm_100a", 1, 4096, 64, "f32"), 152, 65537, (64, 2)), + (("sm_100a", 1, 5120, 32, "f32"), 148, 16173, (64, 2)), + (("sm_103a", 1, 4096, 32, "f32"), 148, 4097, (64, 2)), + (("sm_103a", 1, 4096, 48, "f32"), 152, 14337, (64, 2)), + (("sm_103a", 1, 4096, 48, "f32"), 152, 16173, None), + (("sm_107a", 1, 8192, 128, "f32"), 212, 4097, (128, 2)), + (("sm_107a", 1, 8192, 128, "f32"), 216, 1025, None), + (("sm_107a", 1, 8192, 128, "f32"), 216, 1024, (128, 2)), + ): + assert p(key, part, T) == want, (key, part, T) #: The ranges of :data:`K1_CHAIN_RANGES_BY_SM` (sm_107a FP32 rows), ``sm_count -> {(d_in, d_out): ((min_tokens, @@ -2239,18 +2314,19 @@ def p(key, sms, T): } -def test_k1_chain_ranges(): +def test_k1_chain_ranges(monkeypatch): """On sm_107a parts of exactly 212 and 216 SMs the FP32 rows ``4096 / 5120 -> 512 / 1024`` run this package's one-split program of three accumulation chains at the token counts of their ranges in :data:`K1_CHAIN_RANGES_BY_SM` (both ends, one token inside them and the measured token counts inside): :func:`k1_chain_plan` gives the range's - ``(block_n, chains)`` and :func:`k1_plan` that width at one split, at an even number of 64-token K blocks inside the - row's 212-SM two-split range of :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` too; past the chain ranges the two-split - ranges keep two splits. Every range starts and ends at a measured token count and holds no failing one; one token + ``(block_n, chains)`` and :func:`k1_plan` that width at one split, inside the row's 212-SM two-split range of + :data:`K1_TWO_SPLIT_RANGES_BY_SM` too; past the chain ranges the two-split ranges keep two splits at either number + of 64-token K blocks. Every range starts and ends at a measured token count and holds no failing one; one token outside a range, at a failing token count, on another SM count, on another architecture, for a BF16 output and for a batched launch the rows keep the routing of the other tables.""" table = cake_backend.K1_CHAIN_RANGES_BY_SM assert cake_backend.K1_CHAINS_CHOICES == (1, 3) assert set(table) == {("sm_107a", sms) for sms in K1_CHAIN_RANGES} + failing = {} for sms, rows in K1_CHAIN_RANGES.items(): part = ("sm_107a", sms) assert {(k[2], k[3]) for k in table[part]} == set(rows) @@ -2280,7 +2356,7 @@ def test_k1_chain_ranges(): assert cake_backend.k1_chain_plan(K=T, **kw) is None, (key, T) for T in misses: assert cake_backend.k1_chain_plan(K=T, **kw) is None, (key, sms, T) - assert cake_backend.k1_plan(K=T, **kw) is None, (key, sms, T) + failing[(key, sms, T)] = cake_backend.k1_plan(K=T, **kw) for T in (961, 4097, 16172, 40961): for other in (208, 210, 211, 213, 215, 217, 220): assert ( @@ -2295,9 +2371,10 @@ def test_k1_chain_ranges(): assert ( cake_backend.k1_chain_plan(K=T, **dict(kw, arch=arch)) is None ), (key, arch, T) - # on 212 SMs the even block counts of the two-split ranges inside the chain ranges run three chains: 4096 -> 512 at - # 1024 and 8192 tokens (16 and 128 blocks), 4096 -> 1024 at 2048 tokens, 5120 -> 1024 at 1024 tokens, the 512-wide - # rows at 32768 tokens; past the chain ranges the two-split ranges keep two splits (65536 tokens, 1024 blocks) + # on 212 SMs the token counts of the two-split ranges inside the chain ranges run three chains: 4096 -> 512 at + # 1024 and 8192 tokens, 4096 -> 1024 at 2048 tokens, 5120 -> 1024 at 1024 tokens, the 512-wide rows at 32768 + # tokens; past the chain ranges the two-split ranges keep two splits (40962 and 45057 tokens: 641 and 705 blocks; + # 65535 and 65536 tokens: 1024 blocks) p212 = dict(arch="sm_107a", L=1, out_f32=True, sm_count=212) for M, N, T, block_n in ( (4096, 512, 1024, 128), @@ -2310,15 +2387,35 @@ def test_k1_chain_ranges(): assert cake_backend.k1_chain_plan(M=M, N=N, K=T, **p212) == (block_n, 3) assert cake_backend.k1_plan(M=M, N=N, K=T, **p212) == (block_n, 1) for M, N, block_n in ((4096, 512, 128), (4096, 1024, 192), (5120, 512, 128)): - assert cake_backend.k1_chain_plan(M=M, N=N, K=65536, **p212) is None - assert cake_backend.k1_plan(M=M, N=N, K=65536, **p212) == (block_n, 2) + for T in (40962, 45057, 65535, 65536): + assert cake_backend.k1_chain_plan(M=M, N=N, K=T, **p212) is None + assert cake_backend.k1_plan(M=M, N=N, K=T, **p212) == (block_n, 2), T # 5120 -> 1024 has no two-split range past 8192 tokens on 212 SMs: three chains at 32768 tokens assert cake_backend.k1_chain_plan(M=5120, N=1024, K=32768, **p212) == (256, 3) - # on 216 SMs up to 960 tokens the one-chain program of K1_SPLIT_RULES + # on 216 SMs up to 960 tokens the one-chain program of K1_SPLIT_RULES; past the chain ranges, from 40962 tokens + # up, the deterministic program for (d_in, d_out), choice in K1_SM107A_WIDE_FP32_ROWS.items(): kw = dict(arch="sm_107a", L=1, M=d_in, N=d_out, out_f32=True, sm_count=216) assert cake_backend.k1_chain_plan(K=960, **kw) is None assert cake_backend.k1_plan(K=960, **kw) == choice + for T in (40962, 65536): + assert cake_backend.k1_chain_plan(K=T, **kw) is None + assert cake_backend.k1_plan(K=T, **kw) is None, (d_in, d_out, T) + # at a failing token count the rows keep the routing of the other tables: the plan without the chain ranges + # (216 SMs: the deterministic program, from 961 tokens up) + assert failing + monkeypatch.setattr(cake_backend, "K1_CHAIN_RANGES_BY_SM", {}) + for (key, sms, T), plan in failing.items(): + _, L, M, N, _ = key + want = cake_backend.k1_plan( + arch="sm_107a", L=L, M=M, N=N, K=T, out_f32=True, sm_count=sms + ) + assert plan == want, (key, sms, T) + p216 = dict(arch="sm_107a", L=1, out_f32=True, sm_count=216) + assert failing[(("sm_107a", 1, 4096, 1024, "f32"), 216, 16384)] is None + assert failing[(("sm_107a", 1, 5120, 1024, "f32"), 216, 32768)] is None + for M, N in ((4096, 512), (4096, 1024)): + assert cake_backend.k1_plan(M=M, N=N, K=16231, **p216) is None def test_k1_chains_fit(monkeypatch): @@ -2331,7 +2428,7 @@ def test_k1_chains_fit(monkeypatch): key = ("sm_107a", 1, 4096, 512, "f32") tables = {("sm_107a", 212): {key: ((64, 256, (128, 3)),)}} monkeypatch.setattr(cake_backend, "K1_CHAIN_RANGES_BY_SM", tables) - monkeypatch.setattr(cake_backend, "K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM", {}) + monkeypatch.setattr(cake_backend, "K1_TWO_SPLIT_RANGES_BY_SM", {}) kw = dict(arch="sm_107a", L=1, M=4096, N=512, out_f32=True, sm_count=212) assert cake_backend.k1_chain_plan(K=129, **kw) == (128, 3) assert cake_backend.k1_chain_plan(K=128, **kw) is None @@ -2340,11 +2437,11 @@ def test_k1_chains_fit(monkeypatch): def test_k1_chain_plans_run_the_chain_programs(monkeypatch): """A token count of :data:`K1_CHAIN_RANGES_BY_SM` plans the one-split program of three chains of the range's width (``AtomicGemmPlan.chains``, template ``..._c3``, no workspace and no tickets): the 128-wide and 256-wide - transposed-store programs and the 192-wide row-major-store program, at an even block count inside a 212-SM two-split - range too; past the chain range the two-split range plans two splits of one chain. An explicit ``splits`` / - ``block_n`` runs one chain; ``chains=3`` names the three-chain program of any row at one split; a chain count - outside :data:`K1_CHAINS_CHOICES`, ``chains=3`` with more than one split or on fewer than three K blocks raise - ``ValueError``; a BF16 three-chain program is not registered.""" + transposed-store programs and the 192-wide row-major-store program, inside a 212-SM two-split range too; past the + chain range the two-split range plans two splits of one chain. An explicit ``splits`` / ``block_n`` runs one chain; + ``chains=3`` names the three-chain program of any row at one split; a chain count outside :data:`K1_CHAINS_CHOICES`, + ``chains=3`` with more than one split or on fewer than three K blocks raise ``ValueError``; a BF16 three-chain + program is not registered.""" arch = "sm_107a" _register_k1(monkeypatch, arch) kw = dict(transposed_out=True, arch=arch, sm_count=212) @@ -5232,12 +5329,12 @@ def test_k1_212_sm_range_overrides_bind_per_token_count(monkeypatch): assert len(entry_fn.calls) == len(cases) -def test_k1_212_sm_even_block_ranges_bind_per_token_count(monkeypatch, no_k1_chains): +def test_k1_212_sm_two_split_ranges_bind_per_token_count(monkeypatch, no_k1_chains): """The convenience path plans each token count of a K1 call at the part's SM count, so the ranges of - :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` hold there too: on a 212-SM ``sm_107a`` part the weight gradient of each - row of ``K1_212_EVEN_BLOCK_RANGES`` binds that range's two-split program at its first and last token count, and - the deterministic package's plan one token outside it and one token past its first token count (an odd number of - 64-token K blocks), each token count its own cache entry.""" + :data:`K1_TWO_SPLIT_RANGES_BY_SM` hold there too: on a 212-SM ``sm_107a`` part the weight gradient of each row of + ``K1_212_TWO_SPLIT_RANGES`` binds that range's two-split program at its first and last token count and one token + past its first token count (an odd number of 64-token K blocks), and the deterministic package's plan one token + outside it, each token count its own cache entry.""" arch = "sm_107a" entry_fn = _cpu_k1_launches(monkeypatch, arch) facts = det_backend.DeviceFacts( @@ -5246,14 +5343,15 @@ def test_k1_212_sm_even_block_ranges_bind_per_token_count(monkeypatch, no_k1_cha monkeypatch.setattr(det_backend, "device_facts", lambda device=None: facts) monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) bf16 = torch.bfloat16 - for (d_in, d_out, dt), spans in K1_212_EVEN_BLOCK_RANGES.items(): + for (d_in, d_out, dt), spans in K1_212_TWO_SPLIT_RANGES.items(): cake_backend._K1_ATOMIC_CACHE.clear() entry_fn.calls.clear() out = torch.empty(d_out, d_in, dtype=torch.float32 if dt == "f32" else bf16) cases = [] for lo, hi, choice in spans: assert _k1_blocks(lo + 1) % 2 == 1, (d_in, d_out, lo) - cases += [(lo, tuple(choice)), (hi, tuple(choice)), (lo + 1, None)] + two = tuple(choice) + cases += [(lo, two), (hi, two), (lo + 1, two)] cases += [ (T, None) for T in (lo - 1, hi + 1) @@ -5329,14 +5427,17 @@ def test_k1_exact_sm_count_exclusions_bind_per_token_count(monkeypatch): @pytest.mark.parametrize( - "arch, sm_count", [("sm_100a", 152), ("sm_103a", 152), ("sm_107a", 216)] + "arch, sm_count", [("sm_100a", 152), ("sm_103a", 150), ("sm_107a", 216)] ) def test_k1_odd_block_exclusions_bind_per_token_count(monkeypatch, arch, sm_count): """The convenience path plans each token count of a K1 call at the part's SM count, so the entries of :data:`K1_SPLIT_ODD_BLOCK_EXCLUDED_TOKENS` hold there too: for each FP32 row of the architecture in - ``K1_ODD_BLOCK_EXCLUDED_TOKENS``, the weight gradient binds the row's two-split program one token below the first - token count of its range and 64 tokens above it (even block counts), and the deterministic package's plan at that - first token count and 63 tokens above it (an odd block count), each token count its own cache entry.""" + ``K1_ODD_BLOCK_EXCLUDED_TOKENS`` with an entry on the part, the weight gradient binds the row's two-split program + one token below the first token count of its range (65537, or 1024 for sm_107a ``8192 -> 128``) and one token past + the end of that token count's block (an even block count), and the deterministic package's plan at that first token + count and at the last token count of its block (an odd block count), each token count its own cache entry. (On + 152-SM sm_103a parts these rows run the deterministic program from 32768 tokens up: their entries of + :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS`.)""" entry_fn = _cpu_k1_launches(monkeypatch, arch) facts = det_backend.DeviceFacts( index=0, arch=arch, sm_count=sm_count, l2_bytes=126 << 20 @@ -5344,7 +5445,11 @@ def test_k1_odd_block_exclusions_bind_per_token_count(monkeypatch, arch, sm_coun monkeypatch.setattr(det_backend, "device_facts", lambda device=None: facts) monkeypatch.setattr(cake_backend, "_k1_index", lambda *views: 0) bf16 = torch.bfloat16 - keys = sorted(k for k in K1_ODD_BLOCK_EXCLUDED_TOKENS if k[0] == arch) + keys = sorted( + k + for k, entries in K1_ODD_BLOCK_EXCLUDED_TOKENS.items() + if k[0] == arch and any(sm_count >= e[0] for e in entries) + ) assert keys for key in keys: _, _, d_in, d_out, dt = key @@ -5354,7 +5459,9 @@ def test_k1_odd_block_exclusions_bind_per_token_count(monkeypatch, arch, sm_coun cake_backend._K1_ATOMIC_CACHE.clear() entry_fn.calls.clear() out = torch.empty(d_out, d_in, dtype=torch.float32) - cases = ((lo - 1, rule), (lo, None), (lo + 63, None), (lo + 64, rule)) + last = 64 * _k1_blocks(lo) + assert lo in (1025, 65538) and _k1_blocks(lo) % 2 == 1 + cases = ((lo - 1, rule), (lo, None), (last, None), (last + 1, rule)) for T, choice in cases: G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) assert cake_backend.projection_wgrad(G, X, out, deterministic=False) is out @@ -5375,9 +5482,10 @@ def test_k1_chain_ranges_bind_per_token_count(monkeypatch, sm_count): """The convenience path plans each token count of a K1 call at the part's SM count, so the ranges of :data:`K1_CHAIN_RANGES_BY_SM` hold there too: on a 212-SM and a 216-SM ``sm_107a`` part the weight gradient of each row of ``K1_CHAIN_RANGES`` binds the one-split program of three chains at the first and the last token count of a - range and at an odd block count inside it (an even block count inside the row's 212-SM two-split range too), and the - routing of the other tables at the measured token counts where the program failed, each token count its own cache - entry.""" + range and at 4097 tokens inside it (inside the row's 212-SM two-split range too), and the routing of the other + tables at the measured token counts where the program failed (the deterministic program, or a program of one chain: + the 216-SM override of ``4096 -> 1024`` at 16384 tokens, the one-split program of ``5120 -> 1024`` at 32768 tokens + on 216 SMs), each token count its own cache entry.""" arch = "sm_107a" entry_fn = _cpu_k1_launches(monkeypatch, arch) facts = det_backend.DeviceFacts( @@ -5396,8 +5504,20 @@ def test_k1_chain_ranges_bind_per_token_count(monkeypatch, sm_count): if lo <= T <= hi: cases.append((T, (block_n, 1, chains))) misses = K1_CHAIN_MISSES[sm_count].get((d_in, d_out), ()) - cases += [(T, None) for T in misses] - cases = sorted(set(cases)) + for T in misses: + with monkeypatch.context() as m: + m.setattr(cake_backend, "K1_CHAIN_RANGES_BY_SM", {}) + other = cake_backend.k1_plan( + arch=arch, + L=1, + M=d_in, + N=d_out, + K=T, + out_f32=True, + sm_count=sm_count, + ) + cases.append((T, None if other is None else (other[0], other[1], 1))) + cases = sorted(set(cases), key=lambda c: c[0]) assert any(c is not None and c[2] == 3 for _, c in cases) for T, choice in cases: G, X = torch.empty(T, d_out, dtype=bf16), torch.empty(T, d_in, dtype=bf16) @@ -6929,7 +7049,10 @@ def _k1_wgrad_inputs(row, T, out_dtype, device, seed): def _k1_check(out, ref, what=""): - """The deterministic package's K1 budgets against FP64.""" + """The K1 accuracy rule: the deterministic package's budgets against FP64, absolute (BF16 output: every element + within ``1e-2 + 1e-2 * |ref|``; FP32 output: relative Frobenius error at most 5e-5 and largest absolute error at + most 1e-3 of the largest reference magnitude). The deterministic program's error on the same inputs is data, not + part of the rule.""" assert bool(torch.isfinite(out).all()), what if out.dtype == torch.bfloat16: bad = (out.double() - ref).abs() > BF16_ATOL + BF16_RTOL * ref.abs() @@ -6941,48 +7064,6 @@ def _k1_check(out, ref, what=""): ) <= F32_MAX_ABS_OVER_REF_ABSMAX * float(ref.abs().max()), what -def _k1_bf16_misses(out, ref): - """Elements of a BF16 output outside the deterministic package's BF16 budget against FP64.""" - return int(((out.double() - ref).abs() > BF16_ATOL + BF16_RTOL * ref.abs()).sum()) - - -def _k1_bf16_budget_ratio(out, ref): - """The largest error of a BF16 output against FP64 over its BF16 budget (above 1: outside the budget).""" - err = (out.double() - ref).abs() / (BF16_ATOL + BF16_RTOL * ref.abs()) - return float(err.max()) - - -def _k1_within_budget(out, ref): - """Whether an output meets the deterministic package's K1 budget against FP64 (:func:`_k1_check`).""" - if out.dtype == torch.bfloat16: - return _k1_bf16_misses(out, ref) == 0 - return _rel_fro(out, ref) <= F32_REL_FRO_MAX and float( - (out.double() - ref).abs().max() - ) <= F32_MAX_ABS_OVER_REF_ABSMAX * float(ref.abs().max()) - - -def _k1_check_or_as_det(out, G, X, what=""): - """The K1 accuracy rule against FP64 and the deterministic program on the same inputs: a relative Frobenius error - at most 2 % above that program's, and the deterministic package's budget (:func:`_k1_check`), except where that - program misses the budget too (FP32 accumulation over the largest token counts): then, for a BF16 output, a BF16 - budget ratio (:func:`_k1_bf16_budget_ratio`) at most 30 % above that program's. The budget ratio is a maximum over - the elements: at the same relative Frobenius error, two summation orders move it by up to about 25 % either way - on one input, so the token tables decide that case on the median over ten seeds (README).""" - ref = G.double().t() @ X.double() - det = torch.full_like(out, float("nan")) - det_backend.projection_det_wgrad(G, X, det) - torch.cuda.synchronize() - assert bool(torch.isfinite(out).all()), what - rf, rf_det = _rel_fro(out, ref), _rel_fro(det, ref) - assert rf <= 1.02 * rf_det, (what, rf, rf_det) - if _k1_within_budget(out, ref): - return - assert not _k1_within_budget(det, ref), (what, "outside the budget", rf, rf_det) - if out.dtype == torch.bfloat16: - br, br_det = _k1_bf16_budget_ratio(out, ref), _k1_bf16_budget_ratio(det, ref) - assert br <= 1.30 * br_det, (what, br, br_det) - - def _k1_prepare(G, X, out, arch, splits, block_n=None): """This package's prepared weight gradient at an explicit split count (the transposed-store launch).""" block_n = block_n or cake_backend.k1_block_n(int(G.shape[1])) @@ -7028,7 +7109,7 @@ def force(G, X, out, splits, block_n=None): monkeypatch.setattr(cake_backend, "K1_SPLIT_OVERRIDES_BY_SM", {}) monkeypatch.setattr(cake_backend, "K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM", {}) monkeypatch.setattr(cake_backend, "K1_SPLIT_RANGE_OVERRIDES_BY_SM", {}) - monkeypatch.setattr(cake_backend, "K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM", {}) + monkeypatch.setattr(cake_backend, "K1_TWO_SPLIT_RANGES_BY_SM", {}) monkeypatch.setattr(cake_backend, "K1_CHAIN_RANGES_BY_SM", {}) cake_backend.clear_convenience_caches() return cake_backend.atomic_k1_template(block_n, f32) @@ -7169,10 +7250,9 @@ def test_gpu_k1_row_major_program(shape, transposed_out, k1_rule): @pytest.mark.parametrize("shape", K1_CHAIN_SHAPES, ids=lambda s: "x".join(map(str, s))) def test_gpu_k1_chain_programs(shape, monkeypatch): - """The one-split program of three accumulation chains, prepared with ``chains=3``, meets the deterministic - package's budget against FP64 and, from 961 tokens (where the deterministic program splits K), the K1 accuracy - rule against the deterministic program (:func:`_k1_check_or_as_det`); each launch into a NaN-poisoned ``out`` - repeats bit for bit, also replayed from a CUDA graph; the convenience path at a token count of a range of + """The one-split program of three accumulation chains, prepared with ``chains=3``, meets the K1 accuracy rule (the + deterministic package's budget against FP64, :func:`_k1_check`); each launch into a NaN-poisoned ``out`` repeats bit + for bit, also replayed from a CUDA graph; the convenience path at a token count of a range of :data:`K1_CHAIN_RANGES_BY_SM` for the device's SM count stores the same bits; the profiler shows the three-chain program; fewer than three K blocks raise ``ValueError``.""" device, arch = _atomic_device_or_skip() @@ -7201,8 +7281,6 @@ def test_gpu_k1_chain_programs(shape, monkeypatch): prepared.launch() torch.cuda.synchronize() _k1_check(out, G.double().t() @ X.double(), "prepared") - if T >= 961: - _k1_check_or_as_det(out, G, X, "prepared") first = out.clone() for i in range(3): out.fill_(float("nan")) @@ -7225,7 +7303,7 @@ def test_gpu_k1_chain_programs(shape, monkeypatch): "K1_CHAIN_RANGES_BY_SM", {(arch, sm_count): {key: ((T, T, (block_n, 3)),)}}, ) - monkeypatch.setattr(cake_backend, "K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM", {}) + monkeypatch.setattr(cake_backend, "K1_TWO_SPLIT_RANGES_BY_SM", {}) cake_backend.clear_convenience_caches() try: out.fill_(float("nan")) @@ -8511,10 +8589,11 @@ def test_gpu_planned_rows_run_the_programs_of_their_verdicts(): def test_gpu_k1_sm_count_overrides_run_their_programs(): """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_OVERRIDES_BY_SM` (``sm_107a``: 216 SMs): every overridden row prepares and runs the override's - program at the first and the last token count of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` (``4096 -> 1024`` FP32 at - 16172 and 16384 tokens: the 192-wide row-major-store program at one split), meets the deterministic package's - FP64 budget and, where both run one split, gives the bits of the row's :data:`K1_SPLIT_RULES` program (256-wide - transposed tiles) on the same inputs. A part of any other SM count skips.""" + program at the first and the last token count of :data:`K1_SPLIT_RULES_BY_SM_TOKENS` outside the part's ranges of + :data:`K1_CHAIN_RANGES_BY_SM`, which come first (``4096 -> 1024`` FP32 at 16384 tokens: the 192-wide row-major-store + program at one split), meets the deterministic package's FP64 budget and, where both run one split, gives the bits + of the row's :data:`K1_SPLIT_RULES` program (256-wide transposed tiles) on the same inputs. A part of any other SM + count skips.""" device, arch = _atomic_device_or_skip() sm_count = det_backend.device_facts(device).sm_count table = cake_backend.K1_SPLIT_OVERRIDES_BY_SM.get((arch, sm_count), {}) @@ -8528,9 +8607,11 @@ def test_gpu_k1_sm_count_overrides_run_their_programs(): f32 = dt == "f32" rule = cake_backend.K1_SPLIT_RULES[key] for T in cake_backend.K1_SPLIT_RULES_BY_SM_TOKENS: - want = cake_backend.k1_plan( - arch=arch, L=L, M=M, N=N, K=T, out_f32=f32, sm_count=sm_count - ) + kw = dict(arch=arch, L=L, M=M, N=N, K=T, out_f32=f32, sm_count=sm_count) + # the ranges of K1_CHAIN_RANGES_BY_SM come first + if cake_backend.k1_chain_plan(**kw) is not None: + continue + want = cake_backend.k1_plan(**kw) if want is None: # a token count of K1_SPLIT_RULE_EXCLUDED_TOKENS continue assert want == (block_n, splits), (key, T) @@ -8564,8 +8645,8 @@ def test_gpu_k1_sm_count_overrides_through_the_convenience_binding(): """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_OVERRIDES_BY_SM` (``sm_107a``: 216 SMs): the convenience path of every overridden row binds the override's program (its :data:`_K1_ATOMIC_CACHE` entry) at the first and the last token count of - :data:`K1_SPLIT_RULES_BY_SM_TOKENS` and gives the bits of the prepared launch, the tickets zero after. A part of - any other SM count skips.""" + :data:`K1_SPLIT_RULES_BY_SM_TOKENS` outside the part's ranges of :data:`K1_CHAIN_RANGES_BY_SM` and gives the bits of + the prepared launch, the tickets zero after. A part of any other SM count skips.""" device, arch = _atomic_device_or_skip() sm_count = det_backend.device_facts(device).sm_count table = cake_backend.K1_SPLIT_OVERRIDES_BY_SM.get((arch, sm_count), {}) @@ -8578,9 +8659,11 @@ def test_gpu_k1_sm_count_overrides_through_the_convenience_binding(): _, L, M, N, dt = key f32 = dt == "f32" for T in cake_backend.K1_SPLIT_RULES_BY_SM_TOKENS: - plan = cake_backend.k1_plan( - arch=arch, L=L, M=M, N=N, K=T, out_f32=f32, sm_count=sm_count - ) + kw = dict(arch=arch, L=L, M=M, N=N, K=T, out_f32=f32, sm_count=sm_count) + # the ranges of K1_CHAIN_RANGES_BY_SM come first + if cake_backend.k1_chain_plan(**kw) is not None: + continue + plan = cake_backend.k1_plan(**kw) if plan != (block_n, splits): # an excluded token count continue G, X = _bf16((T, N), device, g), _bf16((T, M), device, g) @@ -8610,30 +8693,26 @@ def test_gpu_k1_sm_count_overrides_through_the_convenience_binding(): def _k1_token_range_edges(arch, sm_count): """``(key, choice, T, want)`` for every row of the part's ranges in :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, - :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` and :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM`: the first and the last token - count of each range (``want``: the program of the part's table, or of the override or two-split range) and one - token outside it where neither the window nor another range of the row holds (``want`` ``None``: the - deterministic program), and one token past the first token count of a two-split range (an odd number of 64-token K - blocks: ``want`` ``None``); ``choice`` is the planner's.""" + :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` and :data:`K1_TWO_SPLIT_RANGES_BY_SM`: the first and the last token count of + each range (``want``: the program of the part's table, or of the override or two-split range), one token past the + first token count of a two-split range (an odd number of 64-token K blocks; ``want``: that range's program), and one + token outside a range where neither the window nor another range of the row holds (``want`` ``None``: the + deterministic program); ``choice`` is the planner's.""" first, last = cake_backend.K1_SPLIT_RULES_BY_SM_TOKENS part = (arch, sm_count) rows = cake_backend.K1_SPLIT_RULES_BY_SM_TOKEN_RANGES.get(part, {}) over = cake_backend.K1_SPLIT_RANGE_OVERRIDES_BY_SM.get(part, {}) - even = cake_backend.K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM.get(part, {}) + two = cake_backend.K1_TWO_SPLIT_RANGES_BY_SM.get(part, {}) table = cake_backend.K1_SPLIT_RULES_BY_SM.get(part, {}) cases = [] - for key in sorted(set(rows) | set(over) | set(even)): + for key in sorted(set(rows) | set(over) | set(two)): _, L, M, N, dt = key spans = [(lo, hi, table[key], False) for lo, hi in rows.get(key, ())] spans += [(lo, hi, v, False) for lo, hi, v in over.get(key, ())] - spans += [(lo, hi, v, True) for lo, hi, v in even.get(key, ())] - for lo, hi, _, parity in spans: - for T in sorted({lo, hi, lo - 1, hi + 1} | ({lo + 1} if parity else set())): - held = [ - v - for a, b, v, par in spans - if a <= T <= b and not (par and _k1_blocks(T) % 2) - ] + spans += [(lo, hi, v, True) for lo, hi, v in two.get(key, ())] + for lo, hi, _, split in spans: + for T in sorted({lo, hi, lo - 1, hi + 1} | ({lo + 1} if split else set())): + held = [v for a, b, v, _ in spans if a <= T <= b] if not held and first <= T <= last: continue plan = cake_backend.k1_plan( @@ -8692,20 +8771,35 @@ def _k1_with_chain_ranges(arch, sm_count, key, T, want): return (chain[0], 1), chain[1] +#: Edge token counts of the GPU edge test at which the deterministic program itself misses the BF16 budget for some +#: seed (measured, 10 seeds), so the budget does not choose the route there: ``(arch, sm_count, key, T) -> T'``. The +#: edge asserts the route only, and the accuracy rule runs at ``T'``, the nearest measured token count of the same +#: range at which the deterministic program met the budget on every seed on the measured parts (``None``: none). +K1_P1_EDGES = { + ("sm_100a", 148, ("sm_100a", 1, 4096, 64, "bf16"), 65535): 41985, + ("sm_103a", 148, ("sm_103a", 1, 4096, 32, "bf16"), 65535): 45057, + ("sm_103a", 148, ("sm_103a", 1, 5120, 48, "bf16"), 49152): None, + ("sm_103a", 148, ("sm_103a", 1, 5120, 64, "bf16"), 47104): 45057, + ("sm_103a", 148, ("sm_103a", 1, 5120, 64, "bf16"), 49153): None, + ("sm_103a", 148, ("sm_103a", 1, 5120, 128, "bf16"), 65537): None, +} + + def test_gpu_k1_sm_count_token_ranges_run_at_their_edges(): """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` or - :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` (``sm_107a``: 212 SMs; ``sm_100a`` / ``sm_103a``: 152 SMs): at the first - and the last token count of every - range of a row the prepared launch runs the program of the part's table in :data:`K1_SPLIT_RULES_BY_SM` (of the - range, for an override or two-split range) and meets the deterministic package's FP64 budget - (:func:`_k1_check_or_as_det`); one token outside a range (and outside the window), and one token past the first - token count of a two-split range (an odd number of K blocks), the row runs the deterministic program, with the - deterministic package's bits. On a part of exactly the SM count of a table of - :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM` (148-SM ``sm_100a`` / ``sm_103a`` parts) the reverse: the deterministic - program with its bits at both ends of a range, the row's program one token outside it. An edge inside a range of - :data:`K1_CHAIN_RANGES_BY_SM` for the part runs that range's program of several chains instead - (:func:`_k1_with_chain_ranges`), under the same FP64 rule. A part of any other SM count skips.""" + :data:`K1_TWO_SPLIT_RANGES_BY_SM` (``sm_107a``: 212 SMs; ``sm_100a`` / ``sm_103a``: 152 SMs): at the first and the + last token count of every range of a row, and one token past the first token count of a two-split range (an odd + number of K blocks), the prepared launch runs the program of the part's table in :data:`K1_SPLIT_RULES_BY_SM` (of + the range, for an override or two-split range) and meets the K1 accuracy rule (:func:`_k1_check`); one token outside + a range (and outside the window) the row runs the deterministic program, with the deterministic package's bits. On + a part of exactly the SM count of a table of :data:`K1_SPLIT_RULE_EXCLUDED_TOKENS_BY_SM` (148-SM ``sm_100a`` / + ``sm_103a`` parts) the reverse: the deterministic program with its bits at both ends of a range, the row's program + one token outside it. An edge inside a range of :data:`K1_CHAIN_RANGES_BY_SM` for the part runs that range's + program of several chains instead (:func:`_k1_with_chain_ranges`), under the same FP64 rule. At an edge of + ``K1_P1_EDGES`` the deterministic program itself misses the BF16 budget for some seed: the edge asserts the route + only, and the FP64 rule runs at the entry's token count of the same range instead (none: no check). A part of any + other SM count skips.""" device, arch = _atomic_device_or_skip() sm_count = det_backend.device_facts(device).sm_count cases = _k1_token_range_edges(arch, sm_count) @@ -8732,8 +8826,21 @@ def test_gpu_k1_sm_count_token_ranges_run_at_their_edges(): ), (key, T) _assert_ran(prepared.launch, prepared.module_name) torch.cuda.synchronize() - _k1_check_or_as_det(out, G, X, f"{M} {N} {dt} {T}") ran += 1 + at = K1_P1_EDGES.get((arch, sm_count, key, T), T) + if at is None: + continue + if at != T: + G, X = _bf16((at, N), device, g), _bf16((at, M), device, g) + out = torch.full((N, M), float("nan"), device=device, dtype=out_dtype) + prepared = cake_backend.prepare_projection_wgrad( + G, X, out, deterministic=False + ) + assert isinstance(prepared, cake_backend.PreparedAtomicGemm), (key, at) + assert (prepared.plan.block_n, prepared.plan.splits) == plan, (key, at) + _assert_ran(prepared.launch, prepared.module_name) + torch.cuda.synchronize() + _k1_check(out, G.double().t() @ X.double(), f"{M} {N} {dt} {at}") else: assert plan is None, (key, T) assert not isinstance(prepared, cake_backend.PreparedAtomicGemm), (key, T) @@ -8751,7 +8858,7 @@ def test_gpu_k1_sm_count_token_ranges_run_at_their_edges(): def test_gpu_k1_sm_count_token_ranges_through_the_convenience_binding(): """Unforced (``deterministic=False``, no split count), on a part of exactly the SM count of a table of :data:`K1_SPLIT_RULES_BY_SM_TOKEN_RANGES`, :data:`K1_SPLIT_RANGE_OVERRIDES_BY_SM` or - :data:`K1_SPLIT_EVEN_BLOCK_RANGES_BY_SM` (``sm_107a``: 212 SMs; ``sm_100a`` / ``sm_103a``: 152 SMs): the + :data:`K1_TWO_SPLIT_RANGES_BY_SM` (``sm_107a``: 212 SMs; ``sm_100a`` / ``sm_103a``: 152 SMs): the convenience path binds the part's program (the range's, for an override or two-split range) at the first and the last token count of every range of a row (its :data:`_K1_ATOMIC_CACHE` entry) and gives the bits of the prepared launch where at most two partials meet (the