Repository navigation
[Perf] [MI355X] glm53: Shorter attention chain, one-pass 8-row front, cheaper routing hand-off (2/2) - #69
Merged
Conversation
This was referenced Oct 28, 2026
Closed
Closed
jhinpan
approved these changes
Oct 28, 2026
jhinpan
left a comment
jhinpan
left a comment
Collaborator
There was a problem hiding this comment.
I approve this PR. I tested it, with #68, on current main (90c2355). The merge has no conflicts. The glm53 kernel and layout files are the same as on this head.
Setup: the same as on #68.
-
Bitwise: All digests are equal to #68 and to
main: 41,760 at TP8 and 10,560 at TP4. -
Tests: The kernel tests pass (3 of 3). The e2e file at TP8 passes (8 of 8), with the fault test. GSM8K is 0.995 without MTP and 0.98 with MTP. The accept length is 3.25. KL is 0.034 or less.
-
Kernel time (TP8, E4M3, model order), µs per launch, #68 → this PR. The second percentage is #68 + this PR against
main.Context 2 rows 4 rows 8 rows 8k 60.19 → 54.95 (−8.7%; −15.4%) 78.82 → 72.08 (−8.6%; −16.4%) 127.78 → 117.15 (−8.3%; −14.0%) 256k 69.95 → 64.75 (−7.4%; −14.4%) 97.27 → 88.97 (−8.5%; −14.9%) 153.85 → 143.94 (−6.4%; −13.8%) -
AgentX (the same setup as on #68): Against #68, this PR decreases ITL by 6.3% / 5.6% / 5.3% at c1 / c2 / c4 on the same requests. With #68, against
main, ITL decreases by 11.0% / 11.2% at c1 / c2, and E2EL decreases by 9.1% / 9.7%. TTFT does not change. The gain stays at long context: ITL decreases by 11.9% below 64k and by 10.4% at 192k or more. At c4, the stack decreases ITL by only 4.9%. A probable cause: the prefills of other sessions run between the decode steps.
Not blocking: The TP4 4-row indexer builds spill 11 / 16 VGPRs, as your PR shows. GLM-5.3 runs at TP8 in the cookbook and in InferenceX. Thus, the net gain at TP4 is sufficient.
Not checked: the register survey. I did not build the 40 keys again.
This was referenced Oct 28, 2026
4 of 5 tasks
This was referenced Oct 28, 2026
added 6 commits
October 28, 2026 21:42
The stage placement is now one function, layout.stage_bases. One cache task normalized, roped and stored every row's KV latent (and index key) in turn, and the split tile behind it on its CTA started late; each row now has its own cache task, on its q_norm CTA after its q_norm task, each row's sums in their order. q_norm takes the first CTAs past qkv_a's, so it polls from the launch start, q_b the next ones, and split starts past q_b's, so its tiles land on CTAs qkv_a frees; o starts at split's first CTA. The MoE stages keep their CTAs.
… the q wait A split tile loaded its rows' slots and the rows this launch appends only after q arrived, and each wave loaded its 8 keys' cache rows one after another, then polled each new key's row in turn. The slots are now loaded at the task's start; every key's cache row is issued first; and if any key of the task is new, the CTA polls every new row once, behind the cache rows' loads, into the q region of LDS, from which each wave copies its keys' rows. The conversions and stores follow, all before the q wait. Same values.
On an indexer launch qkv_a tasks 0-7 also ran the index K projection's tile and tasks 0-1 the weights projection's: those CTAs ended about 7 us after every other front CTA and held back the cache task's index keys. These 10 tiles are now further tasks of the qkv_a stage, each on its own CTA, normalizing the input rows itself, with the same units in the same order. At 8 rows a front tile normalized its rows four at a time, issued its weights again for the second four, and ran its MFMAs, reduction and stores once per group. It now issues the weights once and normalizes all 8 rows into LDS before one pass. Above 4 rows the index tiles issue half their weights before staging the rows and half after, so the weights and 8 rows fit in registers; the MFMAs sum the units in the same order.
…routing The router published its FP8 activation blocks one tagged word (4 values) an 8-byte store, and every up/gate CTA polled each row's words and scales in batches, 2 or 3 polls in series. The blocks now go out two tagged words a 16-byte store, and a CTA polls every row's words and all the scales in one batch. The mailbox and its values are unchanged. A CTA that ran a router task reaches up/gate after the activations are out: it now routes first, issues its first routed tile's weights, and stages the activations behind them. The other CTAs stage them while the router still runs, as before.
From the router on, at 8 rows, the wave index is read once into a scalar register, so the per-wave unit, row-group and address terms of the router, up/gate and down are scalar arithmetic and the 8-row builds stop spilling the thread id to scratch. Below 8 rows it stays a vector value.
Up/gate and down loaded each unit's f32 block scale with a vector load at a wave-uniform address: a VGPR and a vector-memory instruction a unit, in the weights' in-order queue. They are now scalar loads from the constant address space (up/gate loads the gate and up rows' scales and picks one by lane at the MFMA), volatile so that they issue with their unit's weights instead of being sunk to the MFMA.
heslami
force-pushed
the
perf/glm53-attn-front
branch
from
October 29, 2026 04:59
cfba6f0 to
9f17c49
Compare
jhinpan
pushed a commit
that referenced
this pull request
Oct 29, 2026
…se (#76) ## Motivation `PeerBuffer.close()` frees its uncached buffer (`hipExtMallocWithFlags(..., hipDeviceMallocUncached)`) with `hipFree` and unmaps the peers' buffers. On MI355X (ROCm 7.2), once a launch has polled such a buffer, ordinary allocations that later take its physical pages can hold stale data: 64-byte runs of zeros, or of values the tensor held before. A process that closes a runtime and then builds new tensors can get wrong values. Only processes that close runtimes and keep running are exposed. Serving builds every runtime once and never calls `close()`. The callers are the glm53 and glm53_flash runners' `close()`: - glm53_flash: since #62 its scratch is a one-rank `PeerBuffer` too, and `check_layer.py` and `time_layers.py` close runtimes between checks; - glm53: frees its peer buffer, and with #68 its uncached scratch; - deepseek_v4, deepseek_v41, minimax_m3 and minimax_m31 never close theirs. A system-scope L2 writeback and invalidate on every CU (`buffer_wbl2 sc0 sc1`, `buffer_inv sc0 sc1`), before or after the free, doesn't prevent it. Streaming 64 MB through the cache after the free does, and ordinary or fine-grained buffers aren't affected. So the stale lines are in the L2s, and no free is safe in every case. ## Modifications - `common/peer.py`: `PeerBuffer.close()` no longer calls `hipIpcCloseMemHandle` and `hipFree`. It drops its view and keeps the local allocation and the peers' mappings in a module-level `_RETIRED` list until the process exits. It is still idempotent, and the handshake-failure path still calls it. - No caller changes: every runner releases its buffers through `PeerBuffer.close()`. - `test/unit/test_common_peer.py` (CPU): `close()` keeps the buffer and its peer mappings, calls no HIP function, and a second `close()` does nothing. With `main`'s `close()` it fails at the first HIP call (`hipIpcCloseMemHandle`). **Follow-up:** after #73 lands, a small PR adds the close-and-rebuild sequence of the negative control below to glm53's kernel tests. The hazard shows only some of the time, so that is a stress check, not a pass/fail gate. Cost: a closed runtime's uncached memory stays allocated until the process exits (2 MiB granules, per runtime): - glm53_flash (scratch and peer buffer): 4–26 MiB at TP2, 4–18 MiB at TP4, 6–18 MiB at TP8. The 8-row sparse-attention launch is the largest. - glm53 (peer buffer): 2–8 MiB at TP4, 4–14 MiB at TP8. With #68's uncached scratch: 8–48 MiB at TP4, 8–44 MiB at TP8. Only processes that close and rebuild runtimes pay this, i.e. the kernel checks. ## Accuracy Tests **Negative control, glm53 at #73, TP8.** It is #73's tree (not `main`), without and with this change: four repeats each, alternating, in one allocation. That tree carries #68's uncached scratch, and #73's `check_layer.py` provides the trigger. Its full dump sequence closes a runtime after each config and builds the next one's caches: ``` torchrun --nproc-per-node 8 test/kernels/glm53/check_layer.py --kind index --launches 3 \ --rows 2,4,8 --rowsets p0,p1,v4 --kv fp8 --dump OUT torchrun --nproc-per-node 8 test/kernels/glm53/check_layer.py --kind model --launches 8 \ --rows 2,4,8 --rowsets p0,p1,v4 --kv bf16 --dump OUT ``` The comparison is over the `digests_*_r<rank>.json` files the dumps write, 16,920 digests a run. | | without this change | with this change | | --- | --- | --- | | runs whose ranks disagree on a cache digest | 4 of 4: `bf16-model/w2p1`'s layer-3 index cache in all four, `fp8-index/w2v4`'s layer 0–2 MLA caches in three | 0 of 4 | | digests that differ from #69's `check_layer.py` on the same kernel (it builds the runtimes before the weights, and doesn't trigger it) | 308 in one run, 636 in each of the other three | 0 in every run | **1-GPU reproduction (out of tree).** One rank of the TP8 launch runs alone on one GPU, with the peers' parts pre-filled. Each config's caches as built are compared with an identical construction. Same code objects in every arm: - buffers freed at close: wrong in 32 of 34 runs; - buffers kept until the process exits (retired by the runner's `close()`, or never closed): wrong in 0 of 19 runs. The driver replaces `PeerBuffer` with its own buffers, so this form doesn't exercise `common/peer.py` itself; the TP8 control above does. The command uses the out-of-tree driver `nd_loop11.py` over #73's `check_layer.py`: ``` LOOPBACK=8 TOOLS=test/kernels/glm53 torchrun --nproc-per-node 1 nd_loop11.py \ --kind index --launches 3 --kv fp8 --rows 2 --close-barrier --verify-creation --out OUT ``` **Not a miscompile.** A check of every shipped build for the kernel-argument miscompile was clean. **glm53_flash on `main`.** Its tools don't reproduce the hazard. At TP2, 17 runs of 76 close-and-rebuild rounds in total, with no rank wrong: - DSA, FFN, dense and KDA launches; 2–8 rows; 3–30 launches; - the DSA caches built right after `close()`, a synchronize and a barrier, with 0–3 s sleeps; - each compared with a rebuild. With this change: - `test/kernels/test_glm53_flash_kernels.py`: 8 of 8 pass, at TP2, TP4 and TP8; - every launch's outputs and writes are bitwise equal across two runs, for every kind, FP8 layout and width: 996, 2,988 and 5,976 tensors at TP2, TP4 and TP8; - code objects, compile-only, every kind at TP2, TP4 and TP8, every width, both FP8 layouts, stop on and off (192 builds): - 190 are byte-identical to `main`'s; - the other two, the 2-row sparse-attention builds at TP2 and TP4 (shuffled FP8, stop on), vary from build to build on `main` itself: two variants each in five builds; - every object this head built for them equals one `main` built. CPU unit tests: the same result on `main` and this head apart from this PR's two new tests, which pass (and fail with `main`'s `close()`), on a CPU-only host with SGLang `ab580b2e06`: 260 passed, 3 skipped, and the same 8 failures and 3 collection errors (SGLang contract checks on that commit; AITER imports that need a GPU). Setup: MI355X. glm53 on `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025` (`sha256:fb07e0ef78fe`) with SGLang `5044779f`; glm53_flash on `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014` (`sha256:e4498f5d595f`) with SGLang `ab580b2e06`. ## Speed Tests and Profiling No launch or serving path changes, and serving never calls `close()`. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change: `test/unit/test_common_peer.py` checks that `close()` keeps the memory. The hazard itself is probabilistic: the negative control above needs #73's `check_layer.py`, and the follow-up adds it as a stress check. - [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. --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan
pushed a commit
that referenced
this pull request
Oct 29, 2026
## Motivation **Stacked on #69** (and #68 under it). Merge those first; this PR's diff is its last four commits, and it extends #68's `test/kernels/glm53/` tools. A routed decode step of GLM-5.3 runs layers 3 to 77 as launches (75 of them). Everything else in the step is still SGLang's: the embedding and its all-reduce, the dense layers 0 to 2, the final norm, the LM head and the logits' all-gather. This PR runs every layer and the tail as launches. That makes 79 launches a step: layer 0 from the step's tokens, layers 1 and 2 (dense), layers 3 to 77, and one tail launch. Result on MI355X, TP8, FP8 E4M3 KV cache, #69 → this PR: - MTP 3-1-4 at the golden accept length 2.99, per cycle: 7.47 → **7.14 ms** at one request (−4.5%; 400 → 419 tok/s per user), 11.05 → **10.77 ms** at two (−2.6%). - MTP off, median TPOT: 4.78 → **4.45 ms** at one request (−6.9%), 4.95 → **4.56** at two (−7.7%), 6.18 → **5.87** at four (−5.1%). At 16 requests, past the largest width, both run SGLang's path (15.91 / 15.91). - With #68 and #69, against `main` (as #68 measured it): MTP C=1 cycle 8.49 → 7.14 ms (−15.9%), C=2 12.57 → 10.77 (−14.4%); MTP-off TPOT at one request 5.48 → 4.45 ms (−18.9%). - Kernel-only (`time_layers.py`), the four new launches add 227 / 276 / 400 µs to the step at 2 / 4 / 8 rows (TP8), in place of SGLang's embedding, dense layers, final norm, LM head and all-gather. Accuracy does not decrease: see Accuracy Tests. Scope: - GPU: MI355X (gfx950), TP 4 or 8, as before. - Steps: decode and MTP verify steps of 1 to 8 rows (1 to 4 at TP4), as before. - Coverage: all 78 layers, the embedding, the final norm, the LM head and the logits' all-gather. Output assembly and sampling stay SGLang's. - SGLang: no change. ## Modifications Four commits: the new builds (1), the coverage with its checks and bounds (2 and 5), the first launch's wait deadline (3), and the kernel checks and timing of the new builds (4). ### 1. Dense, token and tail builds of the launch (`kernel.py`, `layout.py`, `runner.py`) Three new builds, each selected by a `Glm53Build` field that is off by default (`dense`, `tokens`, `head`, with `vocab_rows`). With those fields off, the indexer and shared launches compile to the same code objects as #69, byte for byte (Speed Tests, Registers). - **`dense`**: a dense layer (0 to 2). GLM-5.3's dense layers are indexer layers, so the launch is the indexer launch with the layer's FP8 SwiGLU MLP in place of the MoE, on the experts' W8A8 path: - the MLP's intermediates a rank are 6 expert-shaped slices (1,536 = 6 × 256 at TP8, 3,072 = 6 × 512 at TP4); - every row routes to all 6 with weight 1; the router GEMV and the top-8 are skipped; - activations are quantized to FP8 per 128, as SGLang's dense FP8 linears do, and the mid is quantized per 128 before down; - up / gate and down take every row as the columns of one MFMA per slice. - **`tokens`** (with `dense`): layer 0, the step's first launch. It takes the step's token ids (int32, −1 for a pad row). The rank whose vocabulary shard holds a token pushes that embedding row to every rank, as tagged bf16 pairs, and the launch reads them where it reads its input rows (the qkv_a and index front tiles, and the attention residual). One rank holds each row, so this equals SGLang's embedding all-reduce. - **`head`**: the tail. The final RMSNorm (bf16, as SGLang's), this rank's bf16 LM-head rows read in place, and the logits' all-gather into f32 [rows, vocab]. The gathered logits use two symmetric slots by launch parity, so a rank one tail ahead never overwrites logits a slower peer is reading. The normed rows are written too, for the hidden states. - The tail has no hand-off mailboxes. Its launch scratch (#68's uncached one-rank `PeerBuffer`) is at least 256 bytes, since an uncached allocation of 0 bytes has no address. ### 2. The spec covers layers 0 to 2 and the tail (`spec.py`, `weights.py`, `check.py`) - **The forward.** Layer 0 takes the step's tokens; SGLang's embedding is not run. Layers 1 and 2 are dense launches. Layer 3, which reuses the previous top-k, reads layer 2's from the launches' own buffer, so the spec's copy of SGLang's top-k into it goes. The tail runs after layer 77; SGLang's final norm is not run. - **The logits.** The spec binds the logits processor's `_get_logits` to its own at build time. On a routed step it returns the tail's f32 logits [rows, vocab], exactly what `_get_logits` returns after SGLang's LM head and all-gather; otherwise it calls SGLang's. `LogitsProcessor.forward`'s pruning, output assembly, logprob paths and sampler stay SGLang's (see Reviewer notes). - **Hidden states.** The forward hands the logits processor the tail's normed rows, which is what SGLang's `DeepseekV2ForCausalLM.forward` hands it (the final norm's output). So `return_hidden_states` and MTP's capture get the final norm's rows. A unit test checks what the forward hands over, and an e2e case checks the returned rows against the sampled tokens. - **Step input.** Layer 0's token ids are copied into a per-width buffer with the step's other tables, in the graph when captured. Every launch argument is a per-width buffer or SGLang's own tensor, read at launch: no table of addresses is kept, so an eager step at a captured width cannot change what a replay reads. - **Weights.** The dense MLPs are copied once into the experts' layout (81 MiB a rank at TP8, logged with the other copies). The embedding and the LM head are read in place. - **Refusals** (model reasons): - a dense MLP that is not block FP8, or whose intermediates do not slice into the experts' layout; a dense layer without its indexer; - an embedding or LM head that is quantized, padded, has added rows, or is not sharded over the TP group in rank order; an embedding with another output dtype; - an f32 LM head or f32 LM-head output (`use_fp32_lm_head`, `use_fp32_lm_head_output`); - logits gathered other than over the TP group (DP attention, DP LM head); - a logit scale or soft-capping. - A step with input embeddings runs the model's own forward. - **`SGLANG_READS`** gains the logits processor's `_get_logits` and the flags above, the vocab-parallel embedding's shard indices, the dense MLP's `gate_up_proj` / `down_proj`, `self_attn.o_proj` and `reduce_results`, and `residual.batch.clear`; the reads of the dense layers' SGLang path go. - **`SGLANG_AMD_MONO_CHECK`:** - layer 0 is checked against SGLang's embedding plus layer 0; - each dense layer's attention: its o_proj output (the launch's TP partials, summed in rank order before the residual add) against SGLang's o_proj output (a forward hook), all-reduced; - each dense layer's MLP: its input against the input of SGLang's MLP (a forward pre-hook), its mid against an f32 gate | up on its own FP8 activations, its output against the f32 down of its own FP8 mid, and its distance to an f32 MLP over SGLang's FP8 MLP's on the same input; - the tail's logits against SGLang's final norm, LM head and all-gather on the same rows, and against an f32 tail. - `check.REPORTED` lists each check's outputs; a reported output it doesn't list raises, and a unit test keeps every bound to a listed output. ### 3. A step's first launch waits up to half the watchdog (`spec.py`) As #61 did for glm53_flash: the runtime of layer 0's launch gets `min(max(ctx.deadline_s, 0.5 * watchdog_timeout), guard.MAX_DEADLINE_TICKS // guard.CLOCK_HZ)` as its wait deadline. Layer 0 is now always a step's first launch, where the ranks meet after their hosts' work between steps (first-use JIT builds, for example), with no RCCL collective before it. Every other runtime keeps `ctx.deadline_s`. `_build` passes the cap, so the unit test runs on a CPU-only host. | `--watchdog-timeout` | every other launch | layer 0's launch | | --- | --- | --- | | 300 s (SGLang's default) | 60 s | 150 s | | 1,200 s (the e2e test's) | 60 s | 600 s | | 20 s (the e2e fault test's) | 10 s | 10 s | | 11,100 s, or `inf` | 60 s | 5,497 s (the cap) | ### 4. Kernel checks and timing (`test/kernels/glm53/check_layer.py`, `time_layers.py`, `test_glm53_kernels.py`) - `check_layer.py` builds the new kinds on random weights: `dense`, `embed` (layer 0 from token ids, then dense layers), `head` (the tail alone) and `full` (layer 0, the dense layers 1 and 2, the model order, then the tail). `--writes` and `--replays` take them as they take the others. - `--check-dense`: the last dense launch's MLP, from its own mailboxes, against f32 references on its own inputs, within the spec's `dense.*` bounds; rank 0's down scales ×2, or its gate | up scales ×2, must fail them. - `--check-embed`: layer 0 from the token ids equals the same layer as a dense launch fed the gathered embedding rows: outputs, MLA and index cache rows, top-k, bitwise. A changed token id must change the output. - `--check-head`: the tail's logits against a torch tail and an f32 tail within the spec's `head.*` bounds; the final norm ×1.01 must fail `head.logits.rel`. - `time_layers.py` times `--kind full` (the served step's 79 launches), `dense`, `embed` and `head`. - `test_glm53_kernels.py` runs the new checks at TP4 and TP8, and the writes check and a 20-replay graph of the new kinds. ### 5. Bounds (`NUMERICS`) Relative bounds, about 3× the largest value `SGLANG_AMD_MONO_CHECK` measured in serving (`dense.attn.rel`: 2.2×, below): the e2e file's four check cases, TP8, 1 to 8 rows, decode and MTP verify steps, bf16 and E4M3 MLA cache, positions to about 108K keys (about 200 dense checks and 68 tail checks a run). No existing bound changes. | Bound | Value | Serving max | What it checks; what fails it | | --- | --- | --- | --- | | `dense.attn.rel` | 0.15 | 0.068 (108K keys; 0.052 on E4M3 verify; ≤ 0.020 in decode and bf16 verify) | The dense launch's o_proj output against SGLang's: attention and o_proj, on every dense layer. The dense launches' o_proj scales ×1.2 give 0.197 to 0.210 on layers 0, 1 and 2 | | `dense.mlp_in.rel` | 0.09 | 0.029 (layer 0 at 108K keys; ≤ 0.014 elsewhere) | The dense launch's MLP input (post-attention norm of its FFN input) against SGLang's MLP input: attention, o_proj and the residual. The dense launches' o_proj scales ×1.2 give up to 0.37 on layer 0 (see Accuracy Tests) | | `dense.mid.rel` | 1.7e-4 | 4.4e-5 | gate / up and SiLU against f32 on the launch's own FP8 activations. One rank's gate / up scales ×2: 3.45 | | `dense.x.rel` | 2.8e-2 | 9.7e-3 | down and the TP reduce, against the f32 down of the launch's own FP8 mid. One rank's down scales ×2: 0.125 | | `dense.excess.rel` | 0.11 | 0.036 (min −0.046) | The MLP's distance to an f32 MLP on its own input, over SGLang's FP8 MLP's, less one. The two controls above: 2.0 to 9.6 | | `head.logits.rel` | 1.7e-3 | 4.8e-4 | The tail's logits against SGLang's final norm, LM head and all-gather on the same rows. The final norm ×1.01: 9.2e-3 | | `head.logits_vs_fp32.rel` | 7.5e-3 | 2.25e-3 | Against an f32 tail; SGLang's own distance is the same (2.25e-3) | - **`dense.attn.rel` is 2.2× its serving max, not 3×.** 3× would be 0.21, where the ×1.2 control sits (0.197 to 0.210), so it would not catch that error; 0.15 catches it on every check. The long-context maximum (0.068, layer 1 at about 108K keys) comes from top-k differences at the 2,048th key's threshold, where the launch's bf16 index scores and SGLang's FP8-q scores part: `layer.topk.rel` reports them (up to 0.064 in that case), while `layer.topk_own` (a row's own draft-chain keys) stays 0. At positions below 9K keys the maximum is 0.020. - Reported without a bound: `dense.x_vs_fp32.rel` (0.046; SGLang's `dense.stock_vs_fp32.rel` 0.046), `head.normed.rel` (≤ 9.8e-5), `head.top1.agree` (1.0 in every check). - The dense layers' existing `layer.*` bounds hold too: `layer.streams.cosine` ≥ 0.988, `layer.kv.cosine` ≥ 0.9991, `layer.index.cosine` ≥ 0.99986 over all layers. ## Accuracy Tests Setup as #68: MI355X, TP8, 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. GSM8K runs sgl-eval 0.1.3 from `PYTHONPATH`, with a pure-Python stand-in for its compiled `editdistance`, which the image's Python can't load (as in #68). - **E2E** (`test/e2e/test_glm53_server.py`, TP8), this head against #69: | Case | This PR | #69 | | --- | --- | --- | | GSM8K past the largest width (200, thinking, T 1.0) | 0.985 | 0.975 | | GSM8K with MTP, and the accept length | 0.99, 3.11 | 0.985, 3.16 | | Decode against prefill, KL at 2 / 3 / 7 / 8 rows (bound 0.06) | 0.024 / 0.024 / 0.018 / 0.022 | 0.021 / 0.022 / 0.019 / 0.022 | | The four check cases (decode, verify, verify on E4M3, past the parallel top-k) | pass, 0 violations | — | | **New:** returned hidden states (each decode row times the LM head gives the sampled token) | 21 / 21 rows | — | | Injected fault | fails as on `main`, at its last check (below) | — | - The injected-fault case fails here on `main` the same way, only at its last check: every rank stops at the wait deadline with `HIP error: an illegal memory access was encountered` and the server exits without sending a token, but ROCr in our container prints no `Memory access fault ... on address (nil)` line (not even for a bare one-GPU kernel that only stores to 0x0), so the test finds no fault address (`assertTrue(faults)`). In the reviewer's setup on #68 and #69 the line is printed and the case passes. - The GSM8K, KL and hidden-states cases ran on a tree that differs from this head only in check-mode code (the `dense.attn` check and the `dense.mlp_in` bound), which those cases don't run. - The four check cases ran with `dense.attn.rel` reported and unbounded (its bound was set from that run: 0 violations of the other bounds, every `dense.attn.rel` ≤ 0.068), then the decode case again with its bound (passes). - **The dense checks' negative control: a small error on purpose.** The dense launches read their o_proj block scales ×1.2 (a copy; SGLang's unchanged), on the decode check case. `dense.attn.rel` gives 0.197 to 0.210 and fails on all three dense layers, 57 of 57 checks (bound 0.15). `dense.mlp_in.rel` fails only on layer 0 (19 of 19, up to 0.37): on layers 1 and 2 attention is a small part of the FFN input. `layer.streams.cosine` (0.95) doesn't catch it (min 0.959). - **Kernel tests** (`test_glm53_kernels.py`, TP4 and TP8, every width the TP builds): 7 passed (13 min), the four new cases included: - dense MLP, at TP8: mid 1.8e-5 to 2.0e-5, output 3.7e-3 to 4.0e-3, excess −0.034 to +0.017 against the bounds above; rank 0's down scales ×2 give an output rel of 0.125 and an excess of 2.0 to 2.2, its gate / up scales ×2 a mid rel of 3.45 and an excess of 8.6 to 9.6; - layer 0 from the tokens equals the dense launch fed the gathered rows, bitwise (outputs, MLA and index cache rows, top-k), at 1 / 2, 2 / 2, 3 / 4, 4 / 4, 7 / 8 and 8 / 8 live rows; a changed token id changes the output; - the tail, at TP8: logits 4.0e-5 to 6.0e-5 against a torch tail and 2.35e-3 against an f32 tail (the torch tail's own: 2.35e-3), normed rows bitwise, top-1 100%; the final norm ×1.01 gives 9.2e-3; - 20 graph replays of layer 0, the 2 dense launches, 4 MoE launches and the tail equal the eager step, with a pad row and with verify rows; - each new kind writes exactly its live rows' cache rows. - **Bitwise, existing launches:** the indexer and shared code objects are byte-identical to #69's (all 40 keys of #68's survey, below), so every existing launch is unchanged. ## Speed Tests and Profiling Setup as Accuracy Tests. All numbers are from this head (3f5f995) against #69 (cfba6f0), in one allocation. Serving and the kernel-only runs used a tree that differs from this head only in check-mode code (the `dense.attn` check and the `dense.mlp_in` bound), which they don't run. Serving as #68: one server per arm and setting (MTP 3-1-4, and MTP off), all of a setting's points on it, 8,192-token ShareGPT prompts, 512 greedy output tokens, 4 waves of C requests after a warm-up wave, MTP at the golden accept length (`SGLANG_SIMULATE_ACC_LEN=2.99`, `real-draft-token`; measured 2.99 on every server). **Out-of-tree streaming client** (`sglang.benchmark.serving` does not import on this image), as in #68. One round; #69's servers here agree with #69's own numbers within 0.2%. At C≥2 the server prefills a wave's requests one after another, and a request that has started decoding waits out the later prefills. The C≥2 numbers below are those of each wave's last-prefilled request, which decodes beside the others from its first token to its last. TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms: | Requests | #69 | This PR | Change | tok/s per user at AL 2.99 | | --- | --- | --- | --- | --- | | 1 (4 verify rows) | 7.474 | 7.135 | −4.5% | 400 → 419 | | 2 (8 rows) | 11.054 | 10.766 | −2.6% | 270 → 278 | TP8, FP8 E4M3 KV, MTP off, median TPOT in ms: | Requests | #69 | This PR | Change | | --- | --- | --- | --- | | 1 | 4.776 | 4.445 | −6.9% | | 2 | 4.947 | 4.564 | −7.7% | | 4 | 6.184 | 5.867 | −5.1% | | 16 (SGLang's path) | 15.909 | 15.910 | 0.0% | Serving was measured at TP8 on the E4M3 cache. At TP4, one server (MTP, one request: 4 verify rows) booted and served at 9.52 ms a cycle, AL 2.99, without error; TP4 is otherwise measured kernel-only. ### Kernel-only (`time_layers.py`) `torchrun --nproc-per-node <TP> test/kernels/glm53/time_layers.py --kind <kind> --kv fp8 --rows 2,4,8`, µs a launch, the mean of two passes in alternating order (each the median of 7 trials of 50 replays). TP8: | Launch | 2 rows | 4 rows | 8 rows | | --- | --- | --- | --- | | model order (75), #69 → this PR | 54.05 → 54.05 | 71.42 → 70.79 | 115.47 → 115.44 | | **whole step (79), `--kind full`** | 54.19 | 71.30 | 114.69 | | step, 75 launches (#69) → 79 (this PR), ms | 4.05 → 4.28 | 5.36 → 5.63 | 8.66 → 9.06 | | dense (layers 1-2) | 62.58 | 77.11 | 102.50 | | layer 0 from the tokens (alone in its graph) | 72.10 | 89.20 | 112.95 | | tail | 48.92 | 51.23 | 57.06 | - The model order runs byte-identical code objects in both arms; its 4-row difference is the passes' spread. - The serving gain is the SGLang work these launches replace (its embedding all-reduce, the dense layers' unfused path, final norm, LM head GEMM and all-gather) less the 227 / 276 / 400 µs they add. TP4 (2 and 4 rows): #69's 75 launches 70.07 / 102.89 µs a launch (5.26 / 7.72 ms a step); this PR's 79, `--kind full`, 69.95 / 101.22 (5.53 / 8.00 ms), so the four new launches add 271 / 279 µs. ### Registers A compile-only build of every key the spec and the kernel tests build: #68's 40 (TP4 2 / 4 rows, TP8 2 / 4 / 8; indexer and shared; bf16 and E4M3; `stop=True` and `stop=False`) plus 50 new ones (dense and layer 0, both caches; the tail; both `stop`s): 90 builds. - **The 40 existing builds are byte-identical to #69's** (`cmp` of the code objects). - Every symbol is `glm5_monokernel_0`, as on #69. - **The new builds** spill no VGPRs and keep no private segment, except: the 8-row tail (TP8), 8 VGPRs / 36 bytes (`stop=True` and `False`), and the 8-row dense build at TP8 on the bf16 cache (`stop=True`), 1 VGPR / 8 bytes. The E4M3 dense build and every layer 0 build have none. ### Startup and memory - **Weight copies:** the dense MLPs as expert slices, 81.0 MiB a rank at TP8 (162.0 MiB at TP4), in the startup log with the other copies. W_UK / W_UV and the indexer's packed tensors grow by layers 0-2's (131.3 → 136.5 MiB, 6.8 → 7.9 MiB). - **Launch buffers** (TP8): 363.9 → 459.0 MiB (MTP server), 364.1 → 460.7 MiB (MTP off): the new runtimes, with the tail's logits slots. - **KV pool** (TP8, FP8): 2,875,584 → 2,867,840 tokens (−0.27%) with MTP, 2,948,032 → 2,938,816 (−0.31%) without. - **Build time** (FlyDSL cache warm): 4.2 → 8.8 s (MTP, 2 widths), 4.6 → 12.4 s (8 widths). - **A cold boot** (TP8, MTP; SGLang JIT, FlyDSL, Triton and inductor caches empty, on NFS): the launches built in 57.1 s, the server was ready in 728 s and served an 8K request (AL 2.99), with no guard stop. ### Tests - Unit tests on a CPU-only host: 277 passed, against #69's 265; the same 2 contract failures (minimax, minimax_m31) and 3 collection errors as #69. The new cases cover a routed step's launches and what it hands the logits processor, the first launch's deadline at watchdogs of 20, 300, 1,200 and 11,100 s and `inf`, the dense slices (row-major and shuffled), the embedding and LM-head shards, the dense check's read of the launch's o_proj partials, the refusals, the reported outputs, and the tail check against a wrong final norm. - `test_glm53_kernels.py` with 8 GPUs (TP4 and TP8): 7 passed (13 min). ## Reviewer notes - **`_get_logits` is the one SGLang method the spec replaces**, as in #23. The binding is active only inside a routed step; otherwise it calls SGLang's. - The dense layers run at SGLang's own precision (W8A8, FP8 per 128); the LM head stays bf16. - **Step input and graphs.** Every launch argument is a per-width buffer or SGLang's own tensor, read at launch; the spec keeps no table of addresses. A replay rewrites the step's buffers in the graph before it reads them, so an eager step at a captured width cannot change what a replay reads. The tail's logits and normed rows are the runtime's per-width buffers, as a captured graph's outputs are; SGLang consumes them in the step, in stream order. - **The dense layers' tight checks.** `layer.streams.cosine` (0.95) is loose. On the dense layers the tight companions are `dense.attn.rel` (attention and o_proj, against SGLang's o_proj output), `dense.mlp_in.rel` (the MLP's input) and the MLP's own bounds. The o_proj ×1.2 control fails `dense.attn.rel` on all three layers. `dense.attn.rel`'s bound is 2.2× its serving max, so a long-context run that drifts past about 2× today's maximum would report a violation. - `kernel.py` grows by 359 lines (the tail, the token push and the dense MLP branch): candidates for the stage-builder split. - #66 (AITER compiled dispatch, now on `main`) changes the layer's launch call in `gfx950/runner.py`. This PR adds a second launch call (the tail's), which needs the same change when the stack rebases onto `main`. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan
pushed a commit
that referenced
this pull request
Oct 29, 2026
… exact top-2048 on up to 64 slices (#74) ## Motivation **Stacked on #68, #69 and #73** (`[Feat] [MI355X] glm53: Cover layers 0-2 and the logits`). Merge those first. This PR's diff is its last four commits, and it uses #68's `test/kernels/glm53/` tools. On an indexer launch the attention waits for the index chain: index query, key scoring, then the exact top-2048. That covers 18 of the 75 MoE layer launches of a step, plus layer 0 and the dense layers 1–2, which reuse the indexer stages. Two parts of the chain grow with context: - **Scoring reads each key once per row.** #42's per-wave scoring scores one row at a time, so at MTP's 4 verify rows every key of the request is read 4 times. - **The selection is mostly serial.** A row of up to 32,768 keys selects on one CTA. Past that, #24 spreads it over at most 16 slices, and the row's CTA adds their 4,096-bin histograms and, once the slices have gathered, copies their lists. This PR scores a request's rows together, reading each key once from LDS batches of whole pages. It then selects every row's top-2048 on up to 64 slices with all-gathered 6-bit digit histograms. The selection stays exact: the same keys, in the same order and with the same ties, as the base's radix select. Every launch's outputs are bitwise equal to the base's (Accuracy Tests). Result on MI355X, base (#73) → this PR: - TP8, MTP 3-1-4 at the golden accept length 2.99, cycle at one request: - 8K context: 7.16 → 6.92 ms (−3.4%); - 128K context: 8.29 → 7.24 ms (−12.7%). - At two requests, while both decode: - 8K: 10.76 → 10.54 ms (−2.1%); - 128K: 12.66 → 11.48 ms (−9.3%). - TP8, MTP off, median TPOT at 8K: 4.45 → 4.42 ms (−0.5%) at one request, 4.56 → 4.44 (−2.8%) at two, 5.87 → 5.77 (−1.7%) at four. - Kernel-only, `time_layers.py`, E4M3 MLA cache, TP8, µs a launch at 4 rows: - indexer launch: 96.42 → 86.33 at 8K (−10.5%) and 140.00 → 100.00 at 128K (−28.6%); - the 75-launch model order: 71.12 → 68.78 (−3.3%) and 82.88 → 72.20 (−12.9%). The gain grows with context because both rewritten stages scale with the keys a row reads. ## Modifications The first commit adds a check. The next three are the changes, each measured on top of the one before (Speed Tests, "Each change on its own"). ### 0. An exact top-k check (`test/kernels/glm53/check_topk.py`, `check_layer.py`, `test_glm53_kernels.py`) - `check_topk.py` runs indexer launches under torchrun at TP4 or TP8, at chosen positions (`--pos`). `--kind dense` and `--kind embed` run layer 0 and the dense layers. - It uses four index caches: - `check_layer.py`'s random keys; - four keys repeated over every position, so the 2048th score falls in a tie (`ties`); - all-zero keys (`zeros`); - each row's last 4,096 keys scaled 8x, so the selection crowds into them (`recent`). - Row sets: - the served decode points (`serving`, as `time_layers.py`); - one request a row, at `--pos` and below, some under 2,048 keys (`mixed`); - one request's rows, a draft chain as long as the width (`chain`). - Before each step, the cached scales of the step's own keys are poisoned (`check.poison_new_index_scales`). A launch that scores a key this launch writes from the cache therefore fails. - After each launch it checks every live row: - **Scores** (`check.launch_scores`) against an f64 reference. The reference uses the launch's index query (its `index_q` mailbox, rotated and transformed in f32 as the kernel does, then bf16), its head weights, and the row's keys as the launch leaves the index cache. The error is taken over the row's largest reference score; the bound is 1e-4. - **The list** (the top-k at the row's `indptr`), entry by entry, against the radix select's list of the launch's own scores. That list is the keys above the 2048th key's 12-, 24- or 32-bit prefix in position order, then the prefix's keys above the 2048th, then its ties in position order. `missed` also counts the keys of the exact top-2048 (score descending, then position) that the list lacks, as `check.py`'s `layer.topk_exact` does. - **Negative controls** on copies of each row's result: a selected slot replaced by the best unselected key, two entries swapped, and one score's bits flipped must each fail. - `--dump DIR` writes SHA-256 digests of every row's scores and list in `check_layer.py`'s format, so `check_layer.py --compare` can match two trees. - `test_glm53_kernels.py` runs it at TP4 (widths 2, 4) and TP8 (2, 4, 8): - the indexer launch at about 8K and 300K keys (at 300K a scoring CTA's run takes three batches and a slice two LDS chunks); - layer 0 and a dense layer at 8K and 128K. - It passes on the base too: it checks what this PR keeps, the base's list. - **Test fix, in the same commit:** `check_layer.py`'s steps gave each row its request's `req_ids` entry, while the page table has a row per step row. At 8 verify rows, the second request's rows read the first request's page-table row. Each row now gets its own entry, as the spec passes them (`req_ids` = row index, `page_table` = the DSA backend's row per step row). ### 1. Each row's index key on its own task, and the index-query tiles first (`kernel.py`, `layout.py`) - **Index key:** the cache task used to normalize, rope and store a row's index key after the row's KV row. The key is now an `index_key` task of its own, on the row's q_norm CTA. - **q_norm CTAs:** they run their own index-query tile right after q_norm, ahead of the index key and the KV row. - This replaces #68's change 2, where those tiles ran second, on the CTAs of the first extra tiles. - That change's shuffle and its TP4 guard go. - **q_b CTAs:** a q_b CTA issues its index-query tile's weights with q_b's, before the wait, and runs that tile before its q_b tile. - **Values:** the same values, computed in the same order. - **Unpaged fused indexer:** it goes. Every indexer build reads SGLang's paged index cache (`runner.py` passes `index_paged=key.with_indexer`). - **On its own this change costs 0.6 to 3.3 µs** (Speed Tests). It stays because changes 2 and 3 were built and measured on it: the scoring patches this launch's keys from the `index_key` task's pairs. ### 2. A request's rows scored together, each key read once (`kernel.py`, `layout.py`) - **Before (#42):** a CTA scored tiles of 64 keys of one row, a wave every 8th group of 16 keys. A request's 4 verify rows read each key 4 times, and every scoring CTA polled, roped and Hadamard-transformed its row's whole query. - **Groups of rows:** up to 4 consecutive rows of one request score its keys in one pass. - A row belongs to the previous row's request when its page table maps that row's position to the slot that row writes; a page being written belongs to its request alone. - Each row's 2 waves (4 rows) or 4 waves (2 rows) hold its query and head weights in registers. - **Runs of pages:** the request's pages are split into runs of at least 2 pages, one run a CTA, up to 256 runs. They are counted from a CTA that the group's first row offsets. - A CTA copies its run into LDS 8 pages at a time with 16-byte loads. The next batch's loads are in flight while its waves score the current one. - A run's block ids load with the query's poll. - **Keys this launch writes:** - Each `index_key` task also publishes its row's key as cached, FP8 value words and scale, in `index_k8` pairs. A batch holding any of those keys patches them in from the pairs, so the cache store need not be visible yet. - The patch covers every row of the request in this launch, not only the group's. A draft chain of more than 4 rows (MTP with 4 or more steps) spans two groups, and the second group reads the first group's new keys. - The base handles this case by comparing every key's slot with every row's slot (#15, #42). Here, `check_topk.py`'s `chain` rowset at 8 rows checks it. With each group patching only its own rows, that check fails (Accuracy Tests). - **The query once a row:** `index_qh` tasks rope and Hadamard-transform each row's index query once and publish it in bf16; the scoring CTAs poll it into LDS. - A task covers a row's 8 heads, a wave a head. - At TP4, q_b runs on every CTA and produces every index-query tile, so these tasks run after q_b. - **Scores:** every score is computed as before, with the same MFMA chains over the same bf16 query and keys and the same head sums in the same order. The scores are bitwise equal to the base's (Accuracy Tests). - **Removed:** the `index_ready` flag and the `index_k_new` pairs. ### 3. Every row's exact top-2048 on up to 64 slices (`kernel.py`, `layout.py`) - **Before (#24):** a row of up to 32,768 keys selected on one CTA, reading every score at least twice. Past that, up to 16 slices built 4,096-bin histograms of their chunks; one CTA added them and published the threshold bin; the slices gathered their winners and candidates; and the row's CTA copied them and finished the radix select. - **Now** every row selects on up to 64 slices, a slice per 1,024 keys. `layout.select_slices` gives the split, and the kernel and the CPU test both call it. Each slice is a run of the row's keys in position order, held as radix keys in LDS. - **Digit histograms:** a radix select by 6-bit digits from the top. For each digit, every slice publishes the 64-bin histogram of its keys under the prefix found so far (32 tagged pairs of 16-bit counts) and reads every other slice's. It then finds the digit holding the remaining k-th key and counts the keys above it in the slices before its own. - **Placing and resolving:** the prefix stops at 12 or 24 bits once its keys fit the candidate room (as the base's does), or runs to all 32 bits. Each slice then places its keys above the prefix, whose places are known, and stages its keys with the prefix, with their cache slots, for the row's first slice. That slice resolves them in two 10-bit passes, or one 8-bit pass at 24 bits. At 32 bits the ties take their places directly. - **Order:** the list is the base's one-CTA radix select's, entry by entry. That is the keys above the threshold's 12-, 24- or 32-bit prefix in position order, then the prefix's keys above the threshold, then its ties in position order. - **Placement:** the slices start past uk's CTAs. - When every row has at most 256 / S slices, slice c of row s is task c·S + s, so short rows' slices share one range of CTAs. - Otherwise each row takes a run of 64 tasks. - **Scratch:** the histogram, threshold and gather stages' buffers give way to `sel_hist` and `sel_cand`. The indexer runtimes of widths 2, 4 and 8 together go from 163.4 to 135.2 MiB a rank at TP8, and from 76.5 to 64.4 MiB at TP4 (widths 2, 4). - **Tests, in the same commit:** - `test/unit/test_glm53_select.py` covers `select_slices`: the slices tile a row's keys once, in position order, none empty, at most 64, a slice's keys within its LDS room and its 16-bit counts. It also covers the stage order. - `test_glm53_server.py`'s parallel-top-k case asserts that its rows reach all 64 slices. ## Accuracy Tests - **Exact top-2048** (`check_topk.py`, E4M3 MLA cache, 2 indexer launches × 2 steps). - Coverage: positions 8,447, 131,071, 262,143 and 524,287, the four index caches, and the serving, mixed and chain row sets. - TP8, widths 2, 4, 8: 288 checks on this head and 288 on the base, every row exact (missed 0, order exact). The largest score error is 1.8e-6, and every negative control is caught. - Head vs base: the digests of every row's scores and list are equal (41,984 digests, 0 differ). The scores are bitwise the base's, and so is the list, order and ties included. - TP4 (widths 2, 4) on this head: every row exact. - Negative control for the new-key patch: with each scoring group patching only its own rows, the 8-row `chain` set fails at 8K and 128K on every index cache. The launch scores the earlier rows' new keys from their poisoned cached scales, giving a score error of about 1e30. - The kernel tests (below) also run it on layer 0 and a dense layer at 8K and 128K, at TP4 and TP8. - **Bitwise** (`check_layer.py --dump`, `--compare`). - Coverage: kinds `shared`, `index`, `model` (the first 8 launches), `dense` and `full` (layer 0, the dense layers 1–2 and the tail); bf16 and E4M3 MLA cache; every width the TP builds; one request a row, the same with a pad row, and MTP's verify rows; 3 steps. - TP8: 63,648 digests, all equal. TP4: 16,416 digests, all equal except one run of one configuration (bf16 dense, 2 rows with a pad row: 127 digests, the first dense launch's output with its top-k equal). That configuration, rerun 4 times, was equal base vs head and run to run. The negative control (one input value negated) changes 7,056 of 10,944 digests. - Every output and cache row is therefore bitwise equal to the base's, and so is the top-k, at every TP, width and cache dtype the spec builds. - These dumps used `check_layer.py` with its runtimes built before the layers' weights, as #69's tool did. With the base's own `check_layer.py`, which since #73 draws every layer's weights first, two TP8 width-2 configurations (E4M3 index with verify rows, bf16 model with a pad row) vary from run to run on the base itself, with the same kernel code objects. This is under investigation separately and is not from this PR. - **What can differ from SGLang:** nothing new. The serving check's `layer.topk_own` compares with SGLang's selection from FP8-q scores, while the kernel scores the bf16 query. At a threshold tie, the two can pick a different key, as on the base. - **`SGLANG_AMD_MONO_CHECK=1`:** - The e2e file's four check cases pass on this head, each asserting no violation: - decode rows 1 / 3 / 8; - MTP verify rows 4 / 8; - verify rows on the E4M3 MLA cache; - rows past the parallel top-k (about 108K keys at widths 2 and 4, on 64 slices; about 36K at width 8, on 36). - Past 128K keys (two requests of 140K tokens, E4M3 KV, eager), on this head, 0 violations: - decode at 1 and 2 rows: 750 layer checks, 54 of them indexer checks past 131,072 keys; - MTP verify at 4 and 8 rows: 525 checks, 36 past 131,072 keys. - `layer.topk_exact` and `layer.topk_own` are 0 in every indexer check; `layer.topk.rel` is at most 0.048 and `layer.index.cosine` at least 0.99997. - **E2E** (`test/e2e/test_glm53_server.py`, TP8, on this head): - decode against prefill (KL), at 2 / 3 / 7 / 8 rows: passed; - the four check cases: passed (above); - GSM8K past the largest width, without and with MTP, and the MTP accept-length bound: passed. sgl-eval 0.1.3 was put on `PYTHONPATH` with a pure-Python stand-in for its compiled `editdistance`, as for #68; - the injected fault fails here on `main` the same way, only at its last check: every rank stops at the wait deadline with `HIP error: an illegal memory access was encountered` and the server exits without sending a token, but ROCr in our container prints no `Memory access fault ... on address (nil)` line (not even for a bare one-GPU kernel that only stores to 0x0), so the test's `assertTrue(faults)` fails. In the reviewer's setup on #68 and #69 the line is printed and the case passes. - **Long context, decode against prefill at 128K** (an out-of-tree client: 4 prompts of 131,072 ShareGPT tokens with a GSM8K question, 64 greedy tokens, the decode's top-20 log-probabilities against a prefill of prompt and output; E4M3 KV, MTP off): - base: KL mean 0.081, top-1 agreement 0.941; - this PR: KL mean 0.063, top-1 agreement 0.956; - both arms have the same largest KL (1.87) over 252 positions. ## Speed Tests and Profiling **Setup:** as #68: one node with 8 × MI355X; image `lmsysorg/sglang-rocm:v0.5.22-rocm720-mi35x-20261025` (squashfs `sha256:fb07e0ef78fe`), AITER 285579b6, FlyDSL 0.4.0; `zai-org/GLM-5.3` at `ba32f8c1` (FP8); the e2e test's server flags with `--kv-cache-dtype fp8_e4m3`. - The e2e and check runs overlay SGLang's `python/sglang` at `5044779f`. - The serving runs use the image's SGLang with `--enable-amd-monokernel` and the plugin's spec registration from before #48. - Kernel-only numbers don't depend on SGLang. - All numbers are from this head against its base: - serving, e2e, the check runs, the long-context runs and the bitwise dumps: this head (b6fdf8e) against the base (3f5f995); - kernel-only, compile survey, exact top-k and kernel tests: the same four commits on #73's previous head (809f6ec), whose kernel code is the same; #73's last change touched only `check.py` and `spec.py`. **Serving:** as #68. - One server per arm and setting: MTP 3-1-4 at 8K, MTP 3-1-4 at 128K, MTP off at 8K. All of a setting's points ran on the same server. - Prompts of 8,192 or 131,072 ShareGPT tokens, no shared prefix, radix cache flushed before each wave. 512 greedy output tokens with EOS ignored. 4 waves (3 at 128K) of C concurrent requests after a warm-up wave. - MTP runs at the golden accept length (`SGLANG_SIMULATE_ACC_LEN=2.99`; measured 2.99 on every server). A cycle is a request's decode time over its verify steps. - One round. **The client is out of tree**: `sglang.benchmark.serving` does not import on this image, so a small streaming client against `/generate` measures TTFT, TPOT and the verify count from `meta_info`. TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms (median over requests): | Context | Requests | base | This PR | Change | | --- | --- | --- | --- | --- | | 8K | 1 | 7.164 | 6.920 | −3.4% | | 8K | 2 | 10.760 | 10.539 | −2.1% | | 128K | 1 | 8.287 | 7.238 | −12.7% | | 128K | 2 | 12.655 | 11.483 | −9.3% | The requests of a two-request wave prefill one after the other. The first-prefilled request waits out the other's prefill (380 ms at 8K, 7.9 s at 128K) before both decode, so its cycle is the one above plus that wait over its verify steps. The table gives the other request, which decodes beside the first from its first token to its last; the MTP-off rows below are measured the same way. TP8, FP8 E4M3 KV, MTP off, 8K, median TPOT in ms: | Requests | base | This PR | Change | | --- | --- | --- | --- | | 1 | 4.446 | 4.424 | −0.5% | | 2 | 4.563 | 4.437 | −2.8% | | 4 | 5.872 | 5.773 | −1.7% | ### Kernel-only (`time_layers.py`) `torchrun --nproc-per-node <TP> test/kernels/glm53/time_layers.py --kind <kind> --rows <rows> --kv fp8 --pos <8447 | 131071>`, µs a launch, base → this PR. Each number is the mean of two passes in alternating order in one allocation; each pass is the median of 7 trials of 50 replays. The passes agree within 0.8 µs (1.3 µs on TP4's 4-row indexer launch). TP8, E4M3 MLA cache: | Launch | Context | 2 rows | 4 rows | 8 rows | | --- | --- | --- | --- | --- | | indexer | 8K | 71.54 → 69.99 (−2.2%) | 96.42 → 86.33 (−10.5%) | 145.52 → 132.98 (−8.6%) | | indexer | 128K | 96.66 → 78.65 (−18.6%) | 140.00 → 100.00 (−28.6%) | 223.84 → 173.65 (−22.4%) | | model order (75) | 8K | 54.08 → 53.68 (−0.7%) | 71.12 → 68.78 (−3.3%) | 115.86 → 112.66 (−2.8%) | | model order (75) | 128K | 60.46 → 55.52 (−8.2%) | 82.88 → 72.20 (−12.9%) | 135.27 → 122.45 (−9.5%) | TP4, E4M3 MLA cache, indexer launch: | Context | 2 rows | 4 rows | | --- | --- | --- | | 8K | 87.88 → 88.77 (+1.0%) | 134.13 → 125.26 (−6.6%) | | 128K | 113.92 → 95.55 (−16.1%) | 178.05 → 138.61 (−22.2%) | ### Each change on its own Each commit against the one before it, TP8, E4M3 MLA cache, indexer launch, µs at 2 / 4 / 8 rows. Two passes in alternating order, one allocation. | Change | 8K | 128K | | --- | --- | --- | | 1. index key on its own task, index-query tiles first | +0.65 / +0.93 / +2.00 | +1.68 / +1.13 / +3.24 | | 2. a request's rows scored together | +0.51 / −6.04 / −11.64 | +2.98 / −15.13 / −35.29 | | 3. top-2048 on up to 64 slices | −2.71 / −4.98 / −2.90 | −22.67 / −26.00 / −18.14 | | **total** | **−1.55 / −10.09 / −12.54** | **−18.01 / −40.00 / −50.19** | - Change 1 on its own is slower by up to 3.2 µs. It is kept because changes 2 and 3 were built and measured on it. - Change 2 costs the 2-row launch 0.5 µs at 8K and 3.0 µs at 128K: a 2-row launch has one live row, which gains nothing from reading a key once for several rows. - Change 3 takes most of the 128K gain. ### Registers The coverage PR's compile survey: every key the spec builds (`stop=True`) and the same keys as the tests build them (`stop=False`). That is TP4 (2, 4 rows) and TP8 (2, 4, 8); indexer, shared, dense, layer-0 and tail builds; bf16 and E4M3 MLA cache; 90 builds. - Every symbol is `glm5_monokernel_0`, as on the base. - The shared and tail code objects are byte-identical to the base's. - VGPR spills / private segment bytes, base → this PR, for the builds that have either: | Build | `stop=True` (spec) | `stop=False` (tests) | | --- | --- | --- | | TP8 indexer, 8 rows, bf16 | 4 / 20 → 24 / 100 | 5 / 24 → 20 / 84 | | TP8 indexer, 8 rows, E4M3 | 4 / 20 → 20 / 84 | 5 / 24 → 24 / 100 | | TP8 dense, 8 rows, bf16 | 1 / 8 → 20 / 84 | 0 / 0 → 20 / 84 | | TP8 dense, 8 rows, E4M3 | 0 / 0 → 20 / 84 | 0 / 0 → 20 / 84 | | TP4 indexer, 4 rows, bf16 | 11 / 48 → 15 / 64 | 40 / 164 → 48 / 196 | | TP4 indexer, 4 rows, E4M3 | 16 / 68 → 13 / 56 | 34 / 140 → 48 / 196 | - The 8-row indexer and dense builds spill about 20 VGPRs (84–100 bytes), against 0–5 on the base. They come with change 2. Even so, the 8-row indexer launch is 8.6% faster at 8K and 22% at 128K. - Every other build has no VGPR spill and no private segment on either head. - The layer-0 builds have none at any width. ### Tests - Unit tests on a CPU-only host: the base's results (this head 288 passed, the base 276; the base's 2 contract failures and 3 collection errors). - `test_glm53_kernels.py` with 8 GPUs (TP4 and TP8): 8 passed (17 min 32 s), including the new exact top-k test. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan
pushed a commit
that referenced
this pull request
Oct 30, 2026
## Motivation **Stacked on #73 and #74** (and #68, #69 under them). Merge those first; this PR's diff is its last three commits. It uses #73's tail and token builds and `test/kernels/glm53/` tools. With MTP (EAGLE 3-1-4, SGLang's GLM-5.3 default), each cycle is one verify step and three draft steps of the NextN layer. After #73 the verify step runs entirely as launches. The draft steps still run SGLang's NextN: the embedding, `enorm` / `hnorm` / `eh_proj`, one full decoder layer (the indexer on the draft extend), `shared_head.norm`, the LM head and the logits' all-gather. That is dozens of kernels and three all-reduces a draft step, on 1 to 8 rows. This PR runs each draft step as two launches: the NextN layer, starting from the step's tokens through its `eh_proj` entry, and #73's tail with the NextN's norm. The draft extend (each request's verify rows) runs the indexer kind; the draft decode steps run the shared kind on the top-k the draft extend selected (SGLang's MTP index share). SGLang needs no change: Monokernel API v1 routes a draft model by the target's verify key, which a draft never runs, so the draft spec puts a wrapper in front of the draft's `model.forward`, as #29 does. Result on MI355X, TP8, FP8 E4M3 KV cache, #74 → this PR: - MTP 3-1-4 at the golden accept length 2.99, per cycle: 6.93 → **6.44 ms** at one request (−7.1%; 432 → 465 tok/s per user), 10.49 → **9.96 ms** at two (−5.0%). - Real acceptance (SGLang's default rejection sampling, greedy requests), paired over two rounds ABBA: accept length 3.29 → 3.30 at one request (+0.01 ± 0.03), 3.36 → 3.33 at two (−0.03 ± 0.02), within the servers' A/A spread (Accuracy Tests). Scope: - GPU: MI355X (gfx950), TP 4 or 8, as the target. - Steps: draft extends of 1 or 2 requests (4 or 8 rows; 1 request at TP4) and draft decode steps of 1 to 4 requests. Other draft steps, including draft decode steps of 5 to 8 requests, run SGLang's path (Reviewer notes). - A draft decode step whose proposal is sampled (rejection sampling with a batch that is not all greedy) runs SGLang's path; the draft extend's proposal is SGLang's on the tail's logits. - SGLang: no change. ## Modifications Three commits: the NextN builds (kernel and runtime), the draft spec and its forward, and the kernel checks. ### 1. The NextN entry, and a NextN launch from the step's tokens (`kernel.py`, `layout.py`, `runner.py`, `weights.py`) - **`Glm53Build.eh`** (with `tokens`): the layer starts from the NextN entry, `eh_proj(enorm(e) | hnorm(h))`. It is off by default; with it off, every existing build compiles to the same code object as the base's (Speed Tests, Registers). - e: the token ids' embedding rows, pushed by the rank whose vocabulary shard holds each, as #73's layer 0 does, into their own symmetric buffer (`ein`). - The entry: each rank computes its `HIDDEN / tp` columns of x from its rows of `eh_proj` (a row view of SGLang's bf16 weight), in 16-column tiles of 2 tasks each (the e half and the h half, f32 sums; the h half's partial goes to the e task's CTA through `ehp`), rounds them to bf16 and pushes them to every rank as tagged pairs (`xin`). The layer then reads x where layer 0 reads its embedding rows (the q_a / kv_a and index front tiles, and the attention residual). - The entry's tasks run on the last CTAs: at TP8 96 tasks on CTAs 160–255, at TP4 192 tasks on CTAs 64–255. The stages that share those CTAs (qkv_a and q_b tiles) wait for x anyway; only their weight prefetch starts later. - Everything after x is the indexer or shared kind's code, unchanged. - **Arguments:** the h rows ride in `cur_pos`, which a layer launch does not read; the entry's tensors (eh_proj's rows, the two norms, the embedding rows) ride in the op's own copy of the pointer table (`EH_ARGS`, after the indexer's entries), which the shared kind otherwise does not have. - **`Glm53Epochs`:** several runtimes' guard blocks in one tensor, so one `add_` advances all their step epochs. - **`head_weights(norm=)`:** #73's tail on another norm (the NextN's `shared_head.norm`). ### 2. The MTP draft steps on the NextN launches (`draft.py`, `tables.py`, `spec.py`, `check.py`, `__init__.py`) - **The draft spec** (`glm53_nextn_hip`), registered for `GlmMoeDsaForCausalLMNextN`. - **Routing.** `_build` puts a wrapper in front of the draft's `model.forward`; SGLang's runtime then installs its own forward over it, which hands every draft step to the wrapper (they never match the target's verify key). The wrapper decides on static properties of the step only (forward mode, batch size, rows, and whether the step reuses a top-k), so a captured graph replays the decision made at capture: - a draft extend (`DRAFT_EXTEND_V2`) of whole requests at 4 or 8 rows (4 at TP4) runs the NextN indexer launch; - a draft decode step (`DECODE`) of 1 to 4 rows that reuses the draft extend's top-k runs the NextN shared launch at the next width (2 or 4); - every other step runs SGLang's forward. - A draft-shaped step the launches cannot read (no hidden states, hidden states of another shape or dtype, no reused top-k) is refused during graph capture and runs SGLang's path eagerly, logged once a key. Each key the launches take is logged at its first launch. - The draft's attention backend is checked once per step mode, as the runtime checks a routed step's. - **A step.** The step tables (one launch, below), the NextN launch, the seed, then the tail: - the draft extend writes the NextN MLA and index-K rows of every row, then SGLang's seed (`dsa_seed_topk_capture`: each row's top-k as KV slots, then −1); - a draft decode row attends exactly the seed's valid entries, compacted in order (its count comes from the seed, not from its position); - the tail runs on `shared_head.norm` and the LM head SGLang shares with the target; the logits processor returns its f32 logits through the same `_get_logits` binding as #73's target (the draft extend graph's selected rows when `draft_extend_select_index` is set), and stores the tail's normed rows as the step's hidden states, which the next draft step reads. - The embedding and the LM head are SGLang's, read in place; SGLang shares them with the target after the draft is built, so a step binds the ones current at that step (a pointer compare on the host). - **Sampling.** The logits are the tail's, so SGLang's draft-extend proposal on them is SGLang's, greedy or sampled. A draft decode step whose proposal SGLang samples (rejection sampling, which this image enables by default on HIP for EAGLE topk 1, with a batch that is not all greedy) runs SGLang's path: its graph runner captures that variant separately. - **SGLang's path for the whole draft** (logged at startup): a speculative token map, no MTP index share (`index_share_for_mtp_iteration`), a seed width other than 2048, an indexer that is not SGLang's gfx950 fused one, a TP size the target does not run, or an `eh_proj` that is not a bf16 `nn.Linear`. - **Step tables in one launch** (`tables.py`, Triton): `step_rows` fills a width's `Glm53Step` (slots, positions clamped to the context, `indptr`, request rows, token ids, and with a seed the compacted top-k) in one launch, for the target's `_step` and the draft's; `write_seed` writes the seed. The words equal the torch ops they replace (unit tests under Triton's interpreter, kernel checks on the GPU). The target's step epochs advance in one launch (`Glm53Epochs`). - **Weights.** No new copies: eh_proj's rows are a view of SGLang's weight; the NextN layer's tensors are read and packed as a target layer's (`layer_weights`, with its W_UK / W_UV, router and indexer copies logged at startup). - **`SGLANG_AMD_MONO_CHECK`** (eager): each routed draft step against SGLang's NextN on the same inputs and caches: - the entry (`draft.eh`: x against SGLang's `eh_proj` on its `enorm | hnorm` rows; `draft.eh_vs_fp32`: against an f32 entry); - the NextN layer on the launch's own x, as a target layer's check but reported as `nextn.*`; - a draft decode row's keys against the seed's valid entries (`draft.topk_extra`), and the seed SGLang gets against the draft extend's top-k (`draft.seed`); - the tail against SGLang's `shared_head.norm`, LM head and all-gather and an f32 tail (#73's `head.*` bounds, logged as `draft tail`); - the draft tokens: the tail's top-1 against SGLang's whole NextN (its own entry, the layer routed as SGLang routes, its norm and LM head) on the same inputs (`draft.token.agree`, reported, no bound), and SGLang's logit gap where they differ (`draft.token_gap`). - `SGLANG_AMD_MONO_DRAFT_FAULT={swap_norms,count_by_position,seed_row0,rope_shift}` is a test fault for the negative controls (Accuracy Tests). - **`SGLANG_READS`** of `draft.py`: the NextN model's modules, the forward batch's draft fields (`spec_info.hidden_states`, `dsa_topk_indices`, `dsa_seed_topk_capture`, `reuse_dsa_topk_indices`, `sampling_info.is_all_greedy`), the logits metadata's `draft_extend_select_index`, the spec arguments read, and `get_dsa_mtp_topk_width`. ### 3. Kernel checks of the NextN launches (`test/kernels/glm53/check_nextn.py`, `test_glm53_kernels.py`) - `check_nextn.py`, under torchrun at TP4 or TP8, on random weights in SGLang's layouts, NextN index at 4 and 8 rows (draft extend rows, 4 a request; 4 at TP4) and NextN shared at 2 and 4, each with and without a pad row, bf16 or E4M3 MLA cache: - **entry:** the NextN launch against the plain launch of its kind fed the NextN launch's own x (its gathered `xin` rows): output rows, MLA cache rows, index cache rows and top-k, bitwise; x within 1e-4 of a torch entry and 7e-3 of an f32 entry; a pad row's x zero. The enorm and hnorm weights swapped must fail both. - **tables:** `step_rows` and `write_seed` against the torch ops, word for word (SGLang's padding slot 0, pad rows, positions past the context, −1 inside a seed); the seed sized by position must differ. - **replay:** a draft step (NextN launch and tail) in a graph, 20 replays bitwise equal to the eager step, no error word. - `test_glm53_kernels.py` runs it at TP4 and TP8 on both MLA caches. - CPU unit tests (`test/unit/test_glm53_draft.py`): the registry entry; the steps the wrapper routes, at TP8 and TP4, positional and keyword; the steps it leaves to SGLang (other modes, more rows, a sampled draft decode); the capture-time refusal and the eager fallback logged once; the runtime's forward over the wrapper; the backend check once per mode; every SGLang call site of `model.forward` the wrapper relies on; the tail's logits through `_get_logits`; each setting that keeps the draft standard; the entry, seed and decode-row metrics with their negative controls (real `check.py` functions); the step tables against the torch ops under Triton's interpreter, in a child process (`test/unit/triton_interpreter.py`), as the target's `_step` tests now run too. ## Accuracy Tests The draft is not bitwise equal to SGLang's NextN, and is not meant to be: it runs the target layers' numerics (bf16 attention against the FP8 cache, f32 routing). Everything else is bitwise or exact: - **NextN launch = plain launch fed its own x, bitwise** (`check_nextn.py`, entry): outputs, MLA rows, index rows and top-k, at TP8 (NextN index 4 / 8 rows, NextN shared 2 / 4, each with and without a pad row; 56 checks per cache) and TP4 (index 4, shared 2 / 4; 38 checks per cache), bf16 and E4M3 MLA cache: all pass. The entry against a torch entry: 3e-6 to 3e-5; against f32: 2.3e-3, the torch entry's own distance; pad rows zero; the swapped-norm control gives 0.14 against both bounds. The step tables equal the torch ops word for word; 20 graph replays of a draft step are bitwise equal to the eager step. - **Every existing build is unchanged:** the compile survey's 90 existing keys (indexer, shared, dense, token and tail; TP4 2 / 4, TP8 2 / 4 / 8; both caches; `stop` on and off) give the same code objects as the base. Seven of them compile nondeterministically on both trees (the base alone gives two different objects across reruns, with equal register counts); in reruns each head object equals one the base produced. - **`SGLANG_AMD_MONO_CHECK=1` with MTP on** (TP8, FP8 E4M3 KV, eager, 1 and 2 requests at 1K and 8K, 48 tokens each; 240 draft steps: draft extends at 4 and 8 rows, draft decode steps at 1 and 2): **0 violations**, target layers included. | Metric | Measured | Bound | Negative control (must fail) | | --- | --- | --- | --- | | `draft.eh.rel` | max 4.5e-4 | 1.4e-3 | `swap_norms`: 0.57–0.79 | | `draft.eh_vs_fp32.rel` | max 2.8e-3 (SGLang's own: 2.8e-3) | 8.8e-3 | `swap_norms`: 0.57–0.79 | | `draft.topk_extra.max_abs` | 0 | 0 | `count_by_position`: 1–2 | | `draft.seed.max_abs` | 0 | 0 | `seed_row0`: 4–2048 | | `nextn.streams.cosine` | min 0.963 | 0.89 | `rope_shift`: min 0.62 (110 of 264 steps) | | `nextn.kv.cosine` | min 0.99936 | 0.995 | `rope_shift`: 0.36–0.86 | | `nextn.index.cosine` | min 0.99965 | 0.999 | `rope_shift`: from 0.85 | | `nextn.topk_exact` / `nextn.topk_own` | 0 / 0 | 0 / 0 | the target's | | draft tail `head.logits.rel` / `_vs_fp32` | max 4.4e-4 / 2.2e-3 | #73's 1.7e-3 / 7.5e-3 | #73's | | `draft.token.agree` (reported) | drafts 1 / 2 / 3: 95.9 / 95.5 / 99.1% of rows | — | — | Each control ran on its own server and failed only its own checks (`rope_shift` also fails `nextn.kv` and `nextn.index`). - **Why the stream bound is loose:** a target layer's stream is the residual of every layer before it, so one layer's error moves it little. The NextN stream is a fresh x plus one layer's attention and MLP. It carries the full distance of SGLang's NextN attention, which runs FP8 q and P on this cache where the launches run bf16. The tight checks on the same path are the bitwise kernel check and the entry against f32. - The draft tokens differ from SGLang's NextN on 1–5% of rows, at SGLang logit gaps up to 1.25; these are the two attentions' numerics. - **Accept length at real acceptance**, paired (TP8, FP8 E4M3 KV, 8,192-token ShareGPT prompts, 512 greedy tokens, SGLang's default rejection sampling; two rounds ABBA, one server per arm and round, the same prompts on every server): | Requests | n | base | this PR | change (± s.e.) | A/A base / this PR | | --- | --- | --- | --- | --- | --- | | 1 | 80 | 3.291 | 3.304 | +0.013 ± 0.032 | −0.032 ± 0.053 / −0.017 ± 0.047 | | 2 | 160 | 3.363 | 3.328 | −0.035 ± 0.022 | +0.011 ± 0.027 / −0.030 ± 0.042 | - **E2E** (`test/e2e/test_glm53_server.py`, TP8, this head): - `TestMonokernelCheck`: 4 passed. Its MTP cases (bf16 and E4M3 KV) now also check every draft step (`check.draft`) and assert drafts 1, 2 and 3 were checked; no violation. - `TestMonokernelMtp`: passed. GSM8K 0.985, accept length 3.28, and the draft extends of 1 and 2 requests ran the NextN launches (the test now asserts their log lines; these requests sample, so under rejection sampling their draft decode steps run SGLang's path). - `TestMonokernelDecode` (2) and `TestMonokernelHiddenStates`: passed. - `TestMonokernelFault` fails here on `main` the same way, only at its last check: at the 10 s deadline every rank stops with `HIP error: an illegal memory access was encountered` and the server exits before any token of the failed step reaches the client, but ROCr in this container prints no `Memory access fault ... (nil)` line, so `assertTrue(faults)` fails; in the reviewer's setup on #68 and #69 the line is printed and the case passes. - **GSM8K with MTP at two requests at a time** (greedy, thinking on, 200 questions, as the e2e test's GSM8K but at two requests, so every draft step of 1 or 2 requests runs the launches): 0.990, accept length 3.38; the draft extends of 1 and 2 requests and the draft decode steps of 1 to 4 ran the NextN launches. - **TP4:** an MTP server (FP8 KV, golden AL 2.99, one request at 8K) serves with the NextN launches: the draft extend at 4 rows and draft decode steps of 1 to 4 rows ran them; cycle 8.93 ms, accept length 2.99; NextN kernel checks as above. - Unit tests on a CPU-only host: 335 passed (base: 289); the base's 2 failures and 3 collection errors (minimax, deepseek_v4). ## Speed Tests and Profiling Setup as #68 (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`). All numbers are from this head (a5bfe9d) against the base (b6fdf8e, #74's head). Serving as #68: one fresh server per arm, 8,192-token ShareGPT prompts, 512 greedy output tokens, MTP at the golden accept length 2.99 (measured 2.98–2.99 on every server), **out-of-tree streaming client**. One round. At C≥2 the server prefills a wave's requests one after another, and a request that has started decoding waits out the later prefills. The C≥2 numbers below are those of each wave's last-prefilled request, which decodes beside the others from its first token to its last. TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms: | Requests | base | This PR | Change | tok/s per user at AL 2.99 | | --- | --- | --- | --- | --- | | 1 (draft extend 4 rows, draft decode 1 row) | 6.927 | 6.435 | −7.1% | 432 → 465 | | 2 (8 rows, 2 rows) | 10.488 | 9.962 | −5.0% | 285 → 300 | MTP off is unchanged: the target's launches are the base's code objects. Its step tables are one launch instead of the torch ops' kernels, and its epochs one launch. At real acceptance (the ABBA servers above, median TPOT): 2.12 → 1.93 / 1.96 ms at one request, 3.04 / 3.03 → 2.96 / 2.95 ms at two (−2.6%). ### Registers The compile survey (CPU-only, FlyDSL compile-only): every key the spec and the tests build, TP4 (2, 4 rows) and TP8 (2, 4, 8), both caches, `stop=True` (spec) and `stop=False` (tests). The existing keys are unchanged (Accuracy Tests). The NextN keys, VGPR spills / private segment bytes, the ones that have either: | Build | `stop=True` (spec) | `stop=False` (tests) | the plain kind's (`stop=True` / `False`) | | --- | --- | --- | --- | | TP8 NextN index, 8 rows | 12–16 / 52–68 | 12–16 / 52–68 | 20–24 / 84–100 | | TP4 NextN index, 4 rows | 16 / 68 | 56 / 228 | 13–15 / 56–64, 48 / 196 | | TP4 NextN shared, 4 rows | 24–26 / 100–108 | 40–48 / 164–196 | 0 / 0 | Every other NextN key (TP8 index 4 rows, shared 2 / 4 rows; TP4 shared 2 rows) has no VGPR spill and no private segment. Every symbol is `glm5_monokernel_0`, as on the base. The TP4 NextN shared spill at 4 rows is the entry's 16-column tile next to the 4 rows; it serves TP4's draft decode steps of 3 and 4 requests. ## Reviewer notes - **Draft decode steps of 5 to 8 requests run SGLang's path.** In `check_nextn.py`, the NextN shared launch at 8 rows on the E4M3 cache faulted (illegal address) as the ninth NextN model of one process, after eight others were built, run, replayed and closed. Alone in a fresh process it runs clean, both caches, NextN or plain first. The plain shared launch at 8 rows never faulted (the kernel tests, this head and the base). Until that is root-caused, the spec does not build the NextN shared kind at 8 rows. MTP at 1 to 4 requests is unaffected. - **The draft-step wrapper**, as in #29. It relies on SGLang's draft call sites calling the instance attribute `model.forward`; a unit test pins them. It can go when the monokernel API routes a draft model's own step modes: then the spec declares them, and the wrapper and its capture-time refusal go. - **`_get_logits`** is the one SGLang method the draft spec replaces, as #73's target does; the binding is active only inside a routed step. With `draft_extend_select_index` set (the draft extend graph), it returns the selected rows, as SGLang's pruning would hand them. - **Eager and graph steps share the per-width step tables.** Every routed step, eager or replayed, rewrites all of them in the same stream before its launch (the step-table launch is inside the graph), so an eager step between replays changes nothing a replay reads. - **Kernel test harness:** `check_nextn.py` builds several models in one process. For the runs above, both trees (base and head alike, not part of this PR) carried a runner-only patch in which `Glm53Runtime.close()` keeps its uncached buffers. In a process that closes runtimes and builds new ones, freed uncached buffers corrupt later allocations; this is root-caused separately. Serving never closes a runtime. - **Copies:** the draft copies its own rotary table (128 MiB a rank at the served context) and packs its layer's W_UK / W_UV, router and indexer weights, as a target layer does (logged at startup). - `kernel.py` grows by 152 lines (the entry, the `ein` buffer and the pushed-x switch), a candidate for the stage-builder split. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
On
main, after #68 (merged). It uses #68'stest/kernels/glm53/tools.After #68, the attention chain of a glm53 launch (qkv_a → q_norm → q_b → uk → split → uv → o) still waits on stragglers: one cache task does every row's KV row in turn and the split tile behind it starts late; split tiles land on busy CTAs; a split loads its keys and this launch's new rows in serial round trips after q arrives; the indexer's front tiles share CTAs with qkv_a; at 8 rows the front runs twice and the MoE stages spill. The routing hand-off polls the expert activations in two or three batches. This PR removes those. Every output and every write stays bitwise equal to the base at TP4 and TP8, at every width.
Result on MI355X, base (#68) → this PR:
main: MTP C=1 cycle 8.49 → 7.46 ms (−12.1%), C=2 12.57 → 11.01 (−12.4%); MTP-off TPOT at one request 5.48 → 4.78 ms (−12.9%).time_layers.py --kind model(the 75 launches), E4M3 MLA cache, µs a launch at 2 / 4 / 8 rows: 59.23 → 54.21 (−8.5%) / 76.75 → 70.92 (−7.6%) / 126.64 → 115.55 (−8.8%) at TP8; 76.76 → 70.07 (−8.7%) / 110.68 → 102.19 (−7.7%) at TP4, 2 / 4 rows.main, TP8 model order: 64.19 → 54.21 (−15.5%) / 84.96 → 70.92 (−16.5%) / 135.62 → 115.55 (−14.8%).Modifications
Six commits, each measured on top of the one before (Speed Tests, "Each change on its own", which also gives each one's share of the total).
1. One cache task a row, and the attention stages placed past busy CTAs (
layout.py,kernel.py)layout.stage_bases, used by the kernel.2. split loads its keys and this launch's rows before the q wait (
kernel.py)A split tile loaded its rows' slots and the rows this launch appends only after q arrived, and each wave loaded its 8 keys' cache rows one after another, then polled each new key's row in turn (at 4 verify rows the last row's tile holds all four new tokens). Now the slots load at the task's start; every key's cache row is issued first; if any key of the task is new, the CTA polls every new row once, behind those loads, into the q region of LDS (not used yet), and each wave copies its keys' rows from there; then the conversions and stores, all before the q wait. Same values.
3. The indexer's front tiles on their own CTAs, and 8 rows in one pass (
kernel.py,layout.py)4. The expert activations in one poll, after the routing (
kernel.py)The router published its FP8 activation blocks one tagged word (4 values) an 8-byte store, and every up/gate CTA polled each row's words and scales in 2 or 3 batches in series. The blocks now go out two tagged words a 16-byte store (
put2takes int32 bits too), and a CTA polls every row's words and all the scales in one batch. A CTA that ran a router task reaches up/gate after the activations are out: it now routes first, issues its first routed tile's weights, and stages the activations behind them; the others stage them while the router still runs. The mailbox and its values are unchanged.Its own contribution is about zero: model order −0.36 / −0.08 / +0.31 µs at 2 / 4 / 8 rows, within the passes' spread. It stays because changes 5 and 6 were built and measured on top of it.
5. At 8 rows the MoE stages take the wave index as a scalar (
kernel.py)From the router on, at 8 rows, the wave index is read once into an SGPR, so the per-wave unit, row-group and address terms of the router, up/gate and down are scalar arithmetic. Below 8 rows it stays a vector value (scalar, the terms then share the CU's scalar unit in down's MFMA loop, which measured slower).
6. The expert block scales as scalar loads (
kernel.py)Up/gate and down loaded each unit's f32 block scale with a vector load at a wave-uniform address: a VGPR and a vector-memory instruction a unit, in the weights' in-order queue. They are now scalar loads from the constant address space (
sld_f32, an LLVM load from address space 4, volatile so that it issues with its unit's weights instead of being sunk to the MFMA); up/gate loads the gate and up rows' scales and picks one by lane.Accuracy Tests
check_layer.py --dump: kinds shared, index and the model order's first 8 launches; bf16 and E4M3 MLA cache; every width the TP builds; one request a row, the same with a pad row, and MTP's verify rows; 3 steps): TP8 41,760 digests, TP4 10,560, 0 differ. So it is also bitwise equal tomain.test/e2e/test_glm53_server.py, TP8): this PR's launches are bitwise equal to the base's at every TP, width and cache dtype, so the decode-vs-prefill, check and GSM8K cases are carried by [Perf] [MI355X] glm53: Cheaper hand-offs and placement: the shared expert before the routing, uncached scratch, tagged selected tokens (1/2) #68's run on its head. On this head only the injected-fault case ran, because this PR changes the guarded wait sites, so fault injection's site numbering moves. It fails onmainthe same way, both with the SGLang overlay and with the image's own SGLang (b48809dfc2, main'ssgl_amd_mono/__init__.pyfrom before [Feat] [MI355X] dsv4: Add whole-layer DeepSeek-V4-Pro mono kernels (TP8) #48): no fault at address 0x0 reaches the log (the test'sassertTrue(faults)fails). Every rank stops at the wait deadline withHIP error: an illegal memory access was encountered, and the server exits without a token; ROCr in our container prints noMemory access faultline.Speed Tests and Profiling
Setup as #68 (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). All numbers are from cfba6f0 against a95d4d9 (#68's head). The rebase ontomainchanges neitherkernel.pynorlayout.py.Serving as #68: one fresh server per arm and setting, all of a setting's points on it, 8,192-token ShareGPT prompts, 512 greedy output tokens, MTP at the golden accept length 2.99 (measured 2.99 on every server); out-of-tree streaming client (
sglang.benchmark.servingdoes not import on this image). One round (a later second server of the same arm agreed within 0.3%; the kernel-only passes below agree within 0.2 µs).At C≥2 the server prefills a wave's requests one after another, and a request that has started decoding waits out the later prefills. The C≥2 numbers below are those of each wave's last-prefilled request, which decodes beside the others from its first token to its last.
TP8, FP8 E4M3 KV, MTP 3-1-4 at AL 2.99, cycle in ms:
TP8, FP8 E4M3 KV, MTP off, median TPOT in ms:
Serving was measured at TP8 on the E4M3 cache; the bf16 cache and TP4 are measured kernel-only below.
Kernel-only (
time_layers.py)µs a launch, base → this PR, the mean of two passes in alternating order in one allocation.
TP8, E4M3 MLA cache:
TP8, bf16 MLA cache, model order: 57.09 → 54.30 / 76.17 → 71.27 / 128.33 → 115.42. The kernel now runs the bf16 and the E4M3 cache at the same speed (54.21 / 70.92 / 115.55 on E4M3); the base's two caches were 2.1 / 0.6 / −1.7 µs apart.
TP4 (2 and 4 rows), E4M3 unless marked:
Each change on its own
Each commit against the one before it, TP8, E4M3 MLA cache, µs a launch at 2 / 4 / 8 rows (two passes in alternating order, one allocation; TP8 only). The last column is each change's share of this PR's model-order gain.
Registers
The compile-only survey of #68 (40 builds: TP4 2 / 4 rows, TP8 2 / 4 / 8; indexer and shared; bf16 and E4M3;
stop=Trueandstop=False). Every symbol isglm5_monokernel_0, as on the base. VGPR spills / private segment bytes, base → this PR, the builds that have either:stop=True(spec)stop=False(tests)Tests
test_glm53_kernels.pywith 8 GPUs (TP4 and TP8): 3 passed (3 min 20 s).Checklist
pre-commit run --all-filespasses.test/cover the change..github/CODEOWNERSthat names its owner. (No new model.)