Repository navigation
[Perf] [MI355X] glm53: Shared layer and MoE hand-offs, NextN entry, K-split front, indexer launch, MoE expert dedup and weight loads, kernarg reload, LM-head loads - #85
Merged
jhinpan merged 25 commits intoNov 3, 2026
Conversation
heslami
requested review from
HaiShaw,
alexnails,
jhinpan and
kevin-mii
as code owners
October 29, 2026 19:02
This was referenced Oct 29, 2026
heslami
force-pushed
the
perf/glm53-stack-shared-index-reload
branch
3 times, most recently
from
October 30, 2026 16:51
8777795 to
51e916e
Compare
added 21 commits
November 1, 2026 15:29
Every down task polled all of mid as tagged f32 pairs from the scratch and quantized it itself: 256 CTAs reading the same 74 KB (4 rows) or 147 KB (8 rows). Now wave k of down task t quantizes block t + k x 256 (one wave at TP8; two at TP4 and 8 rows, where mid has 288 blocks) and publishes its FP8 bytes and scale (midq, midqs); every down task polls those, a quarter of mid's bytes. The quantization is the same, so every value is unchanged. MXFP4 experts keep their own path.
…and uk In a launch without the indexer the selection is the previous launch's, so a split tile's keys and their KV latent and k_pe cache rows are known at the launch start. Above 4 rows most split tiles share a CTA with a q_norm, q_b or uk task; their rows now go out after the CTA's q_norm and cache tasks, ahead of its q_b or uk task, and the tile reads them where it read the rows it loaded itself. The indexer launches, TP4 (more split tiles than CTAs) and 2 and 4 rows keep the old order. Every value is unchanged.
…at row's first picks Every CTA routed all four rows (a wave a row, 8 top-8 rounds each) before it issued its first up/gate job. With one up/gate task round (TP8), the CTA with tiles c of slot group g now runs row g % 4 at slots 1 + g / 4, 3 + g / 4, 5 + g / 4 and 7 + g / 4: the same (row, slot, c) tiles as before, on other CTAs, with each tile's chunk (and so its XCD) unchanged. Routing wave w takes row (row + w) % 4, so wave 0 routes this CTA's own row; the routing waves run two rounds and keep their state in LDS, the CTA issues the first job's weights, and the other six rounds, the route weights and the key stores run under them. These rounds take the wave maximum of the lane heads into every lane (DPP and cross-row swaps), and the lane holding it records its expert ID; a round whose maximum is in more than one lane hands the row to route_top8, which breaks the tie by ID as before. The picks, the route weights, mid and down are unchanged. The router CTAs issue their shared tile's weights where they run it. Other widths, TP4, dense layers and the NextN indexer launch keep their code.
… tiles Each of a NextN launch's 48 entry tiles (TP8) runs as four K parts of [enorm(e) | hnorm(h)] on 192 CTAs instead of two halves on 96: a part reads 98 KB of eh_proj, stages its half's normalized rows and sends its f32 partial to the first part, which adds them in part order and pushes the bf16 tile. With the indexer the entry starts past the front's index projection tiles (wrapping around the grid), so their weights go out at the launch start. TP4 keeps the two halves (192 tasks already). Registers: the entry stages 4 rows at a time, and at up to 4 rows a NextN layer polls the pushed rows two at a time in its front tiles. The other kinds' code is unchanged.
On MI355X, MTP at two requests at TP4 served as one 8-row launch is 14% a cycle faster than SGLang's forward, so TP4 now routes up to 8 rows as TP8 does: the target's decode and verify steps and the draft extend of two requests. The row cap per TP size goes; no code object changes. The kernel tests run TP4 at 8 rows too.
…nt at each stage FlyDSL 0.4.0's LLVM miscompiles some NextN builds: the SGPR greedy allocator's live-range splitting spills a tuple of pointer arguments with only some of its lanes defined, and later restores read the lanes that were never stored, so a stage addresses one argument's buffer with another's words. Each stage after the NextN entry now loads the pointers it reads from the kernarg segment, behind an empty asm on the segment's address so the loads stay in their stage. No pointer argument then lives in registers across the launch; only the segment's address does. Other kinds build the same code as before.
--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 and runs a step of its NextN launch.
… every CTA Not bitwise. Each q_a / kv_a tile's first 4096 K run on its own CTA, the remaining 2048 on one of the 92 CTAs past the tiles (two tiles' rests a CTA), so every CTA streams at most 64 KB of weights instead of 96 KB on 164 CTAs. Neither part normalizes its input rows: each multiplies bf16(x * gamma), and q_norm and the cache tasks add the two partial sums (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 projections keep their whole-K tiles, on the CTAs with one tile's rest (the NextN entry starts past them); the MoE stages keep their CTAs. A layer from pushed rows (layer 0, the NextN layer) reads them in both parts and in the rstd, its weights issued before the rows' wait. The front asserts FP8 block-128 q_a / kv_a weights.
Each of the 8 rounds took two wave reductions: the key's maximum, then the smallest ID among the lanes that hold it. When one lane holds the maximum (all but exactly equal f32 keys), its head is the pick: a ballot and a lane read replace the second reduction. Same picks in the same order.
… the rstd Not bitwise (the router logits only). A router task staged the RMS-normalized row: the sum of squares' block reduction, two barriers, then the scaled stores. It now stores bf16(a * gamma) as the rows arrive, with each wave's sum of squares beside them, and one barrier serves both: the GEMV reads the row, and the rstd (summed in the same order, so the activation blocks keep their bits) scales the FP8 activation blocks and the 8 logits.
At 8 rows a scoring CTA's run is two requests' rows, and at long context its scoring is per-CTA compute: each wave read and converted every key for one row. Above 4 rows a wave now runs two rows' MFMAs on one LDS read and conversion of each key, a run's keys load nontemporal, the runs split evenly once every CTA has one (they differ by at most a page), and a multi-batch run stages a batch's scores in LDS and stores them after the next batch's keys land, so the wait for those keys does not wait for these stores. A one-batch run stores its scores as it goes. At 2 and 4 rows (and so at TP4) the scoring keeps its code. Above 4 rows the index-weights tiles also keep a quarter of their weights in flight instead of half: with half, those tiles spilled 20 VGPRs at 8 rows and ended 4.5 us after the other front tiles. A timeline build loads the timeline's address once a launch instead of at each mark, where the load waited for every load in flight before it.
…m CTAs With the indexer, the front's 10 index K and weights projection tiles shared their CTAs with a q_a / kv_a tile's K rest, so those CTAs streamed 8192 K against 4096 elsewhere, their front ended last, and the index-query tiles and the scoring waited on them. The index tiles now 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. From about 32K keys every CTA scored a run, and the runs on the q_norm CTAs, behind q_norm, an index-query tile, the index key and the cache task, finished last. The scoring runs now go to CTAs S .. G - 1 only, counted from the same first CTA. Only the indexer builds change (index, dense, embed, NextN index).
…e selection's code With the indexer, a split tile landed its selection and then spent about 3 us gathering its KV rows and polling q. It now polls q and the rows this launch appends before the selection's wait, staging the rows in the tile's reduction and output regions (q is staged by then). Then each wave polls its own 8 selected keys and issues their cache rows as they land, and copies the staged rows to the keys that are this launch's. Without the indexer, split keeps its order. The selection's steps run once a launch and wait on instruction fetch: their time tracks code bytes, not work. While its scores are pending, each selection slice loads the next 64 KB of code after it (one word of each 64-byte line, from s_getpc) into its XCD's L2. TP8 dense launches above 4 rows skip it, where it costs more than it saves. Only the indexer builds change (index, dense, embed, NextN index).
Above 4 rows each scoring wave loads two rows' query fragments once a run, and it read them from the row-major LDS copy: a fragment's 16 lanes read 16 heads 64 words apart, all on one bank. The rows' queries are now staged in MFMA operand order (head half, 32-dim chunk, lane), so each fragment is one conflict-free 16-byte read a lane. The values are the same. At 2 and 4 rows the scoring is unchanged.
…ode end warm_code took its buffer from s_getpc with a fixed 64 KB size, so in a code object whose code ends less than 64 KB after the call (the dense and embed launches) it read past the end, and on MI355X TP4 dense launches faulted with an illegal address. The buffer's size is now the assembler's distance from the s_getpc to .Lfunc_end0, from one inline asm, capped at the requested bytes; loads past it read nothing. Where 64 KB fits, the read is unchanged.
… down At 8 rows up/gate ran each routed slot's tiles once a row, so an expert two rows picked (the shared one always) streamed once a row, and the CTAs holding the shared expert's tiles ran a tile more than the others. Its jobs are now the 8-intermediate tiles of the shared expert and of each distinct routed expert (in order of first use, found by every wave from the top-8s), each with every row that picked it 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 streams each expert once too: its units multiply the tile with every using row and store the products in LDS, and each row's output is summed over its units in the order and wave segments of the one-row form. The rows the multi-column MFMAs read (the activations and the FP8 mid blocks) are padded by 4 words. A column's sums are those of a one-row tile, so mid and the output keep their bits. Up to 4 rows keep the one-row forms, which measured faster there, and dense layers keep theirs.
…tep's expert groups Down's units past the step's expert groups load no weights (their buffer is empty) and store no products, but their MFMAs still ran. A wave's weight batch whose first unit is past them now skips its MFMAs: at 8 rows a step uses about 43 of the 65 possible groups, so the last batches mostly have none. The products and sums are unchanged.
…een lane rows where read ld_fp8_chunk read a lane's 16 bytes of AITER's (16, 16) layout as two 8-byte loads 512 bytes apart. Now a lane of an even 16-lane row loads the whole 16-byte segment of the first K32 block and its partner (lane + 16) the same segment of the second, so every FP8 weight load is one dwordx4 instead of two dwordx2. Two v_permlane16_swap give each lane its halves; they run where an element is read, so the load stays in flight across the wait between issue and use. The MFMA inputs are the same bits.
…routing alike With --similar X the live rows of a request are one random row plus X times their own, so that they pick many of the same experts, as a verify step's rows do. The default (0) keeps the independent random rows and their digests.
…le once At 4 rows (TP8) a CTA's up/gate jobs are one row's, so an expert that several rows picked streamed its tile once per picking row. Each job now takes every row as an MFMA column. The job that comes first among the picking rows' jobs of the expert (the lowest row on a tie) owns it and writes mid for every row that picked it; the other rows' later jobs of it load nothing. A first job is issued before the other rows' picks are known, so it always loads. A column's sums are those of a one-row job, so mid keeps its bits.
The FP8 weight loads take a cache policy (ld_fp8_chunk and the unit helpers' cm). At 2 and 8 rows the attention weights (q_a | kv_a, q_b, W_UK, o_proj) and the experts' up/gate and down load nontemporal, so streaming them does not evict lines that are read again. At 4 rows only up/gate does, and only where a job streams each expert once for every row that picked it (TP8); down's 4-row tiles read the shared expert and repeated experts once a row, and those repeats hit in cache under the default policy. The loaded bits are the same.
added 4 commits
November 1, 2026 15:29
…he kernarg segment at each stage As in NextN launches, each stage of an indexer or dense-MLP launch (layers 0 to 2) now loads the pointers it reads from the kernarg segment, behind an empty asm on the segment's address. No pointer argument then lives in registers across the launch, and these kinds spill far fewer SGPRs. Shared, tail and NextN launches build as before.
A step's output rows were digested only after each layer launch, so a dump of --kind head (or the tail of --kind full) held no digest of the tail's output. Each step with the tail now records its gathered logits a row, so --compare checks the tail bitwise.
…ows' loads ahead of them A wave replaces each unit it multiplies by the next one of its tile sequence, so the next tile's loads are in flight while the current tile reduces and pushes; the rows' loads for the norm go out before the first weights.
…rd into it The up/gate tile helpers of the 8-row schedule and of the 2- and 4-row one were two copies of one definition; there is now one, ahead of the schedules. reload_args unpacks its loads in signature order and returns at once when a kind keeps its arguments in registers, so its call sites need no guard. The block sums and maxima share one exchange, and the q_b and extra CTAs' index-query tiles one unit maker. Every launch builds the same code.
heslami
force-pushed
the
perf/glm53-stack-shared-index-reload
branch
from
November 1, 2026 23:30
51e916e to
b7ae7cb
Compare
4 of 5 tasks
Collaborator
|
Hi @xiaobochen-amd. Can u help cross-review this PR about monokernels on GLM 5.3? Thx! |
@jhinpan LGTM. I re-ran this PR against main ( Speed: ms per verify round, the median of two runs after a warmup.
The gain holds from 1K to 512K: 6–9% with one request and 16–18% with two. Accuracy
|
This was referenced Nov 3, 2026
jhinpan
pushed a commit
that referenced
this pull request
Nov 3, 2026
…ows, keys straight into registers, the scoring layout ahead, first digits from the scoring (#86) ## 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: - **The scan is bound by per-wave compute, not by bytes in flight.** One 8-page batch a CTA in flight already keeps the key loads streaming; the scoring waves spend the time in address math, FP8 → bf16 conversion, the head-half reduction and LDS traffic. At 4 rows (one request's MTP verify rows) each key was read and converted once per row. - **The scan copied every batch through LDS.** At 4 rows and more, each 8-page batch went through an LDS copy and barriers before any wave scored it. - **The scoring stage started with a chain of loads.** Each scoring CTA first worked out which rows score together (each row's bound, and its page-table entry against the previous row's slot) before it could poll its query. - **The selection's first digit re-read every score.** A selection slice of up to 8,192 keys built its first 64-bin histogram in a local pass, and most scores share the one or two bins of their common exponent, so the LDS atomics serialized. 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): - **8K:** 6.018 → 5.985 ms at one request (−0.5%), 8.479 → 8.443 ms at two (−0.4%). - **128K:** 6.313 → 6.134 ms at one request (−2.8%), 9.372 → 9.173 ms at two (−2.1%). - **512K, kernel-only on this branch:** the indexer launch at 4 / 8 rows 132.5 → 108.7 / 218.0 → 185.6 µs (−18% / −15%), the 75-launch model order 76.8 → 70.8 / 117.5 → 110.2 µs; about 0.52 ms a step less at one request (Speed Tests). ## Modifications Four commits, one per change, `kernel.py` and `layout.py`. The head is bitwise against the base PR. | # | Commit | Change | Numerics | |---|---|---|---| | 1 | two-row scoring from 4 rows | A scoring wave's 16-key group addresses come from its page and first key in scalar math. A row's two head halves reduce across lane rows in three permlane swaps, and two rows share one score store. From 4 rows (was: above 4) a wave scores two rows, reading and converting each key once for both, 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 | bitwise | | 2 | key groups straight into registers | From 4 rows each wave of a row set takes every fourth 16-key group of the CTA's run and loads its lane's words straight into registers as MFMA B operands, two groups ahead (four 8-byte loads and the scale a group). The scores are stored as pairs from the registers, the store's mask dropping keys at or past a row's bound. No LDS copy of the run, no batch barriers, no score stage; the direct loads are plain (commit 1's nontemporal loads go with the batch path). A key this launch writes is taken from its row's `index_k8` pairs, 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 unconditional | bitwise | | 3 | the scoring layout ahead | Above 4 rows, with the paged indexer, the CTAs past the q_norm ones work out the scoring layout once their front tile is done, from scalar loads of each row's slot, position and page-table row: the row groups, their requests' first rows, and each row's index bound and page-table row. They keep it in LDS ahead of the arena, and the scoring stage reads it there. Up to 4 rows the scoring stage still computes it | bitwise | | 4 | first digits from the scoring | From 4 rows each scoring wave adds the first radix 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. The run publishes each live row's two 64-bin histograms in the new `score_hist` mailbox. 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 it | bitwise; new `score_hist` mailbox | - **Layout:** `score_hist` holds, 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 * 4` LDS words ahead of the arena. - **What changes in #74's code:** - The scoring (`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 #74's rule: every new key of the request's rows in the launch, from its first row on. - `score_groups` moves 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. - The selection (`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. - **Unchanged:** the shared, NextN shared and tail launches' code objects. ## Accuracy Tests The checks, on MI355X under torchrun on 8 GPUs (TP8) and 4 GPUs (TP4), each against the base PR: - **Dumps, bitwise:** 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/base/tp8 DIR/pr/tp8`. With #76's `peer.py` applied, 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 #76 fixes (TP8: bf16 model and fp8 dense, index and embed; TP4: bf16 model and fp8 embed). - **Top-k exact at 8K to 512K, scores and selections bitwise:** `torchrun --nproc-per-node 8 test/kernels/glm53/check_topk.py --pos 8447,131071,262143,524287 --rows 2,4,8 --dump DIR/topk/<tree>/tp8` on the base PR and on this PR (and `--nproc-per-node 4`), then `check_layer.py --compare` of 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. The `check_layer.py` dump 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. - **NextN launch = plain launch fed its own x, bitwise:** `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. - 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 before #85 | 99.17% | 98.3–99.8 | 3.23 | | With #85, this PR, sgl-project/sglang#47324 and sgl-project/sglang#48224 | 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 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. - 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 before #85 | 0.0199 | 96.34% | | With #85, this PR, sgl-project/sglang#47324 and sgl-project/sglang#48224 | 0.0203 | 96.20% | 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` (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`. TP8, MTP 3-1-4 at golden AL 2.99, FP8 E4M3 KV, cycle in ms: | Context | Requests | the base PR | This PR | Change | | --- | --- | --- | --- | --- | | 8K | 1 / 2 | 6.018 / 8.479 | 5.985 / 8.443 | −0.5% / −0.4% | | 128K | 1 / 2 | 6.313 / 9.372 | 6.134 / 9.173 | −2.8% / −2.1% | - Measured on the same commits in an earlier stack (the base PR's head as `1b93ec8`, this PR's as `522a7ed`): 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.0` under `torchrun --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). `model` is the mean over the 75 MoE launches in model order, run one width a process: | Launch | Rows | 8K | 128K | 512K | | --- | --- | --- | --- | --- | | index | 2 | | | 96.85 → 90.97 (−5.9) | | index | 4 | 79.56 → 77.94 (−1.6) | 92.09 → 85.59 (−6.5) | 132.48 → 108.70 (−23.8) | | index | 8 | 111.78 → 108.41 (−3.4) | 147.81 → 140.51 (−7.3) | 217.98 → 185.56 (−32.4) | | dense | 2 | | | 89.83 → 83.41 (−6.4) | | dense | 4 | 63.08 → 62.09 (−1.0) | 75.53 → 69.66 (−5.9) | 115.72 → 92.02 (−23.7) | | dense | 8 | 79.52 → 75.80 (−3.7) | 113.53 → 105.53 (−8.0) | 181.71 → 149.72 (−32.0) | | model | 4 | 63.88 → 63.57 (−0.3) | 67.06 → 65.55 (−1.5) | 76.84 → 70.81 (−6.0) | | model | 8 | 92.10 → 91.80 (−0.3) | 101.22 → 99.67 (−1.6) | 117.50 → 110.19 (−7.3) | - From these, the 75 MoE launches and the 3 dense ones of a step take about 0.52 ms less at 512K and one request (4 rows), 0.13 ms at 128K and 0.03 ms at 8K. - Commands: `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 4` and `--rows 8` (one width a process), and the same at `--pos 131071` and `--pos 8447` with `--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. | 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 before #85 | 291.1 | 2,890 | 2.484 | 0.372 | | | With #85 and this PR | 345.3 (+18.6%) | 3,135 (+8.5%) | 2.207 (−11.2%) | 0.375 | | 2 | sgl-amd-mono before #85 | 212.2 | 2,911 | 2.491 | 0.318 | | | With #85 and this PR | 282.9 (+33.3%) | 3,160 (+8.5%) | 2.219 (−10.9%) | 0.318 | | 4 | sgl-amd-mono before #85 | 167.1 | 4,069 | 2.706 | 0.256 | | | With #85 and this PR | 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 before #85 (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. **the base PR → this PR:** - TP8 index, dense and NextN index at 8 rows lose their last VGPR spills (1–6 → 0). - TP4 at 8 rows: index 10 → 1 / 2, dense and embed 3–5 → 1. - Two builds gain one: TP4 index at 4 rows 0 → 1 (8 B), and TP8 NextN index at 4 rows, bf16 cache, 0 → 1. - The shared, NextN shared and tail objects are unchanged. **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 - **New keys, as #74 has them.** A key this launch writes is scored from its row's `index_k8` pairs 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 chain` at 8 rows covers it. - **The join rule.** From 4 rows a row joins its request's group only at the next position (`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 keeps `firsts` for 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). - **Two scoring paths remain.** Below 4 rows the batch path (one row a wave, LDS batches) stays; from 4 rows the paired direct path. A runtime switch by key count was not taken: both paths in one object grew it and cost at short context. - **The scoring layout ahead (commit 3)** applies above 4 rows only; at 4 rows it does not pay. It assumes the atom KV cache layout (asserted), as the paged indexer does. - **Commit 4's trade.** It moves the first digit's counting into the scoring CTAs: a small cost per scoring CTA at short context (8K), and the slice's local pass saved at long context. Commit 4 can be dropped alone. - `kernel.py` grows by 317 lines (+587 / −270); `layout.py` +3. ## 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_topk.py` and the kernel tests) - [x] The SGLang commit or PR this was tested against is named above. - [ ] Accuracy and speed results are attached when kernel or forward code changes. (Speed: serving at 8K and 128K, above.) - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
4 of 5 tasks
jhinpan
added a commit
to jhinpan/InferenceX-amd-sgl-mono
that referenced
this pull request
Nov 5, 2026
Pin lmsysorg/sglang-rocm:v0.5.23-rocm720-mi35x-20261105 by digest. It bundles SGLang 1305b1830a and sgl-amd-mono e52803f9d2, which contains sgl-project/sgl-amd-mono#85 and SemiAnalysisAI#86. 按 digest 固定 lmsysorg/sglang-rocm:v0.5.23-rocm720-mi35x-20261105。它打包了 SGLang 1305b1830a 和 sgl-amd-mono e52803f9d2，后者包含 sgl-project/sgl-amd-mono#85 和 SemiAnalysisAI#86。
jhinpan
added a commit
to jhinpan/sglang-fork
that referenced
this pull request
Nov 6, 2026
…he GLM-5.3 agentic recipe Section 3.6 now names lmsysorg/sglang-rocm:v0.5.23-rocm720-mi35x-20261105, the first MI35x image whose SGLang pins sgl-amd-mono e52803f9 (sgl-project#45202), which contains sgl-project/sgl-amd-mono#85 and sgl-project#86.
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
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:
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):
Modifications
24 commits, one per change, in the order below (#). Three only add test options.
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 pathKERNARG_I64lists 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, andxinfrom the reloadedsym, 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 beforetime_layers.py--kind nextn_indexand--kind nextn_sharedtime 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.Drafttakes the rows directlyq_a2,kv_a2hold 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' waits_getpc). TP8 dense launches above 4 rows skip that warmwarm_codetakes its PC and the distance to.Lfunc_end0from 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 unchangedld_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 onedwordx4._RowSwappedholds the raw 16 bytes; reading elementerunsv_permlane16_swapon 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--similarcheck_layer.pyandtime_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 digestsjob_emit, and the other rows' later jobs of it load nothingld_fp8_chunkand the FP8 unit helpers take a cache policy (CM_NTintiles.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 appliesRELOAD_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()rebindsxinonly where x is pushed. Shared-layer and tail launches build exactly as beforecheck_layer.py--dumpdigests the tail's gathered logitsHD_AHEADunits 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 weightsmidqandmidqsaftermid(19 KB at 4 rows, 38 KB at 8);ehpholds three NextN entry parts' partials instead of one;q_a2andkv_a2hold the front's K rests. The MoE group table and down's products live in the LDS X region; no other new scratch.TP_MAX_ROWSand_Built.max_rowsgo;DeviceLimits(widths=ROWS)bounds every routed step at 8 rows.test_glm53_kernels.pyruns TP4 at widths 2, 4 and 8. The README's GLM-5.3 row says up to 8 rows at both TP sizes.qkv_ahas 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.polltakes an optionalbetween()that runs after the first loads are issued and before they are checked (up/gate issues its first job's weights there).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 [Fix] common: Keep a PeerBuffer's uncached memory allocated after close #76's
peer.py: without it, [Feat] [MI355X] glm53: The NextN draft steps on their own launches #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.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/<parent>/tp8 DIR/<tree>/tp8. Commit 20 also with--rowsets v4 --similar 1.0. Negative control: the same dump of [Feat] [MI355X] glm53: The NextN draft steps on their own launches #75 with--perturb.--kind head, commit 23)The front, not bitwise (commits 8 and 10). Measured at TP8 with MTP on and real acceptance, the commit before the front ([Feat] [MI355X] glm53_flash: Cover layers 0-2, 44 and the logits #23's head) against the front's head, and [Feat] [MI355X] glm53_flash: Cover layers 0-2, 44 and the logits #23's head again as an A/A pair:
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.
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.
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).
test/e2e/test_glm53_server.py) at TP8 and at TP4 (--tp-size 4), on the same commits in an earlier stack: 8 passed.TestMonokernelFaultfails as it does on main.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.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:
TP4, two requests, 8K (cycle in ms):
Real acceptance, #75 → the front's head (commits 1–10), TP8, 8K, 40 waves, paired:
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.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.#75 → this PR, at 8 rows:
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
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 throughfreshand not used.owns(e)reads the step's picks after the routing barrier; a row's first job is issued before it and always loads.kernel.pygrows by 1,397 lines (+1,843 / −446).Checklist
pre-commit run --all-filespasses. (ruff, ruff-format, isort and codespell are clean on the touched files)test/cover the change. (check_layer.py --similar, the tail digest, the NextN kinds intime_layers.py; the kernel tests compare every launch).github/CODEOWNERSthat names its owner. (No new model.)