Repository navigation
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
heslami
requested review from
Fridge003,
Qiaolin-Yu,
Ying1123,
hanming-lu,
hnyls2002,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
October 29, 2026 20:57
heslami
requested review from
BBuf,
DarkSharpness,
HaiShaw,
HydraQYH,
alphabetc1,
celve,
huangtingwei9988 and
yuan-luo
as code owners
October 29, 2026 20:57
This was referenced Nov 2, 2026
jhinpan
pushed a commit
to sgl-project/sgl-amd-mono
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>
This was referenced Nov 3, 2026
jhinpan
pushed a commit
to sgl-project/sgl-amd-mono
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>
added 2 commits
November 3, 2026 14:19
draft_topk1_postprocess reduced each row in a split pass and a finalize launch. Each split now folds its (value, lowest index) into a per-row 64-bit key with an atomic max, and the last split of a row to arrive writes the token, the draft token column and the position, and resets the row's state: one launch per top-1, three a GLM-5.3 MTP cycle fewer. Ties still resolve to the lowest index and NaN still reads as -1e30. draft_topk1_argmax_only no longer advances a scratch positions buffer (positions=None).
A greedy EAGLE verify with topk 1 filled the predicts and accept index, reduced the logits rows with torch.argmax and walked the chains in verify_tree_greedy, then broadcast the three results on HIP. greedy_chain_verify does the reduction and the walk in one launch: each CTA folds a vocabulary split's (value, lowest index) into its row's key with an atomic max, and the last CTA to arrive writes the target tokens and walks every request's chain, writing every predict and accept index entry. Its results are views of one buffer, which the HIP greedy branch broadcasts once. Outputs equal row_argmax then verify_tree_greedy's (kernel test).
added 7 commits
November 3, 2026 14:20
… relayed lengths and the draft extend layout A top-1 chain's verify ran three small ops for values the launches around them already hold: - The accept lengths (num_correct_drafts + 1): greedy_chain_verify and simulated_accept now also write them (accept_lens=). - The overlap scheduler's relay of the new lengths (FutureMap.publish's index_put): the commit launch writes them into the relay (prepare_verify_commit_outputs(relay=...)). The EAGLE worker offers FutureRelayBuffers for the FutureMap to adopt, as DSpark's verify does, and marks the lengths published. - The draft extend's layout (prepare_draft_extend_layout): build_chain_tree writes it over the same window (with_extend_layout=True), carried on the verify input to prepare_for_draft_extend. Three fewer launches a decode step; every output equals the op it replaces.
…olumn 0 With more than one draft step, the chain's column 0 (the step's input tokens) was copied in before the steps. The first step's top-1 launch now writes it (draft_topk1_postprocess(first_tokens=...)), one copy fewer a decode step.
…-planned inputs The verify input's tokens and positions are loaded into the graph's buffers when the verify batch is planned, which run_eagle_verify starts only after the draft has written them. The replay then copied the same tensors again. Like DSpark's pre-planned verify (SGLANG_OPT_DSV41_DSPARK_SKIP_PREFILLED_INPUTS), EAGLE's skips that copy when the replayed batch holds the planned tensors: two launches a step.
The draft decode graph ended with positions.sub_ to undo the steps' advance of its positions buffer. Every replay loads the positions again before it runs, so the restore is one launch a step for nothing.
Two kinds of device copies a decode step that read or write nothing new: - Each step of EAGLE's draft graph ran the eager runner's input load, which copied the batch's tensors into the eager buffers, though the draft graph's steps read its own capture buffers at fixed addresses. The draft graph's body now runs its steps under eager_inputs_in_place(), so their load passes the tensors through, as SGLANG_EAGER_INPUT_NO_COPY does. Every other eager forward, captured or not, copies as before. - The draft graph runners copied the hidden states and the DSA seed with copy_() apart from the grouped small copy. Fields up to MAX_FIELD_NUMEL elements (64K, raised from 8K) now ride in the grouped launch; larger ones keep the DMA copy.
…s KV-length copies only when allocating Two device ops a decode step that no one reads: - FutureMap.publish records an event every step, but with the GPU-only relayed lengths (HIP, no backend that needs the CPU lengths) the next step gathers them in forward-stream order and never waits on it. Only a mixed chunk's tails and the confidence relay do, so the scheduler keeps the event for those alone (FutureMap.record_publish). - eagle_prepare_for_decode copies the current and next KV lengths to the device every step, though only an allocation reads them; with whole pages most steps allocate nothing. The copies now go with the allocation.
SGLANG_SIMULATE_ACC_LEN benchmarks on HIP built the simulated accept index with about a dozen torch ops a verify step (fill, arange, index writes, copies), because the fused simulated_accept kernel was gated to CUDA. It is Triton, so it now runs on HIP too. A kernel test checks it against those torch ops.
heslami
force-pushed
the
up/spec-eagle-top1-fewer-launches
branch
from
November 3, 2026 23:13
a8b80c2 to
bfb2243
Compare
This branch has not been deployed
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
With
topk=1, EAGLE drafts a chain, and a greedy verify only needs each row's argmax. At one or two requests, the time between a decode step's graphs is mostly small launches, and each one sits on the critical path. Many of them compute values that a neighbouring launch already holds or could write:sampling_argmax(torch.argmaxon HIP), walks the chains inverify_tree_greedy, and on HIP broadcasts three tensors;num_correct_drafts + 1, the relay'sindex_putof the new lengths andprepare_draft_extend_layoutare one launch each;This PR removes those launches. Each change is on by default for the batches it applies to; every other batch keeps the current code.
Modifications
topk1.py)torch.argmax.chain_verify.py,eagle_utils.py)greedy_chain_verifyreduces every row the same way, and its last CTA walks every request's chain. The results are views of one buffer, so the HIP broadcast is one collective. Used for greedy top-1 trees on contiguous f32 GPU logits.greedy_chain_verifyandsimulated_acceptwriteaccept_lens. The commit launch writes the new lengths into the overlap relay: the worker offersFutureRelayBuffers, as DSpark's verify does.build_chain_treewrites the draft extend's layout.topk1.py,eagle_worker_v2.py)decode_cuda_graph_runner.py)eagle_draft_cuda_graph_runner.py)eager_runner.py,small_copy.py, draft graph runners)eager_inputs_in_place(): only the draft graph's body sets it, and inside it the eager input load passes its tensors through. Every other eager forward copies as before. Fields up to 64K elements (was 8K) ride in the grouped small copy, which takes in the hidden states and the seed.scheduler.py,overlap_utils.py,eagle_utils.py,allocation.py)spec_utils.py)simulated_acceptwas gated to CUDA; HIP ran about a dozen torch ops. This affectsSGLANG_SIMULATE_ACC_LENbenchmarks only.What keeps its current code: sampling, rejection sampling,
topk > 1, non-f32 or strided logits, CPU, the post-verify-graph path, CUDA for commit 8, and every eager forward outside EAGLE's draft graph.Tests:
test_spec_topk1.py: ties across splits, -0.0, relaunch, 1,500 rows, column 0;test_chain_verify.py: against the fills +row_argmax+verify_tree_greedypath, with 0–3 accepted drafts, bf16 ties, batch sizes 1–16, three vocabularies, widths 2 and 4; negative control: a target row's maximum moved;test_simulated_accept.py: against the torch ops it replaces.test_next_verify_prep.py: the relay write againstpublish'sindex_put, and the layout againstprepare_draft_extend_layout;test_small_copy.py: the 64K limit, and hidden rows plus a seed with index fields, bitwise.test_eager_inputs_in_place.py: the input load passes through only inside the context, and still copies inside a capture without it.Related:
row_argmax. Here the reduction and the chain walk are one launch.build_chain_treewrites the layout by default for eligible batches.Accuracy Tests
Every replaced op is bitwise equal to the op it replaces (kernel tests above). Each removed copy writes what its destination already holds, and each moved copy is byte-for-byte the same.
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).
test_simulated_accept,test_gather_spec_extras,test_spec_topk1,test_chain_verify,test_next_verify_prep,test_small_copy): 234 passed.Speed Tests and Profiling
Setup:
SGLANG_SIMULATE_ACC_LEN).--enable-amd-monokernel). The model's forward ran in the GLM-5.3 monokernels of a test build of the sgl-amd-mono plugin that uses [Feat] [AMD] monokernels: Serve a draft spec's own steps, and skip the DSA replay metadata they never read #47102's step-mode names. This PR changes the host and glue work between those launches, so its share of the cycle depends on how fast the forward itself is.AgentX (InferenceX
inferencex-agentx-mvp, seed 42, traces in InferenceX order) with AMD's InferenceX recipe: TP8, FP8 KV, MTP 3-1-4 at a forced accept length of 2.99, concurrency 1, one hour per arm, all arms on the same node. The SGLang base is main + #47101 + #47102; with the merged sgl-amd-mono plugin those two change nothing on these steps (decode cycle +0.24%, 95% CI −0.32 to +0.63).Earlier development measurements on the same stack, rank 0 (the per-commit op counts are in the table above):
Checklist
CI States
Latest PR Test (Base): ❌ Run #41737465832
Latest PR Test (Extra): ❌ Run #41737465691
Latest PR Test (AMD ROCm 10): ❌ Run #41737465696
Tip
Base CI only covers the core cases, and its workload may not exercise your change. Use
/rerun-testto run the tests your change touches, or copy this prompt for an agent: