Repository navigation
[Perf] [MI355X] glm53: Select a long row's indexer top-k on up to 16 CTAs - #24
Merged
Merged
Conversation
jhinpan
requested review from
HaiShaw,
alexnails and
kevin-mii
as code owners
October 22, 2026 02:04
4 of 5 tasks
jhinpan
force-pushed
the
glm53-parallel-topk
branch
from
October 23, 2026 22:16
4f187f6 to
48ff491
Compare
On an indexer layer, the exact top-2048 of each row ran on one CTA, which scanned all of the row's scores at least twice. Its time grew with the context at the speed of one CU and every later stage of the layer waited for it: past 128k keys at batch 4 or more, the decode step was slower than SGLang's path. A row of more than 8 chunks (32768 keys) now selects on up to 16 CTAs of the grid's upper half, each a run of whole chunks. Each histograms its chunks; the row's index_select CTA merges the histograms and publishes the threshold bin; the slices then collect their winners and candidates in position order, and index_select concatenates them in slice order into the lists the one-CTA path builds, before the same radix finish. Shorter rows keep the one-CTA path. The score stage does not change.
index_thresh and index_select read the slices' histograms and staged lists, plain stores released before their ready pairs, behind an agent-scope acquire now, as the indices hand-off does. A parallel row whose threshold bin overflows polls its scores in the overflow pass: its CTA has not read them. Check mode compares each indexer row's selection with the exact top-k of the launch's own scores (layer.topk_exact, bound 0). The parallel check case runs rows of about 108k keys at widths 2 and 4, so slices hold two chunks.
jhinpan
force-pushed
the
glm53-parallel-topk
branch
from
October 25, 2026 03:19
48ff491 to
235bb05
Compare
jhinpan
added a commit
that referenced
this pull request
Oct 25, 2026
## Motivation **Summary.** After #24, the score stage of the GLM-5.3 indexer read the keys at about 0.7 TB/s, and at 5 to 8 rows past 128k keys it took about half of an indexer launch. This PR scores each group of 16 keys on one wave over all 32 heads, with the reads of the next groups in flight. The scores do not change: in one launch, they are bitwise equal to the scores of #24 on the same inputs. On 8x MI355X at TP8, in InferenceX agentic coding with MTP, interactivity p50 is 1.43x, 1.46x and 1.23x that of the normal path with MTP at concurrency 1, 2 and 4 (#24: 1.15x, 1.17x and 1.11x). **On main after #24.** This PR is two commits on main (`240d587`, which has #24 and #56): `e405d6b` is `e0d7655` rebased with the same changed lines (and `e0d7655` is `9733ff4` rebased), and `3419a0e` is the fix from the review. The results below were measured at `9733ff4` unless marked. **Cause.** In the paged score stage of #24: - A CTA scored a tile of 64 keys with its 8 waves: 4 groups of 16 keys, each for the two halves of the 32 heads. The two waves of a group read the same keys, so each key was read two times. - The two halves were added through LDS, with two block barriers a tile. - The block-table read and the key reads of a tile started only after the tile before it ended. Thus each tile waited for two memory round trips: about 2.1 us a tile, 0.7 TB/s at 256k keys. - Each key also compared its cache slot with the slots of all S rows of the launch (the new-key lookup of #15). This cost grows with the width (S = 2, 4 or 8). ## Modifications - `glm53/gfx950/kernel.py`, paged score stage only: - A wave scores a group of 16 keys over all 32 heads: two MFMA chains on one read of the keys. The head sums stay in the lanes, so the stage has no LDS reduction and no barrier between groups. - A wave walks every 8th group of its CTA's tiles; the tiles of a CTA do not change. The key reads of the next 2 groups are in flight, and their block-table reads go one step further ahead. The first block-table reads of a row overlap the staging of its query. - The query fragments and the head weights of the row stay in registers. The values of a row (bound, request, counts) are wave-uniform. - The new-key lookup of #15 runs only for each wave's group that holds any of the last S keys of a row: a launch writes only those keys (a verify row's draft chain, a decode row's own key). A wave's groups are `SCORE_STRIDE` (32768) keys apart, so a wave has at most one such group; the walk asserts `S + 15 < SCORE_STRIDE` (review fix). - Each score is computed in the same order as before: the same MFMA chain, the same sums in the same order. - `test/e2e/test_glm53_server.py` (review fix): the parallel check case of #24 passes 98304 keys (rows of about 108k keys at widths 2 and 4), where each wave's walk refills its buffers and looks blocks ahead. Before, the longest case (about 36k keys) gave each wave at most 2 groups, so the refills never ran. - No new check. With a wrong refill offset (a deliberate error, not in this PR: a refilled buffer reads the keys 16 positions on), the parallel check case fails on the existing stream-cosine bound: 22 violations, lowest cosine 0.914 (bound 0.95). `layer.topk.rel` rises to 0.66 (0.052 without the error). So once a case reaches the walk's refills, the existing bounds catch an error there. The bitwise check needs a debug build (not in this PR). ## Accuracy Tests **Setup.** 8x MI355X (gfx950) at TP8; zai-org/GLM-5.3 @ `ba32f8c1` (FP8); the image and the commits of #24; the server flags of `test/e2e/glm53_server.py`; a new server for each test. **Test 1: bitwise scores in one launch.** A debug build (not in this PR) also runs the per-tile score code of #24 in the same launch, on the same inputs, and writes those scores to a second scratch buffer. A script compares the two buffers. | Steps | Rows | Keys | Keys bitwise equal | |---|---|---|---| | Decode: 1 to 8 rows, 8k to 256k keys | 256 | 29,961,984 | 29,961,984 | | MTP verify: 1 and 2 requests, 8k to 256k keys | 192 | 29,631,744 | 29,631,744 | The exact top-k check of #24 also found 0 wrong rows in both runs. We did not compare two servers: on this image, two prefills of the same prompt give different bits. After a prefill of 131072 tokens, the scores of the first decode step correlate at 0.9988 between two passes (one prefill, cached) and at 0.9457 between two prefills in the same server, with the same build. Between a server with #24 and a server with this PR, they correlate at 0.9460. **Test 2: E2E tests at TP8 (review fix).** We ran `test/e2e/test_glm53_server.py` as written at `3419a0e` on main `240d587`. All 7 tests passed in one run (1214 s): - `TestMonokernelDecode`: average KL 0.018 to 0.022 for each batch (bound 0.06); GSM8K (sgl-eval, thinking on) 0.99. - `TestMonokernelMtp`: GSM8K 0.99, accept length 3.28. - `TestMonokernelCheck`: no violations. `layer.topk_exact` (#24) and `layer.topk_own` are 0 in all 954 indexer checks. - Decode steps: 1425 checks, stream cosine 0.9942 or more, `layer.topk.rel` 0.014 or less. - Verify steps: 1350 checks, stream cosine 0.9976 or more, `layer.topk.rel` 0.017 or less. - Past the parallel top-k, with rows past 98304 keys: 1200 checks, stream cosine 0.9968 or more, `layer.topk.rel` 0.052 or less. - `TestMonokernelFault`: passed. Since #55, the test sets `HSA_DISABLE_COREDUMP_ON_EXCEPTION=1` and turns off SGLang's crash dumps itself. Two earlier runs also passed all 7 tests: at `9733ff4`, before the review fix, and at `9b48912`, the fix on #24's branch before #56 merged. ## Speed Tests and Profiling **Setup.** The setup of #24: one node with 8x MI355X, TP8, image `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261019`, SGLang `e45cbce599`, AITER `2bb549102b`, the InferenceX GLM recipe flags, and a new server for each point on an idle node. The arms differ only in `--monokernel-backend amd` and the checkout. **Kernel timeline: #24 to this PR.** The debug build of #24, random prompts, 1 to 8 rows decoding together, the last launch of each of 6 passes on each of the 8 ranks (medians). The estimated step is 18 indexer launches and 57 launches without the indexer. | Keys per row | Rows | Score stage (us) | Indexer launch (us) | Estimated step (ms) | |---|---|---|---|---| | 8k | 1 | 5 -> 5 (0.9x) | 97 -> 96 | 4.84 -> 4.80 (1.01x) | | 8k | 2 | 12 -> 9 (1.3x) | 108 -> 106 | 4.98 -> 4.96 (1.00x) | | 8k | 4 | 18 -> 14 (1.3x) | 136 -> 134 | 6.61 -> 6.59 (1.00x) | | 8k | 8 | 29 -> 24 (1.2x) | 206 -> 202 | 10.65 -> 10.61 (1.00x) | | 32k | 1 | 11 -> 5 (2.0x) | 119 -> 114 | 5.20 -> 5.12 (1.01x) | | 32k | 4 | 45 -> 29 (1.5x) | 182 -> 172 | 7.45 -> 7.26 (1.03x) | | 32k | 8 | 89 -> 62 (1.4x) | 285 -> 261 | 12.10 -> 11.67 (1.04x) | | 64k | 1 | 15 -> 6 (2.5x) | 129 -> 133 | 5.33 -> 5.49 (0.97x) | | 64k | 4 | 63 -> 35 (1.8x) | 206 -> 182 | 7.87 -> 7.47 (**1.05x**) | | 64k | 8 | 126 -> 74 (1.7x) | 334 -> 281 | 12.95 -> 11.95 (**1.08x**) | | 128k | 1 | 25 -> 8 (3.1x) | 144 -> 126 | 5.60 -> 5.32 (**1.05x**) | | 128k | 2 | 50 -> 19 (2.7x) | 174 -> 144 | 6.19 -> 5.59 (**1.11x**) | | 128k | 4 | 102 -> 41 (2.5x) | 251 -> 191 | 8.71 -> 7.58 (**1.15x**) | | 128k | 8 | 203 -> 87 (2.3x) | 417 -> 306 | 14.56 -> 12.44 (**1.17x**) | | 256k | 1 | 44 -> 12 (3.7x) | 175 -> 148 | 6.20 -> 5.73 (**1.08x**) | | 256k | 2 | 88 -> 27 (3.2x) | 224 -> 166 | 7.05 -> 6.05 (**1.17x**) | | 256k | 4 | 180 -> 60 (3.0x) | 343 -> 225 | 10.37 -> 8.33 (**1.25x**) | | 256k | 5 | 227 -> 78 (2.9x) | 453 -> 305 | 14.94 -> 12.30 (**1.21x**) | The launches without the indexer run the same code in both builds; between the two runs they differ by -1.8% to +3.4%. The differences at 8k keys and at 64k keys with 1 row are in this noise. The score stage's time is its window, from the first start to the last end of its tasks. With this PR, a tile's start and end stamps can come from two waves, so the time of one task is not meaningful, but the window is. **AgentX with MTP.** InferenceX agentic coding (`agentx-harness` 1.0.11, corpus `23f152f6`), 3600 s a point, the same simulated accept length (2.99) in all arms. | Concurrency | Interactivity p50: normal path with MTP / #15 / #24 / this PR | p90 | Total throughput (tok/s/GPU) | Requests done | Input p50 | |---|---|---|---|---|---| | 1 | 204.5 / 0.95x / 1.15x / **292.1 (1.43x)** | 178.5 / 0.94x / 1.14x / **216.1 (1.21x)** | 1921 / 0.98x / 1.11x / **2386 (1.24x)** | 234 / 231 / 246 / 261 | 225k / 221k / 285k / 300k | | 2 | 199.7 / 0.95x / 1.17x / **292.0 (1.46x)** | 149.6 / 0.91x / 1.03x / **184.9 (1.24x)** | 2044 / 0.99x / 1.15x / **2418 (1.18x)** | 377 / 374 / 409 / 427 | 97k / 96k / 111k / 116k | | 4 | 180.7 / 0.86x / 1.11x / **222.6 (1.23x)** | 128.6 / 0.83x / 0.96x / **146.3 (1.14x)** | 2961 / 0.97x / 1.06x / **3657 (1.24x)** | 501 / 493 / 536 / 582 | 91k / 91k / 89k / 112k | Time of a verify step for each generated token (ms), from the server log intervals without prefill, concurrency 1 and 2 (normal path with MTP / #15 / #24 / this PR): | Requests in the step (verify rows) | 0-64k | 64-128k | 128-192k | 192-320k | 320k or more | |---|---|---|---|---|---| | 1 (4) | 4.61 / 3.21 / 3.15 / **3.06** | 4.66 / 3.50 / 3.29 / **3.12** | 4.74 / 3.93 / 3.49 / **3.21** | 4.86 / 4.92 / 4.00 / **3.33** | 4.95 / 5.65 / 4.39 / **3.48** | | 2 (8) | 5.51 / 4.92 / 4.79 / **4.61** | 5.57 / 5.39 / 5.02 / **4.70** | 5.61 / 5.73 / 5.31 / **4.85** | 5.82 / 7.51 / 6.20 / **5.04** | too few steps | With 2 requests past 192k keys, #24 was still slower than the normal path; this PR is 1.15x faster there. **Fixed length without MTP.** Prompts of 131072 and 262144 tokens, 1024 new tokens, 6 prompts for each concurrent request, the radix cache flushed before each concurrency. Decode step without prefill (ms, median of the server log intervals): | Keys | Batch | Normal path | #24 | This PR | |---|---|---|---|---| | 128k | 1 | 11.04 | 6.47 | **6.25** | | 128k | 4 | 12.94 | 9.28 | **8.16** | | 128k | 8 | 15.52 | 14.78 | **12.52** | | 256k | 1 | 11.36 | 7.01 | **6.49** | | 256k | 4 | 13.55 | 10.93 | **8.75** | | 256k | 5 | 15.92 | 15.40 | **12.61** | At 256k keys, the #24 column is the same selection code on `e952b22` (without #15); #24 itself was measured at 128k. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. ## Review and Merge Process 1. The `lint` and `Check PR title` checks must pass. 2. Ask a code owner from `.github/CODEOWNERS` to review. 3. Merge with "Squash and merge": the PR title becomes the commit title on `main`.
This was referenced Oct 28, 2026
jhinpan
pushed a commit
that referenced
this pull request
Oct 29, 2026
… exact top-2048 on up to 64 slices (#74) ## Motivation **Stacked on #68, #69 and #73** (`[Feat] [MI355X] glm53: Cover layers 0-2 and the logits`). Merge those first. This PR's diff is its last four commits, and it uses #68's `test/kernels/glm53/` tools. On an indexer launch the attention waits for the index chain: index query, key scoring, then the exact top-2048. That covers 18 of the 75 MoE layer launches of a step, plus layer 0 and the dense layers 1–2, which reuse the indexer stages. Two parts of the chain grow with context: - **Scoring reads each key once per row.** #42's per-wave scoring scores one row at a time, so at MTP's 4 verify rows every key of the request is read 4 times. - **The selection is mostly serial.** A row of up to 32,768 keys selects on one CTA. Past that, #24 spreads it over at most 16 slices, and the row's CTA adds their 4,096-bin histograms and, once the slices have gathered, copies their lists. This PR scores a request's rows together, reading each key once from LDS batches of whole pages. It then selects every row's top-2048 on up to 64 slices with all-gathered 6-bit digit histograms. The selection stays exact: the same keys, in the same order and with the same ties, as the base's radix select. Every launch's outputs are bitwise equal to the base's (Accuracy Tests). Result on MI355X, base (#73) → this PR: - TP8, MTP 3-1-4 at the golden accept length 2.99, cycle at one request: - 8K context: 7.16 → 6.92 ms (−3.4%); - 128K context: 8.29 → 7.24 ms (−12.7%). - At two requests, while both decode: - 8K: 10.76 → 10.54 ms (−2.1%); - 128K: 12.66 → 11.48 ms (−9.3%). - TP8, MTP off, median TPOT at 8K: 4.45 → 4.42 ms (−0.5%) at one request, 4.56 → 4.44 (−2.8%) at two, 5.87 → 5.77 (−1.7%) at four. - Kernel-only, `time_layers.py`, E4M3 MLA cache, TP8, µs a launch at 4 rows: - indexer launch: 96.42 → 86.33 at 8K (−10.5%) and 140.00 → 100.00 at 128K (−28.6%); - the 75-launch model order: 71.12 → 68.78 (−3.3%) and 82.88 → 72.20 (−12.9%). The gain grows with context because both rewritten stages scale with the keys a row reads. ## Modifications The first commit adds a check. The next three are the changes, each measured on top of the one before (Speed Tests, "Each change on its own"). ### 0. An exact top-k check (`test/kernels/glm53/check_topk.py`, `check_layer.py`, `test_glm53_kernels.py`) - `check_topk.py` runs indexer launches under torchrun at TP4 or TP8, at chosen positions (`--pos`). `--kind dense` and `--kind embed` run layer 0 and the dense layers. - It uses four index caches: - `check_layer.py`'s random keys; - four keys repeated over every position, so the 2048th score falls in a tie (`ties`); - all-zero keys (`zeros`); - each row's last 4,096 keys scaled 8x, so the selection crowds into them (`recent`). - Row sets: - the served decode points (`serving`, as `time_layers.py`); - one request a row, at `--pos` and below, some under 2,048 keys (`mixed`); - one request's rows, a draft chain as long as the width (`chain`). - Before each step, the cached scales of the step's own keys are poisoned (`check.poison_new_index_scales`). A launch that scores a key this launch writes from the cache therefore fails. - After each launch it checks every live row: - **Scores** (`check.launch_scores`) against an f64 reference. The reference uses the launch's index query (its `index_q` mailbox, rotated and transformed in f32 as the kernel does, then bf16), its head weights, and the row's keys as the launch leaves the index cache. The error is taken over the row's largest reference score; the bound is 1e-4. - **The list** (the top-k at the row's `indptr`), entry by entry, against the radix select's list of the launch's own scores. That list is the keys above the 2048th key's 12-, 24- or 32-bit prefix in position order, then the prefix's keys above the 2048th, then its ties in position order. `missed` also counts the keys of the exact top-2048 (score descending, then position) that the list lacks, as `check.py`'s `layer.topk_exact` does. - **Negative controls** on copies of each row's result: a selected slot replaced by the best unselected key, two entries swapped, and one score's bits flipped must each fail. - `--dump DIR` writes SHA-256 digests of every row's scores and list in `check_layer.py`'s format, so `check_layer.py --compare` can match two trees. - `test_glm53_kernels.py` runs it at TP4 (widths 2, 4) and TP8 (2, 4, 8): - the indexer launch at about 8K and 300K keys (at 300K a scoring CTA's run takes three batches and a slice two LDS chunks); - layer 0 and a dense layer at 8K and 128K. - It passes on the base too: it checks what this PR keeps, the base's list. - **Test fix, in the same commit:** `check_layer.py`'s steps gave each row its request's `req_ids` entry, while the page table has a row per step row. At 8 verify rows, the second request's rows read the first request's page-table row. Each row now gets its own entry, as the spec passes them (`req_ids` = row index, `page_table` = the DSA backend's row per step row). ### 1. Each row's index key on its own task, and the index-query tiles first (`kernel.py`, `layout.py`) - **Index key:** the cache task used to normalize, rope and store a row's index key after the row's KV row. The key is now an `index_key` task of its own, on the row's q_norm CTA. - **q_norm CTAs:** they run their own index-query tile right after q_norm, ahead of the index key and the KV row. - This replaces #68's change 2, where those tiles ran second, on the CTAs of the first extra tiles. - That change's shuffle and its TP4 guard go. - **q_b CTAs:** a q_b CTA issues its index-query tile's weights with q_b's, before the wait, and runs that tile before its q_b tile. - **Values:** the same values, computed in the same order. - **Unpaged fused indexer:** it goes. Every indexer build reads SGLang's paged index cache (`runner.py` passes `index_paged=key.with_indexer`). - **On its own this change costs 0.6 to 3.3 µs** (Speed Tests). It stays because changes 2 and 3 were built and measured on it: the scoring patches this launch's keys from the `index_key` task's pairs. ### 2. A request's rows scored together, each key read once (`kernel.py`, `layout.py`) - **Before (#42):** a CTA scored tiles of 64 keys of one row, a wave every 8th group of 16 keys. A request's 4 verify rows read each key 4 times, and every scoring CTA polled, roped and Hadamard-transformed its row's whole query. - **Groups of rows:** up to 4 consecutive rows of one request score its keys in one pass. - A row belongs to the previous row's request when its page table maps that row's position to the slot that row writes; a page being written belongs to its request alone. - Each row's 2 waves (4 rows) or 4 waves (2 rows) hold its query and head weights in registers. - **Runs of pages:** the request's pages are split into runs of at least 2 pages, one run a CTA, up to 256 runs. They are counted from a CTA that the group's first row offsets. - A CTA copies its run into LDS 8 pages at a time with 16-byte loads. The next batch's loads are in flight while its waves score the current one. - A run's block ids load with the query's poll. - **Keys this launch writes:** - Each `index_key` task also publishes its row's key as cached, FP8 value words and scale, in `index_k8` pairs. A batch holding any of those keys patches them in from the pairs, so the cache store need not be visible yet. - The patch covers every row of the request in this launch, not only the group's. A draft chain of more than 4 rows (MTP with 4 or more steps) spans two groups, and the second group reads the first group's new keys. - The base handles this case by comparing every key's slot with every row's slot (#15, #42). Here, `check_topk.py`'s `chain` rowset at 8 rows checks it. With each group patching only its own rows, that check fails (Accuracy Tests). - **The query once a row:** `index_qh` tasks rope and Hadamard-transform each row's index query once and publish it in bf16; the scoring CTAs poll it into LDS. - A task covers a row's 8 heads, a wave a head. - At TP4, q_b runs on every CTA and produces every index-query tile, so these tasks run after q_b. - **Scores:** every score is computed as before, with the same MFMA chains over the same bf16 query and keys and the same head sums in the same order. The scores are bitwise equal to the base's (Accuracy Tests). - **Removed:** the `index_ready` flag and the `index_k_new` pairs. ### 3. Every row's exact top-2048 on up to 64 slices (`kernel.py`, `layout.py`) - **Before (#24):** a row of up to 32,768 keys selected on one CTA, reading every score at least twice. Past that, up to 16 slices built 4,096-bin histograms of their chunks; one CTA added them and published the threshold bin; the slices gathered their winners and candidates; and the row's CTA copied them and finished the radix select. - **Now** every row selects on up to 64 slices, a slice per 1,024 keys. `layout.select_slices` gives the split, and the kernel and the CPU test both call it. Each slice is a run of the row's keys in position order, held as radix keys in LDS. - **Digit histograms:** a radix select by 6-bit digits from the top. For each digit, every slice publishes the 64-bin histogram of its keys under the prefix found so far (32 tagged pairs of 16-bit counts) and reads every other slice's. It then finds the digit holding the remaining k-th key and counts the keys above it in the slices before its own. - **Placing and resolving:** the prefix stops at 12 or 24 bits once its keys fit the candidate room (as the base's does), or runs to all 32 bits. Each slice then places its keys above the prefix, whose places are known, and stages its keys with the prefix, with their cache slots, for the row's first slice. That slice resolves them in two 10-bit passes, or one 8-bit pass at 24 bits. At 32 bits the ties take their places directly. - **Order:** the list is the base's one-CTA radix select's, entry by entry. That is the keys above the threshold's 12-, 24- or 32-bit prefix in position order, then the prefix's keys above the threshold, then its ties in position order. - **Placement:** the slices start past uk's CTAs. - When every row has at most 256 / S slices, slice c of row s is task c·S + s, so short rows' slices share one range of CTAs. - Otherwise each row takes a run of 64 tasks. - **Scratch:** the histogram, threshold and gather stages' buffers give way to `sel_hist` and `sel_cand`. The indexer runtimes of widths 2, 4 and 8 together go from 163.4 to 135.2 MiB a rank at TP8, and from 76.5 to 64.4 MiB at TP4 (widths 2, 4). - **Tests, in the same commit:** - `test/unit/test_glm53_select.py` covers `select_slices`: the slices tile a row's keys once, in position order, none empty, at most 64, a slice's keys within its LDS room and its 16-bit counts. It also covers the stage order. - `test_glm53_server.py`'s parallel-top-k case asserts that its rows reach all 64 slices. ## Accuracy Tests - **Exact top-2048** (`check_topk.py`, E4M3 MLA cache, 2 indexer launches × 2 steps). - Coverage: positions 8,447, 131,071, 262,143 and 524,287, the four index caches, and the serving, mixed and chain row sets. - TP8, widths 2, 4, 8: 288 checks on this head and 288 on the base, every row exact (missed 0, order exact). The largest score error is 1.8e-6, and every negative control is caught. - Head vs base: the digests of every row's scores and list are equal (41,984 digests, 0 differ). The scores are bitwise the base's, and so is the list, order and ties included. - TP4 (widths 2, 4) on this head: every row exact. - Negative control for the new-key patch: with each scoring group patching only its own rows, the 8-row `chain` set fails at 8K and 128K on every index cache. The launch scores the earlier rows' new keys from their poisoned cached scales, giving a score error of about 1e30. - The kernel tests (below) also run it on layer 0 and a dense layer at 8K and 128K, at TP4 and TP8. - **Bitwise** (`check_layer.py --dump`, `--compare`). - Coverage: kinds `shared`, `index`, `model` (the first 8 launches), `dense` and `full` (layer 0, the dense layers 1–2 and the tail); bf16 and E4M3 MLA cache; every width the TP builds; one request a row, the same with a pad row, and MTP's verify rows; 3 steps. - TP8: 63,648 digests, all equal. TP4: 16,416 digests, all equal except one run of one configuration (bf16 dense, 2 rows with a pad row: 127 digests, the first dense launch's output with its top-k equal). That configuration, rerun 4 times, was equal base vs head and run to run. The negative control (one input value negated) changes 7,056 of 10,944 digests. - Every output and cache row is therefore bitwise equal to the base's, and so is the top-k, at every TP, width and cache dtype the spec builds. - These dumps used `check_layer.py` with its runtimes built before the layers' weights, as #69's tool did. With the base's own `check_layer.py`, which since #73 draws every layer's weights first, two TP8 width-2 configurations (E4M3 index with verify rows, bf16 model with a pad row) vary from run to run on the base itself, with the same kernel code objects. This is under investigation separately and is not from this PR. - **What can differ from SGLang:** nothing new. The serving check's `layer.topk_own` compares with SGLang's selection from FP8-q scores, while the kernel scores the bf16 query. At a threshold tie, the two can pick a different key, as on the base. - **`SGLANG_AMD_MONO_CHECK=1`:** - The e2e file's four check cases pass on this head, each asserting no violation: - decode rows 1 / 3 / 8; - MTP verify rows 4 / 8; - verify rows on the E4M3 MLA cache; - rows past the parallel top-k (about 108K keys at widths 2 and 4, on 64 slices; about 36K at width 8, on 36). - Past 128K keys (two requests of 140K tokens, E4M3 KV, eager), on this head, 0 violations: - decode at 1 and 2 rows: 750 layer checks, 54 of them indexer checks past 131,072 keys; - MTP verify at 4 and 8 rows: 525 checks, 36 past 131,072 keys. - `layer.topk_exact` and `layer.topk_own` are 0 in every indexer check; `layer.topk.rel` is at most 0.048 and `layer.index.cosine` at least 0.99997. - **E2E** (`test/e2e/test_glm53_server.py`, TP8, on this head): - decode against prefill (KL), at 2 / 3 / 7 / 8 rows: passed; - the four check cases: passed (above); - GSM8K past the largest width, without and with MTP, and the MTP accept-length bound: passed. sgl-eval 0.1.3 was put on `PYTHONPATH` with a pure-Python stand-in for its compiled `editdistance`, as for #68; - the injected fault fails here on `main` the same way, only at its last check: every rank stops at the wait deadline with `HIP error: an illegal memory access was encountered` and the server exits without sending a token, but ROCr in our container prints no `Memory access fault ... on address (nil)` line (not even for a bare one-GPU kernel that only stores to 0x0), so the test's `assertTrue(faults)` fails. In the reviewer's setup on #68 and #69 the line is printed and the case passes. - **Long context, decode against prefill at 128K** (an out-of-tree client: 4 prompts of 131,072 ShareGPT tokens with a GSM8K question, 64 greedy tokens, the decode's top-20 log-probabilities against a prefill of prompt and output; E4M3 KV, MTP off): - base: KL mean 0.081, top-1 agreement 0.941; - this PR: KL mean 0.063, top-1 agreement 0.956; - both arms have the same largest KL (1.87) over 252 positions. ## Speed Tests and Profiling **Setup:** as #68: one node with 8 × MI355X; image `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025` (squashfs `sha256:fb07e0ef78fe`), AITER 285579b6, FlyDSL 0.4.0; `zai-org/GLM-5.3` at `ba32f8c1` (FP8); the e2e test's server flags with `--kv-cache-dtype fp8_e4m3`. - The e2e and check runs overlay SGLang's `python/sglang` at `5044779f`. - The serving runs use the image's SGLang with `--enable-amd-monokernel` and the plugin's spec registration from before #48. - Kernel-only numbers don't depend on SGLang. - All numbers are from this head against its base: - serving, e2e, the check runs, the long-context runs and the bitwise dumps: this head (b6fdf8e) against the base (3f5f995); - kernel-only, compile survey, exact top-k and kernel tests: the same four commits on #73's previous head (809f6ec), whose kernel code is the same; #73's last change touched only `check.py` and `spec.py`. **Serving:** as #68. - One server per arm and setting: MTP 3-1-4 at 8K, MTP 3-1-4 at 128K, MTP off at 8K. All of a setting's points ran on the same server. - Prompts of 8,192 or 131,072 ShareGPT tokens, no shared prefix, radix cache flushed before each wave. 512 greedy output tokens with EOS ignored. 4 waves (3 at 128K) of C concurrent requests after a warm-up wave. - MTP runs at the golden accept length (`SGLANG_SIMULATE_ACC_LEN=2.99`; measured 2.99 on every server). A cycle is a request's decode time over its verify steps. - One round. **The client is out of tree**: `sglang.benchmark.serving` does not import on this image, so a small streaming client against `/generate` measures TTFT, TPOT and the verify count from `meta_info`. TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms (median over requests): | Context | Requests | base | This PR | Change | | --- | --- | --- | --- | --- | | 8K | 1 | 7.164 | 6.920 | −3.4% | | 8K | 2 | 10.760 | 10.539 | −2.1% | | 128K | 1 | 8.287 | 7.238 | −12.7% | | 128K | 2 | 12.655 | 11.483 | −9.3% | The requests of a two-request wave prefill one after the other. The first-prefilled request waits out the other's prefill (380 ms at 8K, 7.9 s at 128K) before both decode, so its cycle is the one above plus that wait over its verify steps. The table gives the other request, which decodes beside the first from its first token to its last; the MTP-off rows below are measured the same way. TP8, FP8 E4M3 KV, MTP off, 8K, median TPOT in ms: | Requests | base | This PR | Change | | --- | --- | --- | --- | | 1 | 4.446 | 4.424 | −0.5% | | 2 | 4.563 | 4.437 | −2.8% | | 4 | 5.872 | 5.773 | −1.7% | ### Kernel-only (`time_layers.py`) `torchrun --nproc-per-node <TP> test/kernels/glm53/time_layers.py --kind <kind> --rows <rows> --kv fp8 --pos <8447 | 131071>`, µs a launch, base → this PR. Each number is the mean of two passes in alternating order in one allocation; each pass is the median of 7 trials of 50 replays. The passes agree within 0.8 µs (1.3 µs on TP4's 4-row indexer launch). TP8, E4M3 MLA cache: | Launch | Context | 2 rows | 4 rows | 8 rows | | --- | --- | --- | --- | --- | | indexer | 8K | 71.54 → 69.99 (−2.2%) | 96.42 → 86.33 (−10.5%) | 145.52 → 132.98 (−8.6%) | | indexer | 128K | 96.66 → 78.65 (−18.6%) | 140.00 → 100.00 (−28.6%) | 223.84 → 173.65 (−22.4%) | | model order (75) | 8K | 54.08 → 53.68 (−0.7%) | 71.12 → 68.78 (−3.3%) | 115.86 → 112.66 (−2.8%) | | model order (75) | 128K | 60.46 → 55.52 (−8.2%) | 82.88 → 72.20 (−12.9%) | 135.27 → 122.45 (−9.5%) | TP4, E4M3 MLA cache, indexer launch: | Context | 2 rows | 4 rows | | --- | --- | --- | | 8K | 87.88 → 88.77 (+1.0%) | 134.13 → 125.26 (−6.6%) | | 128K | 113.92 → 95.55 (−16.1%) | 178.05 → 138.61 (−22.2%) | ### Each change on its own Each commit against the one before it, TP8, E4M3 MLA cache, indexer launch, µs at 2 / 4 / 8 rows. Two passes in alternating order, one allocation. | Change | 8K | 128K | | --- | --- | --- | | 1. index key on its own task, index-query tiles first | +0.65 / +0.93 / +2.00 | +1.68 / +1.13 / +3.24 | | 2. a request's rows scored together | +0.51 / −6.04 / −11.64 | +2.98 / −15.13 / −35.29 | | 3. top-2048 on up to 64 slices | −2.71 / −4.98 / −2.90 | −22.67 / −26.00 / −18.14 | | **total** | **−1.55 / −10.09 / −12.54** | **−18.01 / −40.00 / −50.19** | - Change 1 on its own is slower by up to 3.2 µs. It is kept because changes 2 and 3 were built and measured on it. - Change 2 costs the 2-row launch 0.5 µs at 8K and 3.0 µs at 128K: a 2-row launch has one live row, which gains nothing from reading a key once for several rows. - Change 3 takes most of the 128K gain. ### Registers The coverage PR's compile survey: every key the spec builds (`stop=True`) and the same keys as the tests build them (`stop=False`). That is TP4 (2, 4 rows) and TP8 (2, 4, 8); indexer, shared, dense, layer-0 and tail builds; bf16 and E4M3 MLA cache; 90 builds. - Every symbol is `glm5_monokernel_0`, as on the base. - The shared and tail code objects are byte-identical to the base's. - VGPR spills / private segment bytes, base → this PR, for the builds that have either: | Build | `stop=True` (spec) | `stop=False` (tests) | | --- | --- | --- | | TP8 indexer, 8 rows, bf16 | 4 / 20 → 24 / 100 | 5 / 24 → 20 / 84 | | TP8 indexer, 8 rows, E4M3 | 4 / 20 → 20 / 84 | 5 / 24 → 24 / 100 | | TP8 dense, 8 rows, bf16 | 1 / 8 → 20 / 84 | 0 / 0 → 20 / 84 | | TP8 dense, 8 rows, E4M3 | 0 / 0 → 20 / 84 | 0 / 0 → 20 / 84 | | TP4 indexer, 4 rows, bf16 | 11 / 48 → 15 / 64 | 40 / 164 → 48 / 196 | | TP4 indexer, 4 rows, E4M3 | 16 / 68 → 13 / 56 | 34 / 140 → 48 / 196 | - The 8-row indexer and dense builds spill about 20 VGPRs (84–100 bytes), against 0–5 on the base. They come with change 2. Even so, the 8-row indexer launch is 8.6% faster at 8K and 22% at 128K. - Every other build has no VGPR spill and no private segment on either head. - The layer-0 builds have none at any width. ### Tests - Unit tests on a CPU-only host: the base's results (this head 288 passed, the base 276; the base's 2 contract failures and 3 collection errors). - `test_glm53_kernels.py` with 8 GPUs (TP4 and TP8): 8 passed (17 min 32 s), including the new exact top-k test. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Summary. With long contexts, the GLM-5.3 monokernels were slower than the normal path, because the exact top-2048 selection of each indexer row ran on one CTA. This PR selects the top-k of a row of more than 32768 keys on up to 16 CTAs. The result does not change: it is the exact top-2048, as before. On 8x MI355X at TP8, in InferenceX agentic coding with MTP, the monokernels now decode faster than the normal path with MTP: interactivity p50 is 15% higher at concurrency 1, 17% at 2 and 11% at 4. With #15 alone, it was 5%, 5% and 14% lower.
On main. This PR is two commits on main (
876ed7f):6bfceb6is48ff491rebased with the same changed lines (and48ff491is4f187f6rebased), and235bb05is the fix from the review ("Review fix" below). Apart from that fix,gfx950/kernel.py,layout.pyand thecommonmodules they import have the same code as at4f187f6, so the results below measured at4f187f6and8068147still apply. The fix's own results are marked.Problem.
e952b22, the monokernels lost to the normal path at concurrency 8: interactivity p50 -14%, p90 -42%. With MTP ([Feat] [MI355X] glm53: Serve MTP verify steps on the monokernels #15), they lost 5% to 14% at concurrency 1 to 4 (table in "Speed Tests and Profiling").Cause. A debug build (not in this PR) records the start and the end of each stage task of each launch. In an indexer launch:
The estimated step is 18 indexer launches and 57 launches without the indexer. From 8k to 256k keys, it grows 4.1 ms at 1 row and 6.6 ms at 4 rows. The measured decode step grows 4.1 ms and 6.3 ms.
Modifications
glm53/gfx950/layout.pyINDEX_SEL_CHUNK = 4096(the keys of a select chunk: 512 threads x 8),INDEX_SEL_PAR_CHUNKS = 8,INDEX_SEL_SLICES = 16andINDEX_CAND_STAGE_CAP = 16384.select_slices(bound)gives the slices of a row and the chunks of each slice. The kernel and the CPU test call the same function.index_select:index_hist(16 tasks a row),index_thresh(1 a row) andindex_gather(16 a row), and their scratch: about 2.4 MiB a row, 33 MiB a rank for widths 2, 4 and 8.glm53/gfx950/kernel.py: a row of more than 32768 keys (8 chunks of 4096) selects on up to 16 CTAs of the upper half of the grid (CTAs 128 to 255), each with a run of whole chunks. Before this change, these CTAs waited idle for the selection.index_hist: each slice (CTA) makes a 4096-bin histogram of its chunks in LDS (the top 12 bits of each score key) and writes it to scratch. A release fence orders the histogram before its ready flag.index_threshruns on the CTA of the row'sindex_select. It adds the histograms of the slices, finds the threshold bin, and publishes the bin, the keys still needed from the bin and the count in the bin.index_gather: each slice writes, in position order, its positions above the bin (winners) and its (key, position) pairs in the bin (candidates) to scratch, and publishes the two counts.index_selectcopies the winners and the candidates of the slices, in slice order, into the lists that the one-CTA path makes. Then it runs the same radix finish on the candidates.index_threshandindex_selectread the slices' histograms and staged lists, plain stores released before their ready pairs, after an agent-scope acquire, as theindiceshand-off does. A parallel row whose threshold bin overflows polls its scores in the overflow pass, because its CTA has not read them before.glm53/check.py,glm53/spec.py(review fix): check mode compares each indexer row's selection with the exact top-k of the launch's own scores (score descending, then position):layer.topk_exact.max_abs, the keys a row misses, bound 0. Every layer check case runs it.test/unit/test_glm53_select.py(new):select_slicescovers each chunk of a row once, in position order, with at most 16 slices and no empty slice (138 cases, 135 distinct bounds); the slice tasks of each width fit in the upper half of the grid.test/e2e/test_glm53_server.py:test_each_covered_layer_past_the_parallel_top_kruns the layer check on rows of about 108k keys at widths 2 and 4 (27 chunks: 14 slices, 13 of two chunks) and of about 36k keys at width 8 (9 one-chunk slices). It asserts thatpast(65536) is past one chunk a slice._check_layersgetspast(an indexer check must reach this many keys) and requireslayer.topk_exactin the log.No change. Rows of 32768 keys or fewer, the decode steps of GLM-5.3-Flash, DeepSeek-V4.1 and MiniMax-M3, and the normal path.
Accuracy Tests
Setup. 8x MI355X (gfx950) at TP8; zai-org/GLM-5.3 @
ba32f8c1(FP8); the image and the commits in "Speed Tests and Profiling"; the server flags oftest/e2e/glm53_server.py; a new server for each test.Test 1: exact top-k. A debug build (not in this PR) dumps the scores and the selected positions of each row of a launch. A script computes the exact top-2048 of the dumped scores (score descending, then position ascending) and compares the two sets. The prompts are disjoint windows of real text (SGLang's source). 312 rows (8 ranks; 8k to 256k keys; 1 to 8 rows): 0 rows differ, and all the scores of each row come from the same launch.
Test 2: E2E tests at TP8 (review fix). We ran
test/e2e/test_glm53_server.pyas written at235bb05on main876ed7f. All 7 tests passed in one run (1234 s):TestMonokernelDecode: average KL 0.015 to 0.025 for each batch (bound 0.06); GSM8K (sgl-eval, thinking on) 0.995.TestMonokernelMtp: GSM8K 0.995, accept length 3.31.TestMonokernelCheck: no violations.layer.topk_exactandlayer.topk_ownare 0 in all 918 indexer checks.layer.topk.rel0.015 or less.layer.topk.rel0.017 or less.layer.topk.rel0.051 or less.layer.topk.relis the share of SGLang's selection that a row misses. It has no bound: SGLang scores with FP8 queries and the kernel with bf16 queries, so keys near the threshold can swap. It is highest in the case whose prompt repeats one paragraph thousands of times; we did not measure why.layer.topk_exactchecks the selection itself against the launch's own scores.TestMonokernelFault: passed. Since [Test] glm53/glm53_flash: Turn off crash dumps in the fault tests #55, the test setsHSA_DISABLE_COREDUMP_ON_EXCEPTION=1and turns off SGLang's crash dumps itself.An earlier run of the same file at
8068147, before the review fix, also passed all 7 tests.Test 3: the exact check fails on a wrong selection. A deliberate error (not in this PR): the slices stage their candidates with the low 20 bits of each key flipped, so the radix finish takes the lowest keys of the threshold bin instead of the highest. Every position stays valid. On the parallel check case,
layer.topk_exactreports 204 violations (up to 1422 keys of a row's exact top-k missed), and the test fails. No other bound catches it: the lowest stream cosine is 0.991 (bound 0.95), andlayer.topk_ownstays 0. A first version of this control, which dropped the last slice's winners, made the launch publish stale positions and fault on memory instead.The prompts of the GSM8K and KL tests have fewer than 32768 keys, so these tests run the one-CTA path: they show that this path does not change. Test 1 and the new layer check cover the new path.
Speed Tests and Profiling
Setup.
lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261019(sha256:8cb29fc407c2).e45cbce599, AITER2bb549102b, FlyDSL 0.3.4.1, torch 2.9.1+rocm7.2.0.ba32f8c1.--mem-fraction-static 0.80 --chunked-prefill-size 131072 --dsa-prefill-backend tilelang --dsa-decode-backend tilelang), with--max-running-requestsand--cuda-graph-max-bs-decodeat 2x the concurrency. MTP: EAGLE, 3 draft steps, top-k 1, 4 draft tokens. The arms differ only in--monokernel-backend amdand the checkout.amd-smibefore and after).Review fix (
235bb05). Fixed length without MTP: 131072-token prompts, 512 new tokens, 4 prompts for each concurrent request, before (6bfceb6) and after the fix, a new server for each, two rounds with the second in reverse order. Median ITL (ms):At concurrency 4 the TPOT median (30 to 40 ms) includes the other requests' 128k-token prefills, so ITL is the decode measure. In a compile-only build of the 22 GLM-5.3 keys, the 11 builds without the indexer are unchanged, and the indexer builds have 0.3% to 1.4% more instructions. At TP8 no build spills more VGPRs; at TP4, width 4 spills 10 instead of 9.
Kernel timeline: #15 to this PR. The debug build of the cause section, on #15 (
91f0cc7) and on this PR (8068147): a new server for each, random prompts, 1 to 8 rows decoding together, the last launch of each of 6 passes on each of the 8 ranks (medians). The estimated step is 18 indexer launches and 57 launches without the indexer.The launches without the indexer run the same code in both builds; between the two runs they differ by -3.2% to +3.1%. Thus the differences at 8k and 32k keys are in the noise of this measurement. At 32k keys (9 chunks, just past the one-CTA path), the selection takes 2 to 5 us longer at 4 and 8 rows. The score stage does not change.
AgentX with MTP. InferenceX agentic coding (
agentx-harness1.0.11, corpus23f152f6), 3600 s a point. All arms use the same simulated accept length (2.99), so the accept length does not differ between the arms.In the same time, this PR completes more requests (for example 246 against 234 at concurrency 1), so its sessions go further; at concurrency 1 and 2, its inputs are longer.
Time of a verify step for each generated token (ms), from the server log intervals without prefill, all three concurrencies:
(normal path with MTP / #15 / this PR). With 2 requests past 192k keys, the monokernels are still slower than the normal path. At 5 to 8 rows, the score stage takes about half of the indexer launch (timeline above); a follow-up PR changes it.
Fixed length without MTP. 131072-token prompts, 1024 new tokens, 6 prompts for each concurrent request, the radix cache flushed before each concurrency. TPOT median (ms):
e952b22AgentX without MTP. The same selection code on
e952b22(without #15), concurrency 4 and 8, 3600 s a point. Interactivity p50 / p90 (tok/s/user):e952b22Checklist
pre-commit run --all-filespasses.test/cover the change..github/CODEOWNERSthat names its owner.Review and Merge Process
lintandCheck PR titlechecks must pass..github/CODEOWNERSto review.main.