Repository navigation
[Perf] [MI355X] glm53: Long-context indexer: two-row scoring from 4 rows, keys straight into registers, the scoring layout ahead, first digits from the scoring - #86
Merged
jhinpan merged 4 commits intoNov 3, 2026
Conversation
heslami
requested review from
HaiShaw,
alexnails,
jhinpan and
kevin-mii
as code owners
October 29, 2026 19:22
heslami
force-pushed
the
perf/glm53-index-longctx-stack
branch
4 times, most recently
from
November 1, 2026 23:30
33703da to
150e395
Compare
This was referenced Nov 2, 2026
Collaborator
|
Same here. Hi @xiaobochen-amd! Can u help cross-review this PR about monokernels on GLM 5.3? Thx for your help! |
@jhinpan LGTM. I re-ran this PR on top of #85 with the same setup: MI355X, TP8, FP8 KV, MTP 3-1-4 with the accept length fixed at 2.99. Speed: ms per verify round, the median of two runs after a warmup.
The gain grows with context, as the description says: under 1% up to 8K, and about 8% at 512K. Accuracy
|
jhinpan
pushed a commit
that referenced
this pull request
Nov 3, 2026
…-split front, indexer launch, MoE expert dedup and weight loads, kernarg reload, LM-head loads (#85) ## Motivation On main. One commit per optimization, 24 in all, then one cleanup commit that defines the up/gate tile helpers once and folds `reload_args`' guard into it (−202 lines; every compiled kernel is byte-identical to the commit before it). A GLM-5.3 decode step runs 75 MoE layer launches plus layer 0, the dense layers 1–2 and the tail; with MTP each cycle also runs the NextN draft launches. This PR shortens the hand-offs and the weight streams of those launches: - **Shared layer and MoE.** down read all of mid as tagged f32 on every CTA; it now reads mid quantized to FP8 once. At 8 rows the split tiles' KV rows waited behind q_b and uk; they now go out at the launch start. At 4 rows every CTA routed all four rows before its first up/gate job; a CTA now runs one row's jobs, its first after that row's first two picks. TP4 now runs 8-row steps on the kernels. - **NextN.** The entry runs as four K parts on 192 CTAs instead of two halves on 96. Each stage reloads its pointer arguments from the kernarg segment, which avoids a FlyDSL 0.4.0 LLVM miscompile (ROCm/FlyDSL#1344). - **The front** (q_a | kv_a) ran as 164 whole-K tiles while 92 CTAs waited; each tile's K is now split between its own CTA and a spare one. Top-8 rounds reduce the expert ID only on ties, and the router's RMSNorm moves into its logits. - **The indexer launch.** Above 4 rows two rows score a wave; the index projection tiles run alone and the scoring skips the q_norm CTAs; split stages q and the new rows before the selection's wait; the selection warms its code in L2, bounded at the code's end. - **The MoE at 8 rows** runs each expert the step uses once in up/gate and down. FP8 weights load 16 bytes a lane, weights read once a launch load nontemporal, and at 4 rows a repeated expert streams once. - **SGPR spills.** The indexer and dense-MLP launches reload their pointer arguments at each stage, as the NextN launches do. - **The tail** keeps its LM-head loads in flight across its tiles. **Every output and every write stays bitwise equal to #75, except the K-split front's** (commits 8 and 10, below). Every other commit is bitwise against its parent at TP8 and TP4, every width, on the bf16 and the E4M3 MLA cache. The NextN entry's partial-sum mailbox changes layout (commit 4). **The one non-bitwise part, the front:** each front part multiplies bf16(x · gamma) rather than the normalized row, and the consumers scale the sum by the rstd, so q_an, the KV rows, the router scores and the outputs move by about one bf16 rounding; the router fold changes the router logits only. Its accuracy evidence is GSM8K with MTP against an A/A pair and KL against prefill at TP8 (Accuracy Tests). Result on MI355X, #75 → this PR (TP8, FP8 E4M3 KV, MTP 3-1-4 at the golden accept length 2.99): - **Per cycle at 8K: 6.467 → 6.017 ms at one request (−7.0%) and 9.967 → 8.560 ms at two (−14.1%)**, measured step by step from #75 through commit 22 (Speed Tests). The tail's in-flight LM-head loads (commit 24) and the bounded code warm (commit 15) were not in that serving tree, so they are not in this number. - **TP4, two requests, 8K:** −10.4% for commits 1–5, −5.0% for the MoE dedup (16, 17), −4.3% for the 4-row expert-once and nontemporal loads (19–21). - Real acceptance (TP8, 8K, paired): TPOT −4.5% at one request and −7.8% at two for commits 1–10 (Speed Tests). ## Modifications 24 commits, one per change, in the order below (#). Three only add test options. | # | Area | Commit | Change | Numerics | |---|---|---|---|---| | 1 | Shared layer and MoE | mid as FP8 once | Wave k of down task t quantizes mid block t + k × 256 and publishes its FP8 bytes and scale (`midq`, `midqs`). Every down task polls those instead of all of mid. At TP4 and 8 rows (288 blocks) two waves a task quantize. MXFP4 experts keep their path | bitwise | | 2 | Shared layer and MoE | early split rows | Above 4 rows, in launches without the indexer (TP8; at TP4 there are more split tiles than CTAs), each split tile's keys and KV cache rows are loaded right after the CTA's q_norm and cache tasks, ahead of its q_b or uk task | bitwise | | 3 | Shared layer and MoE | per-row up/gate jobs at 4 rows | With one up/gate task round (TP8, 4 rows), the CTA holding tiles c of slot group g runs row g % 4 at slots 1 + g/4, +2, +4, +6: the same tiles as before, each on its XCD. Its routing waves run the first two top-8 rounds, park their state in LDS, issue the first job's weights, and run the other six rounds under them. A tied round hands the row to the existing top-8 | bitwise | | 5 | Shared layer and MoE | TP4 at 8 rows | The per-TP row cap goes. TP4 routes target decode / verify steps and draft extends of up to 8 rows. Draft decode steps keep their 4-row limit at both TP sizes. No code object changes | — | | 4 | NextN entry and its kernarg reload | NextN entry in four parts | Each 16-column entry tile runs as four K parts of `[enorm(e) | hnorm(h)]` on 192 CTAs (TP8). The parts' f32 partials are added in part order by the first part, which pushes the tile. With the indexer the entry starts past the front's index projection tiles. TP4 keeps the two halves (192 tasks already). The entry normalizes 4 rows at a time, and at up to 4 rows a NextN launch polls its pushed rows 2 at a time in the front tiles | bitwise, except the `ehp` mailbox layout | | 6 | NextN entry and its kernarg reload | NextN kernarg reload | `KERNARG_I64` lists the kernel's Int64 arguments in signature order (argument i is the 8 bytes at 8·i of the kernarg segment). A NextN launch takes the segment's address in its prologue (`llvm.amdgcn.kernarg.segment.ptr`); `reload_args()` rebinds the 35 pointer arguments the stages read, and `xin` from the reloaded `sym`, as invariant scalar loads behind an empty asm on the segment's address, so LLVM cannot merge one stage's loads with another's. Each of the 14 stages after the NextN entry calls it at its start; LLVM drops the loads a stage does not use. Other kinds trace exactly as before | bitwise | | 7 | NextN entry and its kernarg reload | NextN kinds in `time_layers.py` | `--kind nextn_index` and `--kind nextn_shared` time one NextN launch a step, `check_nextn.py`'s draft layer, at the served decode points' rows, in the same graph replays as the other kinds. `check_nextn.Draft` takes the rows directly | test only | | 8 | K-split front (the one non-bitwise part) | K-split front | Each q_a / kv_a tile's first 4096 K run on its own CTA, and the remaining 2048 on one of the 92 CTAs past the tiles, two tiles' rests a CTA. Neither part normalizes: each multiplies bf16(x · gamma). q_norm and the cache tasks add the two partials (`q_a2`, `kv_a2` hold the rests) and scale them by the row's rstd, which they compute from x before their wait. The index projection tiles keep whole K, on the CTAs with one rest, and the NextN entry starts past them. The MoE stages keep their CTAs. A launch from pushed rows (layer 0, NextN) reads them in both parts and in the rstd, with its weights issued before the rows' wait | not bitwise: q_an, the KV rows, the router scores and the output move by about one bf16 rounding (GSM8K and KL in Accuracy Tests) | | 9 | K-split front (the one non-bitwise part) | top-8 tie shortcut | A round takes the key maximum. When one lane holds it (a ballot), that lane's head is the pick; only an exact tie runs the ID reduction. The picks come out in the same order | bitwise | | 10 | K-split front (the one non-bitwise part) | router norm in its logits | The router task stores bf16(a · gamma) as its rows arrive, with each wave's sum of squares beside them. One barrier serves both. The rstd, summed in the same order, scales the FP8 activation blocks (bitwise) and the 8 logits | not bitwise: router logits only | | 11 | Indexer launch | two-row scoring above 4 rows | Above 4 rows a scoring wave runs two rows' MFMAs on one LDS read and FP8 → bf16 conversion of each key, a run's keys load nontemporal, runs split evenly once every scoring CTA has one, and a multi-batch run stages a batch's scores in LDS and stores them after the next batch's keys land. The index-weights tiles keep a quarter of their weights in flight before the rows' staging instead of half. Timeline builds load the timeline's address once a launch. At 2 and 4 rows the scoring code is unchanged | bitwise | | 12 | Indexer launch | index tiles and scoring placement | With the indexer the 10 index K / weights projection tiles run alone on the grid's last 10 CTAs and the K rests go two a CTA on the CTAs before them; the NextN entry, which starts past the index tiles, starts at CTA 0. The scoring runs go to CTAs S .. G − 1 only | bitwise | | 13 | Indexer launch | split before the selection, selection code warm | With the indexer, split polls q and the rows this launch appends before the selection's wait, staging the rows in its reduction and output regions; each wave then polls its own 8 selected keys and issues their cache rows as they land. Each selection slice, while its scores are pending, loads the code after it into its XCD's L2 (one word a 64-byte line, from `s_getpc`). TP8 dense launches above 4 rows skip that warm | bitwise | | 14 | Indexer launch | the two-row scoring's queries in MFMA operand order | Above 4 rows the scoring stages its rows' queries in LDS in MFMA operand order (head half, 32-dim chunk, lane), so each wave reads a query fragment as one conflict-free 16-byte load a lane; row-major, a fragment's 16 lanes read 16 heads 64 words apart, all on one bank, twice a wave | bitwise | | 15 | Indexer launch | the code warm ends at the code's end | `warm_code` takes its PC and the distance to `.Lfunc_end0` from one inline asm (`s_getpc_b64`, a label, `s_mov_b32 sN, .Lfunc_end0-.Lwarm_code`; the assembler resolves the distance). The buffer's size is min(distance, 64 KB), and loads past it read nothing. With a fixed 64 KB the dense and embed launches read up to about 8 KB past their code, and TP4 dense launches faulted (illegal address). Where 64 KB fits, the read is unchanged | bitwise | | 16 | MoE dedup and weight loads | each expert once in up/gate and down | Above 4 rows (not dense layers) every wave finds the step's distinct routed experts from the top-8s (lane p holds pick p, a ballot marks first uses). up/gate jobs are the 8-intermediate tiles of the shared expert and of each distinct expert, each with every using row as an MFMA column, dealt round-robin over the CTAs; every CTA routes, issues the activations' poll and its first job's weights, then checks the poll. down's units take one expert's tile against every using row and store the products in LDS; a row's output is summed over its units in the one-row form's order and wave segments. The rows the multi-column MFMAs read are padded by 4 words against bank conflicts | bitwise | | 17 | MoE dedup and weight loads | dead-batch skip | A wave's down weight batch whose first unit lies past the step's expert groups skips its MFMAs | bitwise | | 18 | MoE dedup and weight loads | 16-byte FP8 loads | `ld_fp8_chunk`, shuffled layout: a lane of an even 16-lane row loads the whole 16-byte segment of a row's first K32 block, and its partner 16 lanes up the same segment of the second block, as one `dwordx4`. `_RowSwapped` holds the raw 16 bytes; reading element `e` runs `v_permlane16_swap` on words (e % 2, e % 2 + 2), so the load stays in flight across the task's wait. The packed layout and the callers are unchanged | bitwise | | 19 | MoE dedup and weight loads | `--similar` | `check_layer.py` and `time_layers.py --similar X`: the live rows of a request are one random row plus X times their own, so they pick many of the same experts, as a verify step's rows do. The default keeps the independent rows and their digests | test only | | 20 | MoE dedup and weight loads | 4-row repeated expert once | At 4 rows with one up/gate task round (TP8, not dense, not the NextN indexer launch), each job takes every row as an MFMA column. The owner of an expert is the row whose job of it comes first, the lowest such row on a tie, found by a ballot over the step's picks; the owner writes mid for every row that picked it, through the 8-row path's `job_emit`, and the other rows' later jobs of it load nothing | bitwise | | 21 | MoE dedup and weight loads | nontemporal once-read weights | `ld_fp8_chunk` and the FP8 unit helpers take a cache policy (`CM_NT` in `tiles.py`). At 2 and 8 rows: q_a \| kv_a, q_b, W_UK and o_proj, and the experts' up/gate and down (routed, shared and the dense layers' slices). At 4 rows: up/gate only, where commit 20 applies | bitwise | | 22 | Kernarg reload for the indexer and dense launches | kernarg reload for indexer and dense launches | `RELOAD_ARGS` (NextN, indexer or dense-MLP launch, layer 0 included) selects the launches whose 14 stages reload their pointer arguments from the kernarg segment, as the NextN launches already do; `reload_args()` rebinds `xin` only where x is pushed. Shared-layer and tail launches build exactly as before | bitwise | | 23 | LM-head tail | tail digest in `check_layer.py` | `--dump` digests the tail's gathered logits | test only | | 24 | LM-head tail | tail loads in flight | A tail wave keeps `HD_AHEAD` units in flight across its tiles: each unit it multiplies is replaced by the next one of its tile sequence, so the next tile's weights load while the current tile reduces and pushes. The norm's row loads go out before the first weights | bitwise | - **Layout:** `midq` and `midqs` after `mid` (19 KB at 4 rows, 38 KB at 8); `ehp` holds three NextN entry parts' partials instead of one; `q_a2` and `kv_a2` hold the front's K rests. The MoE group table and down's products live in the LDS X region; no other new scratch. - **Spec and draft:** `TP_MAX_ROWS` and `_Built.max_rows` go; `DeviceLimits(widths=ROWS)` bounds every routed step at 8 rows. `test_glm53_kernels.py` runs TP4 at widths 2, 4 and 8. The README's GLM-5.3 row says up to 8 rows at both TP sizes. - **Placement:** `qkv_a` has a task on every CTA; q_norm, q_b, split, uk and uv follow at CTA 0 as before, and o starts at split's first CTA. With the indexer the index projection tiles take the grid's last 10 CTAs; the NextN entry starts at CTA 0 there. - **`poll`** takes an optional `between()` that runs after the first loads are issued and before they are checked (up/gate issues its first job's weights there). - **What changes in #74's code:** the scoring keeps its per-row path at 2 and 4 rows and adds a two-row path above 4; a run's first CTA and page range are counted over the scoring CTAs (S .. G − 1). The selection gains one call, the code warm. Split's indexer path is reordered; without the indexer split keeps its order. ## Accuracy Tests All on MI355X in the image below, under torchrun on 8 GPUs (TP8) and 4 GPUs (TP4). - **Bitwise, each commit against its parent** (every commit but 8 and 10). Commits 1–21 were measured with these tools on the same commits in an earlier stack, whose kernel files are this branch's except commit 15's bounded code warm. Commits 15 and 22 were measured as this branch's head against that stack's head. The dump set has no head kind, so commit 24's tail logits were not in this run. Runs at 2 rows need #76's `peer.py`: without it, #75 run twice gave 1,693 differing digests at TP8, and a few 2-row configs differ run to run at TP4 too; with it, 0. - Run as, on each tree, for each kind:launches in `shared:3 index:3 model:8 dense:3 embed:3` and each kv in `bf16 fp8`: `torchrun --nproc-per-node 8 test/kernels/glm53/check_layer.py --kind K --launches N --rows 2,4,8 --rowsets p0,p1,v4 --kv KV --dump DIR/<tree>/tp8` (and `--nproc-per-node 4` into `tp4`); then `python3 test/kernels/glm53/check_layer.py --compare DIR/<parent>/tp8 DIR/<tree>/tp8`. Commit 20 also with `--rowsets v4 --similar 1.0`. Negative control: the same dump of #75 with `--perturb`. | Area (commits) | TP8 | TP4 | | --- | --- | --- | | Shared layer, MoE and NextN entry (1–5) against #75 | 0 differ | 0 differ | | NextN reload (6) | 0 differ | 0 differ | | Front (8–10) | differs by design (commits 8 and 10); every CHECK passes | same | | Indexer launch (11–15) | 0 differ | 0 differ | | MoE dedup (16, 17) | 0 differ | 0 differ | | 16-byte loads (18) | 0 differ | 0 differ | | 4-row expert once, nontemporal (20, 21) | 0 differ | 0 differ | | Bounded code warm (15), kernarg reload (22): head against the earlier stack's head | 0 differ | 0 differ | | Tail loads (24): the tail's logits (`--kind head`, commit 23) | not in this run | not in this run | - **The front, not bitwise (commits 8 and 10).** Measured at TP8 with MTP on and real acceptance, the commit before the front (#23's head) against the front's head, and #23's head again as an A/A pair: - **GSM8K** (1,319 questions, 8-shot, two requests): 0.9234 (#23), **0.9166 (front)**, 0.9136 (#23 again). McNemar, one-sided: front vs #23 p = 0.145; the A/A pair p = 0.065; all above 0.025. - **KL decode ‖ prefill** (MTP off, 64 prompts × 128 tokens, top-20): 0.0848 (#23), **0.0766 (front)**, 0.0721 (#23 again). Top-1 agreement 0.931 / 0.932 / 0.936. - On both, the front's value lies within #23's own A/A spread. - AIME 2025, 30 problems × 16 samples, with the model card's sampling (temperature 1.0, top_p 0.95, thinking on, up to 163,840 tokens). TP8, FP8 KV, MTP 3-1-4. Generations average about 23K tokens and some reach the cap, so the indexer runs over long KV. | | Accuracy | 95% CI | Accept length | |---|---|---|---| | SGLang without monokernels | 99.17% | 98.1–100 | 3.21 | | sgl-amd-mono main | 99.17% | 98.3–99.8 | 3.23 | | With this PR, #86, sgl-project/sglang#47324 and its follow-up | 98.33% | 96.9–99.6 | 3.23 | The stack has 8 wrong answers out of 480, against 4 for each of the others (Fisher p = 0.39; against main, paired by problem, p = 0.36). Samples are independent draws. The reference served 16 requests at a time, the monokernel builds 2. - Decode drift at long context: KL of decode against SGLang's stock prefill of the same tokens (top-20), and top-1 agreement. TP8, FP8 KV, MTP off; 128 prompts × 256 greedy tokens at each of 8K, 32K, 64K and 128K, cut from the AIME reasoning traces; two servers per build. | | KL decode ‖ prefill | Top-1 agreement | |---|---|---| | SGLang without monokernels | 0.0249 | 95.94% | | sgl-amd-mono main | 0.0199 | 96.34% | | With this PR, #86, sgl-project/sglang#47324 and its follow-up | 0.0203 | 96.20% | The stack against main: KL +1.9% (95% CI −1.0 to +5.0), top-1 −0.14 points, both within the spread between two servers of the same build (up to 4% in KL and 0.25 points in top-1). The front alone (commit 10 against commit 7): KL −1.0% (−4.8 to +3.0). With MTP on at 64K: KL −5.7% (−12.8 to +2.1). - **Stress** (TP8, the front's head): 4 runs × 20,000 graph replays, 0 bad. - **E2E** (`test/e2e/test_glm53_server.py`) at TP8 and at TP4 (`--tp-size 4`), on the same commits in an earlier stack: 8 passed. `TestMonokernelFault` fails as it does on main. ## Speed Tests and Profiling Setup as #75: image `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025` (squashfs `sha256:fb07e0ef78fe`) with SGLang's `python/sglang` overlaid at `5044779f`, AITER 285579b6, FlyDSL 0.4.0, `zai-org/GLM-5.3` at `ba32f8c1`, the e2e test's server flags with `--kv-cache-dtype fp8_e4m3`. **Serving ladder**, TP8, 8K, golden AL 2.99, FP8 E4M3 KV, MTP 3-1-4; one server per arm; the cycle at two requests is the last-prefilled request's. Measured step by step on the same commits in an earlier stack: | Step (this PR's commits) | Cycle, 1 request (ms) | Δ | Cycle, 2 requests (ms) | Δ | | --- | --- | --- | --- | --- | | #75 | 6.467 (A/A 6.444) | | 9.967 (A/A 9.986) | | | shared layer, MoE, NextN entry, TP4 at 8 rows (1–5) | 6.296 | −0.171 (−2.6%) | 9.585 | −0.382 (−3.8%) | | NextN reload and the front (6–10) | 6.232 (first sweep 6.236) | −0.064 (−1.0%) | 9.257 (first sweep 9.320) | −0.328 (−3.4%) | | indexer launch (11–15) | 6.108 | −0.124 (−2.0%) | 9.189 | −0.068 (−0.7%) | | MoE dedup (16, 17) | 6.103 | −0.005 (−0.1%) | 8.780 | −0.409 (−4.5%) | | 16-byte loads (18) | 6.126 | +0.023 (+0.4%) | 8.776 | −0.004 (0.0%) | | 4-row expert once, nontemporal (19–21) | 6.049 | −0.077 (−1.3%) | 8.562 | −0.214 (−2.4%) | | kernarg reload (22) | 6.017 | −0.032 (−0.5%) | 8.560 | −0.002 (0.0%) | | **#75 → (22)** | | **−0.450 (−7.0%)** | | **−1.407 (−14.1%)** | - 128K, the front's step → the indexer launch (11–15): 6.534 → 6.417 ms (−1.8%) at one request, 10.195 → 10.070 ms (−1.2%) at two. - The tail (23, 24) and the bounded code warm (15) are not in this ladder. **TP4, two requests, 8K** (cycle in ms): | Step | Before | After | Δ | | --- | --- | --- | --- | | #75 → shared layer, MoE, NextN entry, TP4 at 8 rows (1–5) | 15.857 | 14.213 | −10.4% | | indexer launch → MoE dedup (16, 17) | 14.043 | 13.346 | −5.0% | | 16-byte loads → 4-row expert once, nontemporal (19–21) | 13.319 | 12.741 | −4.3% | **Real acceptance**, #75 → the front's head (commits 1–10), TP8, 8K, 40 waves, paired: | Requests | AL | TPOT (ms) | Δ TPOT | Paired Δ | | --- | --- | --- | --- | --- | | 1 | 3.265 → 3.269 | 1.936 → 1.836 | −5.2% | −4.5% | | 2 | 3.290 → 3.310 | 2.979 → 2.671 | −10.3% | −7.8% | **AgentX** (InferenceX `inferencex-agentx-mvp`, seed 42, traces in InferenceX order) with AMD's InferenceX recipe: TP8, FP8 KV, MTP 3-1-4 at simulated accept length 2.99, one hour per point. Main and the stack ran back to back on the same node. | Concurrency | | P90 interactivity (tok/s/user) | Throughput (tok/s/GPU) | ITL p50 (ms) | TTFT p50 (s) | | --- | --- | --- | --- | --- | --- | | 1 | SGLang without monokernels | 181.5 | 1,944 | 4.808 | 0.351 | | | sgl-amd-mono main | 291.1 | 2,890 | 2.484 | 0.372 | | | With this PR and #86 | 345.3 (+18.6%) | 3,135 (+8.5%) | 2.207 (−11.2%) | 0.375 | | 2 | sgl-amd-mono main | 212.2 | 2,911 | 2.491 | 0.318 | | | With this PR and #86 | 282.9 (+33.3%) | 3,160 (+8.5%) | 2.219 (−10.9%) | 0.318 | | 4 | sgl-amd-mono main | 167.1 | 4,069 | 2.706 | 0.256 | | | With this PR and #86 | 181.0 (+8.4%) | 4,369 (+7.4%) | 2.326 (−14.1%) | 0.263 | - Paired by request (same conversation and turn), ITL falls 11.6% at concurrency 1 (95% CI 9.0–12.3), 10.9% at 2 (8.9–11.6) and 10.0% at 4 (9.1–11.9); TTFT is unchanged at every point. Throughput is a fixed one-hour window, so a faster server gets further into the traces. - SGLang without monokernels ran at concurrency 1 only, on the same node. Paired by request, ITL is 49.0% lower on main (95% CI 48.1–50.4) and 54.7% lower with the stack (54.4–54.8). ### Registers and lint From the code objects of the spec's builds (`stop=True`): `/opt/rocm/llvm/bin/llvm-readelf --notes <key>.co`, fields `.vgpr_count`, `.vgpr_spill_count`, `.sgpr_count`, `.sgpr_spill_count` and `.private_segment_fixed_size`. Two figures a cell are bf16 / E4M3 MLA cache. **#75 → this PR, at 8 rows:** - TP8 index: VGPR spills 24 / 20 → 2 / 1, private segment 100 / 84 → 12 / 8 B. TP8 dense: 20 → 1. - TP8 NextN shared: 12 → 0; NextN index: 12 → 6 / 4. - TP8 shared: SGPR spills 253 / 295 → 23 / 18. - TP8 embed: VGPRs 228 → 256, no spills. - TP4 index: VGPR spills 28 → 10; dense 24 / 20 → 4 / 5; NextN shared 20 → 0. - The one 8-row build that gains spills is TP4 embed: 0 → 3 VGPR spills (16 B). - SGPR spills fall on every 8-row build, from 300–1,700 to about 100. **Lint:** #84's `test/kernels/test_kernarg_lint.py` (the GLM-5.3 index and shared launches at TP4 and TP8, widths 2, 4 and 8, both caches): 24 / 24 pass. ## Reviewer notes - **The front:** the rstd is computed twice, in q_norm and in the cache task, from x, before their waits; the extra x read overlaps the front. With all 256 CTAs in the front, q_norm and the cache tasks sit behind a front tile on CTAs 0..S−1, which the indexer commits then place around. The pad-row skip the development tree carried at 2 rows is not part of this PR. - **Early split rows assume the selection is the previous launch's**; the indexer launches select in-launch and use commit 13's path. - **The code warm** reads code bytes through a buffer resource built from `s_getpc`, sized by the assembler's distance to `.Lfunc_end0` (capped at 64 KB), so it stays inside the kernel's code; the words are kept live through `fresh` and not used. - **Two scoring paths.** The per-row path stays for 2 and 4 rows; the two-row path is used above 4 rows only. - **The expert dedup is 8 rows only;** the dead-batch skip needs its group count. The group table is per CTA: every CTA routes all rows, so no cross-CTA hand-off is added and the tie order is unchanged. At 4 rows, `owns(e)` reads the step's picks after the routing barrier; a row's first job is issued before it and always loads. - **Nontemporal loads** return the same bits; W_UV, the indexer's weights, the router, the NextN entry's eh projection and the LM head keep the default policy, and TP4 at 4 rows keeps it everywhere. - **The kernarg reload** only moves where the pointers come from. FlyDSL 0.4.0's code objects differ from build to build (ROCm/FlyDSL#1345); #84 makes them reproducible and adds the lint. - **The spec** still does not build the NextN shared launch at 8 rows (#75's limit for draft decode steps). - `kernel.py` grows by 1,397 lines (+1,843 / −446). ## Checklist - [ ] `pre-commit run --all-files` passes. (ruff, ruff-format, isort and codespell are clean on the touched files) - [x] Tests under `test/` cover the change. (`check_layer.py --similar`, the tail digest, the NextN kinds in `time_layers.py`; the kernel tests compare every launch) - [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>
Collaborator
|
cc @heslami Please resolve conflicts here. And also cc @Jacob0226, pls feel free to give this PR a double check if u got time. |
This was referenced Nov 3, 2026
added 4 commits
November 3, 2026 12:10
… and lane-swap head sums At long context the indexer's scoring scan is bound by per-wave compute, not by key bytes in flight. Four changes cut its instructions: - A scoring wave's 16-key group addresses come from its page and first key in scalar math; only the lane's key offset is added per lane. - A row's two head halves reduce across lane rows in three permlane swaps, and two rows share one score store. - From 4 rows (one request's verify rows) a wave scores two rows, as 8-row scoring did: each key is read from LDS and converted once for both rows, with the rows' queries staged in MFMA operand order. A run's keys load nontemporal, long rows' runs split evenly, and every run stages its scores in LDS and stores them as pairs after each batch. - A lone live row (2 rows: a request and its pad row) is scored by all 8 waves. The scores and the selection are unchanged, bit for bit.
…registers With two rows a wave, the scan's batch path spent about 30% of each 8-page batch on its LDS copy and barriers. From 4 rows each wave of a row set now takes every fourth 16-key group of the CTA's run and loads its lane's words straight into registers, as MFMA B operands: four 8-byte loads and the scale, two groups ahead. A group's registers take the group two on once its MFMAs have read them; its scale is moved out first, else LLVM copies the refill and waits for it. The scores are stored as pairs from the registers, and the store's mask drops keys at or past a row's bound. No LDS copy of the run, no batch barriers and no score stage remain. A key this launch writes is taken from its row's index_k8 pairs, for every row of the request in the launch, as the batch path does. The first head term is an fma with 0, which rounds as 0 + product. Below 4 rows the batch path reads all its block ids before its first load, and its loads are unconditional: words past the run load the cache's first 16 bytes, which are never read back. The scores and the selection are unchanged, bit for bit.
At the scoring stage each scoring CTA first worked out which rows score together (each row's index bound, and its page-table entry against the previous row's slot), a chain of dependent loads ahead of its query's poll. Above 4 rows, with the paged indexer, the CTAs past the q_norm ones now do this once their front tile is done, from scalar loads of each row's slot, position and page-table row, and keep the layout in LDS ahead of the arena: the row groups, their requests' first rows, and each row's index bound and page-table row. The scoring stage reads it from there. Up to 4 rows the scoring stage still computes it itself; with it ahead, those builds' registers come out worse. The scores and the selection are unchanged, bit for bit.
…runs At long context a selection slice's first radix digit took a local pass over its 8,192 keys, and most of their scores fall in one or two of its 64 bins, so its atomics serialized. From 4 rows the scoring now builds that histogram: each scoring wave adds the first digit (radix_key >> 26) of every score it stores to LDS bins per row and slice segment. A request's rows join a group only at the next position, so a run spans at most two of each row's slices, and the run publishes each live row's two 64-bin histograms in the new score_hist mailbox (1 MB at 8 rows). A slice polls its runs' segments, which it finds with one run lookup and a compare, and publishes their sum as its first-digit histogram before reading its scores, so its first exchange skips the local pass. Its first chunk of scores loads while the segments are summed. Above 4 rows the scoring layout is planned at the front by the CTAs that score; the q_norm CTAs, which score nothing, plan it at the scoring stage, since their selection slices read it. The digits, counts and placement are unchanged, so the selection is unchanged, bit for bit.
heslami
force-pushed
the
perf/glm53-index-longctx-stack
branch
from
November 3, 2026 20:10
150e395 to
4e6488c
Compare
Contributor
Author
|
Rebased on main now that #85 has landed. The diff is this PR's four commits, and the GLM-5.3 code is unchanged from the build measured below. Ready for review. Ran AgentX on MI355X with AMD's InferenceX recipe (TP8, MTP, InferenceX trace order, 1 hour per point), the build before #85 vs #85 + this PR on the same node. Decode ITL drops 10–12% at c1, c2 and c4, P90 interactivity goes up 19% at c1 and 33% at c2, and TTFT is unchanged. The table is in the PR description. With #85, this PR, sgl-project/sglang#47324 and sgl-project/sglang#48224, c1 decode ITL is 2.14 ms, within 3% of TileRT's published 2.08 ms on InferenceX (their run uses 16 GPUs with separate prefill, so only decode ITL compares). |
4 of 5 tasks
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
In long-context decode the indexer launch is what grows with context. It scores every cached key of each row and then selects each row's exact top-2048, and it runs on 18 of a step's 75 MoE layer launches, on layer 0, on the dense layers 1–2 and in every NextN draft step. In AgentX-style serving most decode steps sit at 256K keys and above. There the scoring scan and the selection's first radix digit dominate the launch:
This PR shortens all four. Every output, every write, every score and every selection stays bitwise equal to the base PR, at TP8 and TP4, every width, on the bf16 and E4M3 MLA cache, at 8K to 512K keys. The one addition is a new scratch mailbox,
score_hist.Result on MI355X, the base PR → this PR (TP8, FP8 E4M3 KV, MTP 3-1-4 at the golden accept length 2.99), measured on the same commits in an earlier stack, whose kernel files are these PRs' except the base PR's bounded code warm (Speed Tests):
Modifications
Four commits, one per change,
kernel.pyandlayout.py. The head is bitwise against the base PR.index_k8pairs, for every row of the request in the launch. Below 4 rows the batch path reads all its block ids before its first load and its loads are unconditionalradix_key >> 26) of every score it stores to LDS bins per row and slice segment. A request's rows join a group only at the next position, so a run spans at most two of each row's slices. The run publishes each live row's two 64-bin histograms in the newscore_histmailbox. A selection slice polls its runs' segments (one run lookup and a compare), publishes their sum as its first-digit histogram before reading its scores, and its first exchange skips the local pass; its first chunk of scores loads while the segments are summed. Above 4 rows the q_norm CTAs, which score nothing, plan the scoring layout at the scoring stage, since their selection slices read itscore_histmailboxscore_histholds, per row and scoring CTA, two 64-bin histograms as tagged pairs: 1 MB at 8 rows, 512 KB at 4. Above 4 rows the scoring layout takes(5 + 2 S + 3) // 4 * 4LDS words ahead of the arena.score_group): from 4 rows the paired path replaces the batched one; below 4 rows the batch path stays, with its block ids read first. The new-key patch keeps [Perf] [MI355X] glm53: Score a request's rows together and select the exact top-2048 on up to 64 slices #74's rule: every new key of the request's rows in the launch, from its first row on.score_groupsmoves to the launch level and takes the rows' bounds and page-table lookups as arguments, so the front and the scoring stage share it. From 4 rows a row joins its group only at the next position.index_slice): from 4 rows the first digit's local pass is replaced by the sum of the slice's runs' segments. The later digits, the counts and the placement are unchanged.Accuracy Tests
The checks, on MI355X under torchrun on 8 GPUs (TP8) and 4 GPUs (TP4), each against the base PR:
shared:3 index:3 model:8 dense:3 embed:3and each kv inbf16 fp8,torchrun --nproc-per-node 8 test/kernels/glm53/check_layer.py --kind K --launches N --rows 2,4,8 --rowsets p0,p1,v4 --kv KV --dump DIR/<tree>/tp8(and--nproc-per-node 4intotp4); thenpython3 test/kernels/glm53/check_layer.py --compare DIR/base/tp8 DIR/pr/tp8. With [Fix] common: Keep a PeerBuffer's uncached memory allocated after close #76'speer.pyapplied, 0 of 33,552 digests differ at TP8 and 0 of 1,944 at TP4. Without it, only a few 2-row configs differ, which is the known 2-row run-to-run variation that [Fix] common: Keep a PeerBuffer's uncached memory allocated after close #76 fixes (TP8: bf16 model and fp8 dense, index and embed; TP4: bf16 model and fp8 embed).torchrun --nproc-per-node 8 test/kernels/glm53/check_topk.py --pos 8447,131071,262143,524287 --rows 2,4,8 --dump DIR/topk/<tree>/tp8on the base PR and on this PR (and--nproc-per-node 4), thencheck_layer.py --compareof the two digest sets. The selections are exact on both trees, on both caches and every rowset. Scores and selections are bitwise: 0 of 41,984 digests differ at TP8 and 0 of 20,992 at TP4. Thecheck_layer.pydump set stays under 9K keys, so this is the check of the multi-batch runs, the even split, the direct loads and the first digits.torchrun --nproc-per-node 8 test/kernels/glm53/check_nextn.py --rows 2,4,8 --kv fp8 --only entry(and--kv bf16, and--nproc-per-node 4). Passes at TP8 and TP4 on both caches.The stack has 8 wrong answers out of 480, against 4 for each of the others (Fisher p = 0.39; against the build before #85, paired by problem, p = 0.36). Samples are independent draws. The reference served 16 requests at a time, the monokernel builds 2.
The stack against the build before #85: KL +1.9% (95% CI −1.0 to +5.0), top-1 −0.14 points, both within the spread between two servers of the same build (up to 4% in KL and 0.25 points in top-1). With MTP on at 64K: KL −5.7% (−12.8 to +2.1).
Speed Tests and Profiling
Setup as #75: image
lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025(squashfssha256:fb07e0ef78fe) with SGLang'spython/sglangoverlaid at5044779f, AITER 285579b6, FlyDSL 0.4.0,zai-org/GLM-5.3atba32f8c1, the e2e test's server flags with--kv-cache-dtype fp8_e4m3.TP8, MTP 3-1-4 at golden AL 2.99, FP8 E4M3 KV, cycle in ms:
1b93ec8, this PR's as522a7ed): golden AL 2.99, one server per arm, the cycle at two requests the last-prefilled request's. The gain grows with context: the scoring scan and the first digit scale with the keys a row reads.Kernel-only, on this branch (
time_layers.py --kv fp8 --similar 1.0undertorchrun --nproc-per-node 8, E4M3 cache, µs a launch, the mean of 2 rounds in alternating order, each round the median of 7 trials of 50 replays; the rounds agree within 0.3 µs).modelis the mean over the 75 MoE launches in model order, run one width a process:torchrun --nproc-per-node 8 test/kernels/glm53/time_layers.py --kind {index,dense} --rows 2,4,8 --pos 524287 --kv fp8 --similar 1.0,--kind model --rows 4and--rows 8(one width a process), and the same at--pos 131071and--pos 8447with--rows 4,8.AgentX (InferenceX
inferencex-agentx-mvp, seed 42, traces in InferenceX order) with AMD's InferenceX recipe: TP8, FP8 KV, MTP 3-1-4 at simulated accept length 2.99, one hour per point. The build before #85 and the stack ran back to back on the same node.Registers and lint
From the code objects of the spec's builds (
stop=True):/opt/rocm/llvm/bin/llvm-readelf --notes <key>.co, fields.vgpr_count,.vgpr_spill_count,.sgpr_count,.sgpr_spill_countand.private_segment_fixed_size. Two figures a cell are bf16 / E4M3 MLA cache.the base PR → this PR:
Lint: #84's
test/kernels/test_kernarg_lint.py(the GLM-5.3 index and shared launches at TP4 and TP8, widths 2, 4 and 8, both caches): 24 / 24 pass.Reviewer notes
index_k8pairs for every row of the request in the launch from its first row on (c0), not only the rows of the scoring group. That matters when one request has more rows than a group (an 8-row draft chain: two groups of 4). The direct path keeps the batch path's rule, over the S rows (new_bounds), and the scoring layout of commit 3 carries the requests' first rows (firsts) for it.check_topk.py --rowsets chainat 8 rows covers it.bounds[r] == bounds[r - 1] + 1), so a run spans at most two of each row's selection slices. A request's rows in an MTP verify step or a draft chain are always consecutive; other rows start a new group, which keepsfirstsfor the new keys.score_hist. A scoring CTA publishes, per live row of its group, the 64 first-digit counts of each of its two slice segments as tagged pairs. A slice sums the segments of the runs over its key range; the runs' layout it recomputes is the scoring's own (score_groups,select_slices, the even split).kernel.pygrows by 317 lines (+587 / −270);layout.py+3.Checklist
pre-commit run --all-filespasses. (ruff, ruff-format, isort and codespell are clean on the touched files)test/cover the change. (check_topk.pyand the kernel tests).github/CODEOWNERSthat names its owner. (No new model.)