Repository navigation
[Perf] [AMD] monokernels: Let a plugin skip per-step attention metadata - #44601
Merged
kevin-mii merged 10 commits intoOct 18, 2026
Merged
Conversation
A spec declares the per-step attention metadata its forward never reads (MonoSpec.unread_metadata, a set of StepMetadata). On a step the forward switch routes, the HIP radix backend skips building it: sgl-amd-mono's DeepSeek-V4.1 forward scores every index source with K2i, so the aiter decode schedule of each low ratio goes unread. Co-authored-by: kkHuang-amd <43161300+kkHuang-amd@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
4 tasks done
The backend's added comment line restated is_metadata_unread, and the StepMetadata docstring repeated the unread_metadata field comment. The test names the model without monokernels instead of a bare Lm(). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii
marked this pull request as ready for review
October 16, 2026 00:19
kevin-mii
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
hebiao064,
ispobock and
merrymercy
as code owners
October 16, 2026 00:19
Collaborator
Author
|
/tag-and-rerun-ci |
This was referenced Oct 16, 2026
HaiShaw
approved these changes
Oct 16, 2026
kevin-mii
enabled auto-merge (squash)
October 16, 2026 08:18
kevin-mii
added a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 16, 2026
…never reads (#5) ## Motivation SGLang lets a spec list the per-step attention metadata its forward never reads (`MonoSpec.unread_metadata`, sgl-project/sglang#44601), and the HIP radix backend skips building it on the steps the monokernels run. The DeepSeek-V4.1 forward scores every index source with K2i (`_select`), so the aiter decode schedule of each low ratio goes unread on those steps. Co-authored with @kkHuang-amd: this reshapes 3970534 (`_Built.skips_metadata`) into spec data, the form sgl-project/sglang#44601 reads. ## Modifications - `SPEC.unread_metadata = {StepMetadata.LOW_RATIO_DECODE_SCHEDULE}`, in place of `_Built.skips_metadata`, which no SGLang calls. - `SGLANG_READS` declares `StepMetadata` and `unread_metadata`. This imports `StepMetadata`, so it needs sgl-project/sglang#44601: merge it after that one, then bump SGLang's `srt_hip` pin to the merged commit. ## Accuracy Tests MI350X, TP4, `test/e2e` on merged SGLang main (ec7ab9d0c2): | | SGLang | sgl-amd-mono | result | |---|---|---|---| | main, the pinned plugin | ec7ab9d0c2 | 8f1bff3 | 9 passed | | the skip as first written | ec7ab9d0c2 + kevin-mii/sglang#19 | 3970534 | 9 passed | | this PR | sgl-project/sglang#44601 (9fd4430dd3) | this change (ae2f020, on 3970534) | 9 passed | `test/unit`: 64 passed against sgl-project/sglang#44601. ## Speed Tests and Profiling See sgl-project/sglang#44601: at 1 request, decode time per output token is 0.4-0.5% lower at 8k context and 1.8-2.9% lower at 128k; at 4 requests the change is within noise. ## 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. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: kkHuang-amd <43161300+kkHuang-amd@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
sgl-project/sgl-amd-mono#5 merged ahead of this PR; its DeepSeek-V4.1 spec lists StepMetadata.LOW_RATIO_DECODE_SCHEDULE in MonoSpec.unread_metadata. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
github-actions
Bot
added
the
dependencies
The image pinned 815b104 (sgl-project#44621) while the srt_hip extra pins e07665c, so the image would serve the plugin without the unread decode schedule. Each pin now names the other, since they must move together. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The ROCm image installs SGLang with the srt_hip (or srt_hip_rocm724) extra, which already installs the plugin at its pyproject pin. sgl-project#44621's clone and editable reinstall then put a second pin on top, so the image and pip could serve different plugin commits. Drop it: the extras' pin is the only one. The plugin declares no dependencies, so the extras install matches the old --no-deps install. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Restores sgl-project#44621's install step (editable checkout, registration check, gfx950 gating, opt-out) that the previous commit removed, and takes its commit from the srt_hip extras of the SGLang being built unless SGL_AMD_MONO_COMMIT overrides it, the way CI reads AITER's pin from this Dockerfile. The image and pip then install the same plugin commit, and a bump is one change in pyproject_other.toml. The build fails if the extras pin different or non-SHA commits, and skips SGLang versions that pin none. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The pin moves into one monokernels_hip extra that srt_hip and srt_hip_rocm724 both include, so a bump is one line and the two cannot disagree. The ROCm image's install step takes its commit from the plugin's direct_url.json, which pip writes for git installs, instead of parsing pyproject_other.toml; SGL_AMD_MONO_COMMIT still overrides it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
4 tasks done
4 tasks done
kevin-mii
pushed a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 18, 2026
## Motivation On MI355X at TP4 with one request, a GLM-5.3-Flash decode step on SGLang's path runs 1,301 kernels (torch profiler). Its median TPOT is 7.75 ms. vllm-project/vllm#61662 runs each MoE decoder layer of this model as one persistent launch on gfx950. This PR ports that kernel to the `amd` monokernel backend of SGLang (RFC sgl-project/sglang#43949). With `--monokernel-backend amd`, each covered layer of a decode step runs as one launch. Result on MI355X: - At TP4 with one request, the median TPOT decreases from 7.75 ms to 4.40 ms (1.76x). - From 1 to 8 requests (1 to 4 at TP2), the speedup is 1.31x to 2.09x. - Accuracy does not decrease. On AIME26 (30 problems x 8), pass@1 is 89.17% (off) and 91.67% (on). On GSM8K, it is 97.35% (off) and 97.27% (on). Scope: - GPU: MI355X (gfx950), at TP 2, 4 or 8. - Steps: decode steps of 1 to 8 rows (1 to 4 at TP2). All other steps use SGLang's path. - Layers: the MoE decoder layers 3 to 43 (41 launches: 30 KDA, 11 sparse attention). The dense layers 0 to 2 and the last layer (44) use SGLang's path. - SGLang: this PR needs sgl-project/sglang#44353 (merged as ec7ab9d0c2). It changes no SGLang file. ## Modifications ### How a decode step runs ```mermaid flowchart TB F["SGLang monokernel framework<br/>(sglang PR 44353)"] --> Q{"Decode step of 1 to 8 rows?<br/>(1 to 4 at TP2)"} Q -- "no" --> N["SGLang's path"] Q -- "yes" --> R["glm53_flash spec forward"] R --> A["SGLang: embedding and<br/>dense layers 0 to 2"] A --> L["41 launches, layers 3 to 43<br/>(30 KDA, 11 sparse attention)"] L --> Z["SGLang: layer 44,<br/>final norm and logits"] L -.->|reads and writes in place| M[("SGLang: KV, index and KDA state pools,<br/>page table, loaded weights")] classDef sgl fill:#eef1f4,stroke:#8c959f,color:#1f2328 classDef amd fill:#dafbe1,stroke:#1a7f37,color:#1f2328 class F,Q,N,A,Z,M sgl class R,L amd ``` Grey boxes are SGLang code, and this PR changes none of it. Green boxes are new in this PR. One launch does these operations for one layer, in this sequence: 1. The attention-side mHC. 2. The attention: KDA, or sparse MLA with its kpool indexer and its cache writes. 3. `o_proj` and its TP all-reduce. 4. The FFN-side mHC and norm. 5. The router. 6. The shared and routed MXFP4 experts. 7. The MoE TP all-reduce. ### Changes to SGLang (`python/sglang/srt`) - None. This PR adds and changes files in sgl-amd-mono only. - Without `--monokernel-backend amd`, SGLang runs as before. The spec registers for gfx950 only, so CUDA and other ROCm GPUs do not change. - With the flag, SGLang calls the spec's forward, not `Glm5NextForConditionalGeneration.forward`, for each decode step that the spec covers. SGLang still does the scheduling, the cache management, the sampling and all other steps. - The plugin reads SGLang internals by name. Each module lists these names in `SGLANG_READS`. `test/unit/test_contract.py` checks them on CPU against the installed SGLang. - At startup, the plugin uses more GPU memory for launch buffers and weight copies (see "Startup"). ### Files - `sgl_amd_mono/glm53_flash/gfx950/`: the decoder-layer kernel from vllm-project/vllm#61662 at ca328e297b. vLLM adapted it from ROCm/aiter#6173. The first GLM-5.3-Flash commit copies the files verbatim. The next commit adapts them: - Every wait runs under `common/guard.py` (deadline, peer flags, fail-stop). - The kernel reads the KDA conv and recurrent state pool of SGLang in place. It also reads the DSA MLA and index caches, the compress tails and the 64-token page table in place. - The page table is a launch argument, so each captured graph reads its own table. - The kernel reads the FP8 attention projections in the AITER (16, 16) layout that SGLang's gfx950 block-FP8 linear method makes. It reads the MXFP4 experts in the layout that SGLang's Quark path makes. - Like the vLLM kernel, this kernel quantizes the expert inputs to MXFP8. SGLang's Quark path quantizes them to MXFP4 (W4A4, as the checkpoint declares). Thus the two paths are not bit-exact. - `sgl_amd_mono/glm53_flash/spec.py`: the spec and its forward. The forward is `Glm5NextForConditionalGeneration.forward` for a covered decode step. - A step with fewer rows than a kernel width (2, 4 or 8) runs at the next width. The added rows are pad rows, and pad rows write no cache. - At TP2, an 8-row sparse-attention launch does not fit in the LDS, so steps of 5 to 8 rows at TP2 use SGLang's path. - The framework sets `DeviceLimits.widths` for each ISA, not for each TP size. Thus at TP2, the framework still lists 5 to 8 rows as covered. - The spec refuses unsupported setups at startup. The config and model checks give their reasons: page size, KV dtype, MoE runner, `--enable-linear-replayssm`, rotary layout, and others. - If the kernel cannot read a weight layout (for example a bf16 shared expert), the build stops and gives the reason. - `sgl_amd_mono/glm53_flash/check.py`: with `SGLANG_AMD_MONO_CHECK=1` (eager mode only), each covered layer also runs SGLang's layer on the same input and caches. The check compares the written streams, each state row and each cache row with the bounds in `NUMERICS`. It logs each value that is outside its bound. - `sgl_amd_mono/common/`: three fixes for problems that the GLM-5.3-Flash launches found: - The guard reads its peer flags in loads of 4 dwords or fewer. At TP8, the 8-dword load failed instruction selection. - Peer buffers use whole 2 MiB granules. Before this fix, a peer faulted on guard flags after the last whole granule. - The peer handshake fails on all ranks when the buffer sizes are different. Before this fix, ranks that built buffers in different orders waited until the deadline. - Registration in `sgl_amd_mono/__init__.py`, a README row, a NOTICE entry, and these tests: - `test/unit/test_glm53_flash_spec.py` (CPU) - `test/kernels/test_glm53_flash_kernels.py` (8 GPUs) - `test/e2e/test_glm53_flash_server.py` (TP4) ### Not in this PR - An FP8 KV cache. - Speculative decoding (MTP). - Decode steps of more than 8 rows. ## Accuracy Tests Setup: - MI355X, TP4, amd/GLM-5.3-Flash-Quark-MXFP4 at b5688f25, sgl-eval 0.1.2. - The two arms ran on the same GPUs. Each arm used a fresh server with CUDA graphs and 8 or fewer requests at a time. Thus all decode steps of the "on" arm ran the launches. | Eval | Settings | Monokernels off | Monokernels on | | --- | --- | --- | --- | | GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 97.35% | 97.27% (13 vs 14 discordant, exact McNemar p = 1.00) | | AIME25 (30 x 4 seeds) | thinking, T 0.6, top-p 0.95, 65,536 tokens | pass@1 82.50% | pass@1 91.67% | | RULER2 NIAH at 1M tokens (4 subtasks x 2) | no thinking, greedy, 256 tokens | 40.62% | 44.12% | - With the monokernels, AIME25 is 9.2 points higher. The paired bootstrap 95% CI over the 30 problems is [+1.7, +17.5] points. A sign test over the problems gives p = 0.09 (10 problems better, 3 worse). - 14 of 120 samples reached the token limit in the "off" arm, and 6 in the "on" arm. The median output was 7,348 tokens (off) and 5,051 tokens (on). - A possible cause of this difference is the MXFP8 expert inputs. I did not isolate it. Per layer, against SGLang's own layer on the same input and caches: - Setup: `SGLANG_AMD_MONO_CHECK=1`, eager mode. Batches of 1, 3 and 8 rows (4, not 8, at TP2). Positions up to 2.4k. 779 checks for each TP size. - Result: no value is outside its bound. | Output | TP2 | TP4 | TP8 | | --- | --- | --- | --- | | written streams, cosine | >= 0.974 | >= 0.992 | >= 0.980 | | MLA KV rows, cosine | >= 0.9991 | >= 0.9992 | >= 0.9993 | | index cache values, cosine (scales bit-exact) | >= 0.99998 | >= 0.99998 | >= 0.99998 | | KDA recurrent state, cosine | >= 0.999999 | >= 0.999999 | >= 0.999998 | The streams differ by approximately the error of the MoE: the kernel uses MXFP8 expert inputs, and SGLang's path uses MXFP4. Decode against prefill (`sglang.test.kl_test_utils`, 8 prompts x 512 tokens, mean KL): - SGLang's path: 0.038. - Monokernels on: 0.009 to 0.042 for each batch, over two e2e runs. - The e2e test sets the bound to 0.06. It runs one graph for each batch of 2, 3, 7 and 8 rows. - With an injected fault (one page-table pointer for all graphs), the 3-row batch measured 0.146. GPQA-Diamond did not run, because the dataset is gated on the Hub. ### AIME26 on daf4739 GSM8K is near its ceiling for this model, so it finds only large errors. The SGLang contribution guide asks for a harder eval, and the GLM-5.3 cookbook reports AIME25 and AIME26. Thus this PR also reports AIME26 on its final head. - Setup: MI355X, TP4, SGLang main 4cca7fcd3e, sgl-amd-mono daf4739, sgl-eval 0.1.3 (the version that SGLang pins). Each arm used a fresh server. - 8 or fewer requests ran at a time, so each decode step of the "on" arm ran the launches. - Generation: thinking, T 1.0 and top-p 0.95 (the recommended settings of the model), 65,536 max tokens. | AIME26 (30 problems x 8) | Monokernels off | Monokernels on | | --- | --- | --- | | pass@1 | 89.17% | 91.67% | | majority@8 | 96.67% | 96.67% | | samples at the token limit | 14 of 240 | 12 of 240 | | median output tokens | 8,463 | 6,516 | - Paired over the 30 problems, the "on" arm is 2.50 points higher (bootstrap 95% CI [+0.0, +5.4]). 9 problems are better, 3 are worse and 18 are tied (sign test p = 0.15). - The direction is the same as on AIME25: higher accuracy and shorter outputs with the monokernels. A possible cause is the MXFP8 expert inputs. ## Speed Tests and Profiling Setup: - Hardware: one node with 8 x MI355X (gfx950). - Image: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014` (`sha256:e4498f5d595f`), with ROCm 10.0.0, torch 2.11.0, FlyDSL 0.3.4.1 and AITER e7d2453f (= `AITER_COMMIT_DEFAULT`). - SGLang: sgl-project/sglang#44353 at 04fcb04e34. Its monokernel framework is the same as the merged ec7ab9d0c2. - sgl-amd-mono: this PR on 3970534. The rebase onto 815b104 changed no code. The later commits change no launched code: - One commit moves an import below the TP2 fallback. - daf4739 refuses `--enable-linear-replayssm` and corrects comments and docs. - The e2e tests ran again on daf4739 (see the end). - Server: the cookbook's MI355X high-throughput cell (no speculative decoding), with these differences: - This checkpoint. The cookbook's OneNexus/GLM-5.3-Flash-MXFP4 keeps the shared expert in bf16, and the spec refuses that. - `--kv-cache-dtype bfloat16`. - `--ep-size 1`. - TileLang DSA prefill and decode. - `--moe-runner-backend aiter`. - The only difference between the arms is `--monokernel-backend amd`. Method: - Workload: `sglang.bench_serving`, random dataset with fixed lengths (`--random-range-ratio 1.0`), max(16, 4c) prompts after a warmup. - Each arm ran on a fresh server, in two rounds. The second round swapped the arm order and the GPU set (TP8: the order only). - At load, each server checked that no other process held memory on its GPUs. - The table gives the median TPOT per request in ms over both rounds, with a bootstrap 95% CI. | TP | In/out | Conc. | Monokernels off | Monokernels on | Speedup | | --- | --- | --- | --- | --- | --- | | 2 | 8192/256 | 1 | 9.31 [9.28, 9.31] | 5.68 [5.68, 5.68] | 1.64x | | 2 | 8192/256 | 2 | 9.58 [9.57, 9.60] | 5.73 [5.72, 5.75] | 1.67x | | 2 | 8192/256 | 4 | 10.25 [10.23, 10.27] | 7.15 [7.13, 7.18] | 1.43x | | 2 | 512/128 | 1 | 9.17 [9.14, 9.18] | 5.48 [5.46, 5.48] | 1.67x | | 2 | 512/128 | 4 | 10.31 [10.26, 10.33] | 6.88 [6.86, 6.91] | 1.50x | | 4 | 8192/256 | 1 | 7.75 [7.74, 7.75] | 4.40 [4.40, 4.40] | 1.76x | | 4 | 8192/256 | 2 | 7.98 [7.98, 8.00] | 4.43 [4.42, 4.49] | 1.80x | | 4 | 8192/256 | 4 | 8.39 [8.38, 8.40] | 5.03 [5.02, 5.03] | 1.67x | | 4 | 8192/256 | 8 | 9.23 [9.22, 9.49] | 7.03 [7.01, 7.39] | 1.31x | | 4 | 512/128 | 1 | 7.66 [7.64, 7.66] | 4.16 [4.16, 4.17] | 1.84x | | 4 | 512/128 | 2 | 7.91 [7.91, 7.92] | 4.18 [4.18, 4.19] | 1.89x | | 4 | 512/128 | 4 | 8.35 [8.33, 8.36] | 4.83 [4.82, 4.83] | 1.73x | | 4 | 512/128 | 8 | 9.39 [9.35, 9.41] | 6.76 [6.74, 6.77] | 1.39x | | 8 | 8192/256 | 1 | 7.30 [7.30, 7.30] | 3.77 [3.77, 3.78] | 1.93x | | 8 | 8192/256 | 2 | 7.51 [7.50, 7.53] | 3.79 [3.79, 3.80] | 1.98x | | 8 | 8192/256 | 4 | 7.90 [7.89, 7.90] | 4.25 [4.24, 4.25] | 1.86x | | 8 | 8192/256 | 8 | 8.43 [8.42, 8.62] | 5.32 [5.31, 5.51] | 1.58x | | 8 | 512/128 | 1 | 7.21 [7.20, 7.21] | 3.52 [3.52, 3.53] | 2.05x | | 8 | 512/128 | 2 | 7.41 [7.41, 7.42] | 3.55 [3.54, 3.56] | 2.09x | | 8 | 512/128 | 4 | 7.85 [7.84, 7.86] | 4.04 [4.03, 4.04] | 1.95x | | 8 | 512/128 | 8 | 8.40 [8.39, 8.41] | 5.06 [5.05, 5.07] | 1.66x | After the largest width, both arms use SGLang's path. The control rows stay within 1%: - each TP at 16 requests: 0.999x to 1.004x; - TP2 at 8 requests: 0.998x and 1.004x. ### Cookbook shape on daf4739 The GLM-5.3 cookbook reports speed at random 8192 in / 1024 out. This table uses that shape at TP4 on daf4739. - Each arm ran on a fresh server, in two rounds. The second round swapped the arm order and the GPU set. - TPOT is the median per request over both rounds, with a bootstrap 95% CI. TTFT is the median per request. Throughput is the mean over the rounds. | Conc. | TPOT off (ms) | TPOT on (ms) | Speedup | TTFT off / on (ms) | Output tok/s per GPU off / on | | --- | --- | --- | --- | --- | --- | | 1 | 7.75 [7.74, 7.75] | 4.39 [4.39, 4.39] | 1.76x | 235 / 232 | 31.3 / 53.8 | | 2 | 8.01 [8.00, 8.03] | 4.42 [4.41, 4.43] | 1.81x | 321 / 317 | 59.6 / 104.3 | | 4 | 8.40 [8.40, 8.41] | 5.02 [5.01, 5.02] | 1.67x | 211 / 199 | 115.8 / 190.9 | | 8 | 9.22 [9.21, 9.38] | 7.02 [7.01, 7.17] | 1.31x | 1,116 / 1,112 | 176.6 / 220.3 | | 16 | 10.08 [9.81, 10.30] | 10.07 [9.80, 10.31] | 1.00x | 1,082 / 1,104 | 337.5 / 337.7 | At 16 requests, a step has more than 8 rows, so both arms use SGLang's path. Profile (TP4, one request, torch profiler, rank 0): - Monokernels on: a step runs 184 kernels, and the GPU is busy for 4.28 ms. Of this time, the 41 layer launches use 3.44 ms (84 us each). The trace shows each launch as `down`. - Monokernels off: a step runs 1,301 kernels. Startup: - In a new container, the build of all widths takes 25 to 32 s (FlyDSL and comgr compilation). When the caches hold the kernels, the build takes approximately 2 s. - Extra memory for each GPU: - TP2: 490 MiB of launch buffers and 455 MiB of weight copies. 261 MiB of the copies is the packed shared expert, which SGLang keeps unfused at TP2. - TP4: 214 MiB of launch buffers and 150 MiB of weight copies. - TP8: 184 MiB of launch buffers and 128 MiB of weight copies. Tests on MI355X in the image above, with SGLang main at 4cca7fcd3e: - unit (CPU): 75 passed. They also passed on 04fcb04e34, and on daf4739 with SGLang main 4cca7fcd3e and 2948662026. - kernels: 24 passed on 5a8b90f: - Each GLM-5.3-Flash launch kind at TP 2, 4 and 8, on random weights in SGLang's layouts. All ranks return the same output and write only their valid rows. AITER-shuffled FP8 projections give the row-major result bitwise. - The guard tests. - The DeepSeek-V4.1 kernel tests with the `common` fixes. - e2e (TP4) has four tests: GSM8K, KL for each batch, the check mode, and an injected wait fault that must stop all ranks. - Before this review: 4 passed (GSM8K 0.99, KL 0.009 to 0.018 for each batch). - On daf4739, in one run: 4 passed. GSM8K-200 is 0.99. KL is 0.018, 0.015, 0.010 and 0.027 for batches of 2, 3, 7 and 8 rows. The 779 checks are within `NUMERICS` (lowest stream cosine 0.9904). After the injected stall, the faults are only at address 0x0. The kernel comes from vllm-project/vllm#61662 (Apache-2.0). vLLM adapted it from ROCm/aiter#6173 (FlyDSL GLM monokernel). ### Review fixes (on top of daf4739) - New startup refusals: - an index cache that SGLang keeps row-major (no AITER preshuffle paged MQA); - `--enable-hisparse`; - `--expert-distribution-recorder-mode`; - SGLang's benchmark routing overrides; - `moe_layer_freq` other than 1; - a TP2 server whose decode graphs all exceed the 4-row cap. - `SGLANG_READS` declares every SGLang name that the spec and the weight readers read. - The stage loop no longer overwrites the kernel's symbol. Before, every variant was named `down`. - Cleanup: - startup errors say what they need; - the docs match the code (the README row, the copy lists, NOTICE's provenance); - `common/peer.py`'s 2 MiB unit is named `_GRANULE`. - The README has one GPU test command per model. - CODEOWNERS adds `**/*glm53_flash* @jhinpan`. - Tests on 7944894: - Setup: MI350X, TP4, SGLang main 4eb606b2ae with sgl-project/sglang#44601. - unit: 75 passed. - e2e: 4 passed. KL was 0.009, 0.014, 0.009 and 0.029. GSM8K-200 was 0.97. The lowest stream cosine was 0.986. The injected fault stopped all ranks. - e254332 adds only a startup refusal. ## 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. ## Review and Merge Process 1. The `lint` and `Check PR title` checks must pass. 2. Ask a code owner from `.github/CODEOWNERS` to review. 3. Merge with "Squash and merge": the PR title becomes the commit title on `main`. --------- Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii
added a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 18, 2026
## Motivation A review of the merged DeepSeek-V4.1 port found five server configurations that SGLang accepts and the launches would serve wrongly with no refusal. It also found a gap in the kernel cache key, a guessed context length, and a README launch that differs from the tested one. The spec contract is that any configuration the launches cannot serve as SGLang would fails at startup with its reason. This PR restores that. ## Modifications Refusals (startup reasons; nothing changes for a configuration that already started): - `--expert-distribution-recorder-mode` is refused. EPLB and `--expert-balancedness-report-mode` also turn the recorder on. - K2 routes inside the kernel and the forward never enters the recorder's per-layer context, so routed steps recorded nothing. - With recording on from startup, the normal-MoE warm-up recorded outside a layer and failed the graph warm-up. - `SGLANG_HIP_DSV41_INDEX_K_PAGE8=1` is refused. It stores the ratio-1 index-K pages in aiter's page-8 byte layout with unchanged page size and shapes. The plane checks passed, while K2i read the pages as page-64 and picked the ratio-1 layers' top-512 from the wrong keys. - `SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE=1` is refused. SGLang would quantize the SWA nope payload from bf16 while K1 quantizes from f32, so one cache would mix two quantizations. K1's KV bound is bit-exact to SGLang's default writer. - `SGLANG_SIMULATE_UNIFORM_EXPERTS` and `SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS` are refused. K2 routes with the model's router, so a benchmark run with either one would compare overridden prefill routing with real decode routing. Forward: - Image tokens in the Engram layers are masked as SGLang masks them: on every image-token row whenever `vision_n_layers > 0`. Before, the mask ran only when the batch had mm inputs, a host check that a captured graph freezes at capture (False). So a decode or verify row carrying the image token kept the Engram residual. - K2i's logits rows are sized from SGLang's `context_len` instead of `--context-length` or `max_position_embeddings`. The scorer bounds each row's stores by the row's length, not by the row stride. Shared (`common/plan.py`): - The sources digest also hashes the aiter FlyDSL helpers that the hashed sources import (`buffer_ops`, `dpp_utils` and `kernels_common` today). - FlyDSL's own key follows only the kernel's directory, so an aiter change kept every key and loaded stale binaries. - This also covers the GLM and MiniMax-M3 specs, which digest `common/` the same way. Docs: - The README's launch now matches the serving tests: `SGLANG_USE_AITER=1`, `AITER_BF16_FP8_MOE_BOUND=0` and `--enforce-shared-experts-fusion`. Without the bound variable, decode graph capture fails in AITER's MoE on gfx950. Not changed: `--enable-hisparse`. SGLang already refuses it for DeepSeek-V4.1 (`validate_deepseek_v41_features`), and V4.1-Flash has no ratio-4 layers, the only ones HiSparse remaps. ## Accuracy Tests - MI350X, TP4, SGLang main 4eb606b2ae + sgl-project/sglang#44601, the plugin's serving test `test/e2e/test_deepseek_v41_server.py`: - 9 passed on 9634f59 (the refusals and the Engram mask). - 9 passed on 1ec77c0, the final code (the last commit edits only the README). The new sources digest rebuilt every kernel there. - `test/unit`: 65 passed. `test/unit/test_source_digest.py` is new: it fails on the digest before the fix, when an imported aiter helper changes. ## Speed Tests and Profiling Per-user decode tok/s on the e2e's DSpark server (MI350X, TP4, simulated acceptance 3.51), plugin main (d5dfda0) against this branch (1ec77c0), A-B-B-A: | Requests | main | this PR | Change | | --- | --- | --- | --- | | 1 | 640.6, 640.4 | 640.9, 641.1 | +0.1% | | 4 | 401.9, 410.4 | 399.3, 402.4 | -1.3% | | 8 | 305.3, 303.2 | 295.4, 299.6 | -2.2% | | 16 | 214.5, 213.0 | 211.7, 212.8 | -0.7% | - Acceptance length was 3.50 to 3.52 in every arm. - The only change to a routed step is the Engram mask: four small elementwise launches, about 0.2% of a c8 step (11.8 ms). Main's own two runs differ by up to 2.1% (c4). So I read the differences as run-to-run noise, not measured into an effect. ## Follow-ups (not in this PR) - SGLang framework: Feature probes in `srt/monokernels/guards.py` for HiSparse and the expert-distribution recorder, so every spec refuses them by default. - Memory: the RoPE tables (2 x 512 MiB a rank) could be read from SGLang's `freqs_cis`. The FP8 shared-expert copy (338 MiB) is a duplicate when fusion is off. Per-width scratch (about 277 MiB on a 48-width server) could be one buffer. - Tests: - the replay checks reuse static inputs, so they cannot catch a missing fence; - the draft's attention and DSpark acceptance are checked nowhere; - the fault test cannot tell peer-flag stops from per-rank deadlines. - `_ensure_writable_flydsl_cache` rewrites `FLYDSL_RUNTIME_CACHE_DIR` for the whole process on read-only installs. It could be scoped to the plugin's compiles. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - The digest has a unit test. - The refusals are data that SGLang's generic guard checks; `test_contract` pins every name they read. - Only the e2e exercises the Engram mask. - [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. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
4 tasks done
kevin-mii
pushed a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 18, 2026
## Motivation On MI355X at TP8 with one request, a GLM-5.3 decode step on SGLang's path has a median TPOT of 11.2 ms. The model is the official `zai-org/GLM-5.3` (FP8 block-128). ROCm/aiter#6173 runs each GLM-5 decoder layer as one persistent FlyDSL launch. vllm-project/vllm#60831 integrates a variant of this kernel that keeps only MXFP4 experts. This PR ports the aiter kernel to the `amd` monokernel backend of SGLang (RFC sgl-project/sglang#43949). It keeps the FP8 block-128 expert path that the official checkpoint needs. With `--monokernel-backend amd`, each MoE decoder layer of a covered decode step runs as one launch. Result on MI355X: - At TP8 with one request, the median TPOT decreases from 11.24 ms to 5.60 ms (2.01x). - From 1 to 8 requests at TP8, the speedup is 1.38x to 2.18x. From 1 to 4 requests at TP4, it is 1.45x to 1.82x. - Accuracy does not decrease. On AIME26 (30 problems x 8), pass@1 is 82.92% (off) and 86.67% (on). On GSM8K, it is 97.35% (off) and 97.27% (on). Scope: - GPU: MI355X (gfx950), at TP 4 or 8. - Steps: decode steps of 1 to 8 rows (1 to 4 at TP4). All other steps use SGLang's path. - Layers: the 75 MoE decoder layers 3 to 77. 18 launches run the indexer, and 57 reuse the previous top-k. The dense layers 0 to 2 use SGLang's path. - SGLang: this PR needs sgl-project/sglang#44353 (merged as ec7ab9d0c2). It changes no SGLang file. - Stack: this PR uses the three `common` fixes of #6, which has merged. Its diff now shows only GLM-5.3. ## Modifications ### How a decode step runs ```mermaid flowchart TB F["SGLang monokernel framework<br/>(sglang PR 44353)"] --> Q{"Decode step of 1 to 8 rows?<br/>(1 to 4 at TP4)"} Q -- "no" --> N["SGLang's path"] Q -- "yes" --> R["glm53 spec forward"] R --> A["SGLang: embedding and<br/>dense layers 0 to 2"] A -- "top-k of layer 2" --> L["75 launches, layers 3 to 77<br/>(18 run the indexer,<br/>57 reuse the previous top-k)"] L --> Z["SGLang: final norm and logits"] L -.->|reads and writes in place| M[("SGLang: MLA cache, paged FP8 index cache,<br/>page table, FP8 weights")] classDef sgl fill:#eef1f4,stroke:#8c959f,color:#1f2328 classDef amd fill:#dafbe1,stroke:#1a7f37,color:#1f2328 class F,Q,N,A,Z,M sgl class R,L amd ``` Grey boxes are SGLang code, and this PR changes none of it. Green boxes are new in this PR. One launch does these operations for one layer, in this sequence: 1. The input RMSNorm, the q_a and kv_a projections, the q_a RMSNorm, and q_b with RoPE. 2. The KV RMSNorm, the k_pe RoPE and the MLA cache write. 3. The indexer, on an indexer layer only: - the index K, Q and weight projections; - the index K norm, RoPE, Hadamard transform and cache write; - the index scores and the exact top-2048. A layer without an indexer reads the top-k of the layer before it. 4. The absorbed query (W_UK), the sparse MLA, W_UV and `o_proj`. 5. The attention TP all-reduce and the residual add. 6. The post-attention RMSNorm, the router in f32, and the FP8 quantization of the activations. 7. The top-8 selection and the experts: 1 shared and 8 routed, in FP8 block-128. 8. The MoE TP all-reduce and the residual add. ### Changes to SGLang (`python/sglang/srt`) - None. This PR adds and changes files in sgl-amd-mono only. - Without `--monokernel-backend amd`, SGLang runs as before. The spec registers for gfx950 only, so CUDA and other ROCm GPUs do not change. - With the flag, SGLang calls the spec's forward, not `DeepseekV2ForCausalLM.forward`, for each decode step that the spec covers. SGLang still does the scheduling, the cache management, the sampling and all other steps. - The plugin reads SGLang internals by name. Each module lists these names in `SGLANG_READS`. `test/unit/test_contract.py` checks them on CPU against the installed SGLang. - At startup, the plugin uses more GPU memory. Thus the KV pool is 0.7% smaller at TP8 and 3.4% smaller at TP4 (see "Startup"). ### Files - `sgl_amd_mono/glm53/gfx950/`: the decoder-layer kernel from ROCm/aiter#6173 at 56d6390def. The first GLM-5.3 commit copies the files verbatim. The next commits adapt them: - Every wait runs under `common/guard.py` (deadline, peer flags, fail-stop). - The kernel reads SGLang's FP8 weights in place, in AITER's (16, 16) layout. These are the attention projections, the indexer projections and the experts. They include the two gate/up paths that load expert weights outside the shared unit. - The kernel reads SGLang's bf16 MLA cache, paged FP8 index cache and 64-token page table in place. - The page table and its stride are launch arguments, so each captured graph reads its own table. - Index keys and queries get the 128-point Hadamard transform after RoPE. SGLang's gfx950 fused indexer caches rotated keys, so the cache keeps one format. The spec refuses all other indexers, because they cache unrotated keys. - `sgl_amd_mono/glm53/spec.py`: the spec and its forward. The forward is `DeepseekV2ForCausalLM.forward` for a covered decode step. - Layer 3 reuses the top-k of layer 2. The forward copies SGLang's indices into the compact buffer that the launch reads. - A step with fewer rows than a kernel width (2, 4 or 8) runs at the next width. One row runs at width 2, because a one-row launch does not complete on gfx950. The added rows are pad rows, and pad rows write no cache. - At TP4, a step of more than 4 rows uses SGLang's path. At TP4, an 8-row launch is slower than SGLang's path: TPOT 18.15 ms against 16.95 ms at 8192 / 256 and c8. - The spec refuses unsupported setups at startup and gives the reasons. Some examples: - another checkpoint format, page size, KV dtype or MoE runner; - an unfused shared expert; - FP8 projections outside AITER's (16, 16) layout, for example with `--fp8-gemm-backend triton`; - experts that do not use the AITER MoE runner; - an indexer other than SGLang's gfx950 fused indexer. - Weight copies: - W_UK and W_UV, packed from SGLang's per-tensor FP8 `w_kc` and `w_vc` (the same values that SGLang multiplies by); - the router and the indexer's `weights_proj`, packed for the MFMA tiles; - the indexer's k LayerNorm in f32; - one bf16 rotary table. - `sgl_amd_mono/glm53/check.py`: with `SGLANG_AMD_MONO_CHECK=1` (eager mode only), each covered layer also runs SGLang's layer on the same input and caches. The check compares the stream, the MLA KV row and the index keys (after dequantization) with the bounds in `NUMERICS`. - Registration in `sgl_amd_mono/__init__.py`, a README row, a NOTICE entry, and these tests: - `test/unit/test_glm53_spec.py` (CPU) - `test/e2e/test_glm53_server.py` (TP8) - There is no random-weight launch test like the one for glm53_flash. The e2e check compares each covered layer with SGLang's layer on the real weights. The cost is that the test needs the checkpoint. ### Routing numerics - The kernel routes in f32, as the model defines it: sigmoid plus correction bias for the choice, and sigmoid for the weights. - SGLang's AITER top-k casts the correction bias of GLM-5.3 (approximately 34) to bf16. Thus SGLang's choice and weights are different from f32 routing. - For this reason, the check mode routes SGLang's reference layer in f32. ### Not in this PR - An FP8 KV cache. - Speculative decoding (MTP). - MXFP4 checkpoints. - Decode steps of more than 8 rows (4 at TP4). ## Accuracy Tests Setup (also for the speed tests): - Hardware: one node with 8 x MI355X (gfx950). - Image: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014` (`sha256:e4498f5d595f`), with ROCm 10.0.0, torch 2.11.0, FlyDSL 0.3.4.1 and AITER e7d2453f. - SGLang: main at 4cca7fcd3e. - sgl-amd-mono: this PR. The TP8 runs are older than two commits, and neither commit changes what TP8 launches: - cbc8f15 adds the TP4 row cap. - 02d42fa adds the review fixes: startup refusals, the rotary check over each indexer, msgspec config containers, comments and notices. - The e2e tests ran again on 02d42fa (see the end). - Model: `zai-org/GLM-5.3` at ba32f8c19c (FP8 block-128). - Server: the cookbook's AMD recipe (TP8, TileLang DSA prefill and decode, `--chunked-prefill-size 131072`, no MTP), with `--mem-fraction-static 0.80`, `--max-running-requests 96` and `--cuda-graph-max-bs-decode 96`. SGLang selects the bf16 KV cache, 64-token pages and the AITER MoE runner. TP4 is not a cookbook cell, and it uses the same flags. - The only difference between the arms is `--monokernel-backend amd`. Each arm runs on a fresh server. - Evals: sgl-eval 0.1.2 at TP8, with 8 or fewer requests at a time. | Eval | Settings | Monokernels off | Monokernels on | | --- | --- | --- | --- | | GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 97.35% | 97.27% (10 vs 11 discordant, exact McNemar p = 1.00) | | AIME25 (30 x 4 seeds) | thinking, T 0.6, top-p 0.95, 65,536 tokens | 90.00% | 89.17% (-0.83 points, bootstrap 95% CI [-5.8, +4.2] over problems; 2 better, 4 worse; sign test p = 0.69) | | RULER2 NIAH at 1M tokens (4 subtasks x 2) | no thinking, greedy, 256 tokens | 31.56% | 25.60% | - AIME25: 11 answers (off) and 13 answers (on) reached the token limit. The median length was 7,630 tokens (off) and 9,351 tokens (on). For each problem, the monokernels ran longer on 17 problems and shorter on 12 (sign test p = 0.46). - NIAH: `mk_niah_basic` scores 100 in both arms. Most other answers reached the 256-token limit (4 of 8 samples off, 6 of 8 on). The reasoning lists the needles, and the limit stops the answer. Thus the multi-value subtasks mostly measure where the answer stops. - The arms also use different routing precision: f32 (on) and SGLang's bf16 correction bias (off). See "Routing numerics". Per layer, against SGLang's own layer on the same input and caches: - Setup: `SGLANG_AMD_MONO_CHECK=1`, eager mode. Batches of 1, 3 and 8 rows. Positions up to 2.4k. 1,425 checks for each TP size. - Result: no value is outside its bound. - The TP4 run is older than the TP4 row cap, so it also covers 8-row launches. | Output | TP4 | TP8 | | --- | --- | --- | | written stream, cosine | >= 0.9916 | >= 0.9938 | | MLA KV row, cosine | >= 0.9996 | >= 0.9995 | | index keys after dequantization, cosine | >= 0.99993 | >= 0.99994 | Decode against prefill (`sglang.test.kl_test_utils`, 512 tokens, mean KL for each batch): - 0.009, 0.017, 0.016 and 0.030 for batches of 2, 3, 7 and 8 rows. - The bound is 0.06. ### AIME26 on 02d42fa GSM8K is near its ceiling for this model, so it finds only large errors. The SGLang contribution guide asks for a harder eval, and the GLM-5.3 cookbook reports AIME25 and AIME26. Thus this PR also reports AIME26 on its final head. - Setup: MI355X, TP8, SGLang main 4cca7fcd3e, sgl-amd-mono 02d42fa, sgl-eval 0.1.3 (the version that SGLang pins). Each arm used a fresh server. - 8 or fewer requests ran at a time, so each decode step of the "on" arm ran the launches. - Generation: thinking, T 1.0 and top-p 0.95 (the recommended settings of the model), 65,536 max tokens. | AIME26 (30 problems x 8) | Monokernels off | Monokernels on | | --- | --- | --- | | pass@1 | 82.92% | 86.67% | | majority@8 | 96.67% | 96.67% | | samples at the token limit | 34 of 240 | 29 of 240 | | median output tokens | 12,197 | 11,492 | - Paired over the 30 problems, the "on" arm is 3.75 points higher (bootstrap 95% CI [-1.2, +9.6]). 7 problems are better, 5 are worse and 18 are tied (sign test p = 0.77). The difference is within noise. - 12% to 14% of the samples reach the 65,536-token limit. Thus these scores are lower than scores with a larger limit, and you cannot compare them with the cookbook numbers. Both arms use the same limit. - The arms also use different routing precision (see "Routing numerics"). ## Speed Tests and Profiling Method: - The setup is the same as for the accuracy tests. - Workload: `sglang.bench_serving` with the random dataset at fixed lengths (8192 in and 256 out, 512 in and 128 out). Each concurrency c uses max(16, 4c) prompts after a warmup. - Each arm runs on a fresh server, in two rounds. The second round swaps the arm order. TP4 also swaps the GPU set. - At load, each server checks that no other process holds memory on its GPUs. - The tables give the median TPOT per request over both rounds, with a bootstrap 95% CI. TP8 (GPUs 0-7): | Shape | c | Monokernels off (ms) | Monokernels on (ms) | Speedup | | --- | --- | --- | --- | --- | | 8192 / 256 | 1 | 11.24 [11.23, 11.24] | 5.60 [5.60, 5.60] | 2.01x | | 8192 / 256 | 2 | 12.36 [12.35, 12.37] | 5.78 [5.77, 5.78] | 2.14x | | 8192 / 256 | 4 | 13.30 [13.29, 13.31] | 7.33 [7.32, 7.34] | 1.81x | | 8192 / 256 | 8 | 15.85 [15.81, 17.80] | 11.51 [11.49, 13.47] | 1.38x | | 8192 / 256 | 16 | 18.65 [18.65, 22.51] | 18.94 [18.91, 22.53] | 0.98x | | 512 / 128 | 1 | 10.79 [10.79, 10.79] | 5.43 [5.43, 5.44] | 1.99x | | 512 / 128 | 2 | 11.95 [11.93, 11.97] | 5.48 [5.47, 5.48] | 2.18x | | 512 / 128 | 4 | 12.81 [12.81, 12.83] | 6.86 [6.85, 6.86] | 1.87x | | 512 / 128 | 8 | 15.41 [15.40, 15.41] | 10.93 [10.91, 10.95] | 1.41x | | 512 / 128 | 16 | 17.71 [17.68, 17.73] | 17.68 [17.66, 17.70] | 1.00x | - At c16, a step has more than 8 rows, so both arms use SGLang's path. The 16-row decode steps take the same time in both arms (median ITL 18.6 ms at 8192 / 256). - The 8192 / 256 c16 TPOT is different because of the scheduler. It admitted one burst of prefix-cached prompts as 1 + 10 + 5 requests, not as 1 + 15. This stops approximately 10 running requests once, for approximately 0.1 s. TP4 (round 1 on GPUs 0-3, round 2 on GPUs 4-7): | Shape | c | Monokernels off (ms) | Monokernels on (ms) | Speedup | | --- | --- | --- | --- | --- | | 8192 / 256 | 1 | 12.02 [11.99, 12.04] | 7.09 [7.09, 7.10] | 1.70x | | 8192 / 256 | 2 | 13.05 [13.04, 13.06] | 7.25 [7.24, 7.26] | 1.80x | | 8192 / 256 | 4 | 14.68 [14.65, 14.70] | 10.10 [10.08, 10.11] | 1.45x | | 8192 / 256 | 8 | 16.96 [16.93, 19.52] | 16.95 [16.93, 19.51] | 1.00x | | 8192 / 256 | 16 | 23.19 [23.17, 27.96] | 23.17 [23.06, 27.95] | 1.00x | | 512 / 128 | 1 | 11.61 [11.58, 11.64] | 6.95 [6.95, 6.95] | 1.67x | | 512 / 128 | 2 | 12.63 [12.61, 12.65] | 6.94 [6.94, 6.95] | 1.82x | | 512 / 128 | 4 | 14.33 [14.31, 14.34] | 9.70 [9.68, 9.71] | 1.48x | | 512 / 128 | 8 | 16.74 [16.71, 16.75] | 16.73 [16.68, 16.75] | 1.00x | | 512 / 128 | 16 | 22.26 [22.20, 22.29] | 22.22 [22.18, 22.23] | 1.00x | - At c8 and c16, a TP4 step has more than 4 rows, so both arms use SGLang's path. - Before the TP4 row cap, the 8-row launches made c8 slower: - 8192 / 256: 18.15 ms against 16.95 ms (0.93x). - 512 / 128: 17.44 ms against 16.76 ms (0.96x). ### Cookbook shape on 02d42fa The GLM-5.3 cookbook reports speed at random 8192 in / 1024 out. This table uses that shape at TP8 on 02d42fa. - Each arm ran on a fresh server, in two rounds. The second round swapped the arm order. - TPOT is the median per request over both rounds, with a bootstrap 95% CI. TTFT is the median per request. Throughput is the mean over the rounds. | Conc. | TPOT off (ms) | TPOT on (ms) | Speedup | TTFT off / on (ms) | Output tok/s per GPU off / on | | --- | --- | --- | --- | --- | --- | | 1 | 11.24 [11.24, 11.24] | 5.58 [5.58, 5.59] | 2.01x | 541 / 540 | 10.6 / 20.2 | | 2 | 12.34 [12.33, 12.35] | 5.75 [5.75, 5.76] | 2.15x | 287 / 285 | 19.7 / 40.9 | | 4 | 13.27 [13.25, 13.28] | 7.28 [7.28, 7.29] | 1.82x | 365 / 353 | 31.3 / 49.7 | | 8 | 15.80 [15.79, 16.26] | 11.44 [11.42, 11.89] | 1.38x | 3,067 / 3,059 | 49.1 / 62.5 | | 16 | 18.62 [18.60, 19.51] | 18.60 [18.60, 19.51] | 1.00x | 1,080 / 1,037 | 86.7 / 87.0 | At 16 requests, a step has more than 8 rows, so both arms use SGLang's path. Profile (torch profiler, rank 0, TP8, one request, mean of 10 decode steps): - SGLang's path: a decode step runs 1,821 GPU kernels in 11.70 ms of GPU time. - Monokernels on: a decode step runs 212 kernels in 5.59 ms. The 75 launches use 4.76 ms (63.5 us each). The other 137 kernels (0.83 ms) are the embedding, the dense layers 0 to 2, the final norm and the logits. Startup: - In a new container, the kernels build in approximately 40 s at TP8 (6 kernels) and 30 s at TP4 (4 kernels). - A later server in the same container builds them in approximately 2 s, because the container keeps the compiled code objects (`~/.cache/comgr`). - During the build, PyTorch memory increases by 511.8 MiB for each GPU at TP8 and by 573.1 MiB at TP4. - The weight copies are 491 MiB at TP8 and 622 MiB at TP4: - W_UK / W_UV: 131 or 263 MiB; - the router: 225 MiB; - the rotary tables: 128 MiB; - `weights_proj`: 7 MiB. - The peer buffers are outside PyTorch. - In total, the KV pool decreases by 0.7% at TP8 (from 1,533,824 to 1,522,624 tokens). At TP4, it decreases by 3.4% (from 539,584 to 521,024 tokens). Tests on MI355X in the image above, with SGLang main at 4cca7fcd3e: - unit (CPU): 88 passed on 02d42fa, with SGLang main 4cca7fcd3e and 2948662026. Two of the cases fail on cbc8f15: the layout and indexer refusals, and the rotary check over each indexer. - e2e (TP8): 4 passed on 02d42fa in one run: - GSM8K-200 through the launches: 0.995. - KL: 0.036, 0.025, 0.018 and 0.021 for batches of 2, 3, 7 and 8 rows. - The check mode: all values are within `NUMERICS`. - An injected wait fault stops all ranks, and the faults are only at address 0x0. - The lowest stream cosine of the check mode changes between runs, on one row. It was 0.9938 and 0.9939 in the earlier runs, and 0.9895 (layer 31, position 30) on 02d42fa. All other checks are 0.994 or more. The kernel comes from ROCm/aiter#6173 at 56d6390def (Apache-2.0, FlyDSL Project Contributors). The file headers are kept. vllm-project/vllm#60831 integrates the MXFP4 variant. ### Review fixes (on top of 02d42fa) - main is merged in (#6 with its review fixes, and #10). - New startup refusals: - unfused DSA top-k (`SGLANG_DSA_FUSE_TOPK=0`). With it, layer 2 hands layer 3 request-local positions instead of KV slots, and layers 3 to 5 would attend other requests' rows. - `--enable-hisparse`, whose logical KV slots the kernels write raw. - `--expert-distribution-recorder-mode`. - SGLang's benchmark routing overrides. - a scaled RoPE such as YaRN. SGLang scales the softmax by mscale squared, while the kernel's scale is fixed. - an f32 router gate that bf16 cannot hold exactly. - a served context longer than the copied rotary table. - a TP4 server whose decode graphs all exceed the 4-row cap. - The kernel cache key holds the sources digest, as GLM-5.3-Flash's and DeepSeek-V4.1's do. So an edit to `common/`, or to an aiter helper the kernel imports, rebuilds it. - `SGLANG_READS` declares every SGLang name that the spec reads. - Cleanup: - the lineage comments are gone; - the stale kernel and runner docstrings match the code; - the guard flags size uses `common.guard.FLAG_BYTES`; - a few names are clearer, and `gfx950/` has a package `__init__`. - CODEOWNERS adds `**/*glm53* @jhinpan`, above the Flash line, because the last match wins. - The README has the GLM-5.3 row and its own test command. - Tests on a470273: - `test/unit`: 89 passed, with SGLang main 4eb606b2ae and sgl-project/sglang#44601. - pre-commit passes. - Not executed: the TP8 e2e. It needs 8 x MI355X and the zai-org/GLM-5.3 checkpoint. ## 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. ## Review and Merge Process 1. The `lint` and `Check PR title` checks must pass. 2. Ask a code owner from `.github/CODEOWNERS` to review. 3. Merge with "Squash and merge": the PR title becomes the commit title on `main`. --------- Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
4 tasks done
kevin-mii
pushed a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 18, 2026
## Motivation Adds MiniMax-M3 to the AMD monokernel backend: each sparse MoE layer of a decode or EAGLE3 verify step runs as one launch. ## Modifications - `sgl_amd_mono/minimax_m3`: the spec and its forward. The embedding, the 3 dense layers and the final norm run SGLang's modules; each of the 57 sparse MoE layers is one launch. TP4, MXFP4 experts, 1 to 16 rows a step (EAGLE3: up to 4 requests of 4 tokens). - One launch runs K1 (input norm, fused qkv / index projection, q / k norms, RoPE, K / V and index-K writes, index scores) and then block selection, sparse attention, o_proj, the router, the MXFP4 MoE and both all-reduces. Every wait goes through `common/guard`. - Weight copies made at startup and logged: per-channel FP8 qkv / index and o_proj (the checkpoint keeps them in bf16), the bf16 router gate, the fp32 routing bias and one bf16 RoPE table, 1.7 GiB a rank. - `SGLANG_AMD_MONO_CHECK=1` compares every launch with SGLang's layer on the same input in eager steps; the spec's bounds come from it. - Tests: `test/unit` (spec, refusals on meta-device layers), `test/kernels` (row tables, every width compiles), `test/e2e` (EAGLE3 verify server, HiCache decode server, check mode, fault injection). README. Needs SGLang's page16 SHUFFLE K / V, sgl-project/sglang#41998 on #41708. Until both merge, `--monokernel-backend amd` refuses MiniMax-M3 and names them; their three switches are read with a fallback and move into `SGLANG_READS` after the merge. SGLang's `sgl-amd-mono` pin needs a bump for M3 to ship. ## Accuracy Tests MI355X, TP4, SGLang main a038a0dca0 + sgl-project/sglang#41998 (544bfad891). `amd/MiniMax-M3-MXFP4`, full 5-shot GSM8K (1,319 questions), chat completions, 16384-token budget, temperature 0, real EAGLE3 verification, 4 concurrent requests so every verify step runs the monokernels: | | Run 1 | Run 2 | |---|---:|---:| | Monokernels off | 96.36% | 96.89% | | Monokernels on | 96.51% | 96.82% | - `test/e2e`: 9 passed. - `test/unit`: the MiniMax-M3 tests and contract items pass on SGLang main and on main + #41998. On SGLang main alone, `--monokernel-backend amd` refuses MiniMax-M3 at startup, naming #41998 and #41708. ## Speed Tests and Profiling AgentX on MI355X, TP4, EAGLE3, 3600 seconds per run. | Concurrency | p90 interactivity (tok/s/user), off → on | Total throughput (tok/s/GPU), off → on | |---|---:|---:| | C1 | 289.2 → 409.8 (+41.7%) | 4,802 → 5,771 (+20.2%) | | C4 | 246.4 → 330.3 (+34.1%) | 7,711 → 8,601 (+11.5%) | | C38, HiCache | 50.3 → 54.4 (+8.0%) | 47,040 → 47,720 (+1.4%) | ### Review fixes (on top of 374860b) - main is merged in (GLM-5.3-Flash #6, the DeepSeek-V4.1 fixes #10 and GLM-5.3 #7). The README, the registry and CODEOWNERS keep every model. - The kernel cache key holds the sources digest, as the other models' do. So an edit to `common/`, to `minimax_m3/gfx950` (including `layout.py`'s dict-valued stage bases) or to an imported aiter helper rebuilds the kernels. - New startup refusals: - experts whose gate and up rows SGLang interleaves (a SiTU activation, A16W4). They are marked shuffled and have the separated layout's sizes, while the kernels read separated rows. - `--enable-hisparse`; - `--expert-distribution-recorder-mode`; - SGLang's benchmark routing overrides. - `spec.py`'s `SGLANG_READS` lists what `spec.py` reads: the parallel sizes, the shape key fields, the `"target_verify"` literal, the layer modules, `forward_native`, `out_cache_loc` and `text_config`. It no longer lists names that only `weights.py` reads. - Code that nothing called is removed: `post_attn.py`'s `_f32`, `index_score.py`'s `q_frag_from`, and two runner fields. - Cleanup: - two refusals say what they need; - the docs match the code: the MoE is MXFP8 x MXFP4 (a8w4), and the attention projections run in FP8 (now also in the numerics note and the README); - CODEOWNERS adds `**/*minimax_m3* @zijiecode`. - Tests: - Setup: MI350X, TP4, SGLang main + sgl-project/sglang#44601 + #41998 (which carries #41708). - `test/unit`: 119 passed with and without #41998, on main with every model. - `test/e2e`: the decode server's 3 tests passed. On 2c6213a, the final M3 code (8b37f17 adds only main's other models), mean KL was 0.081 (bound 0.12) and GSM8K-200 was 0.885 (bound 0.87), and the tested batches ran the monokernels. Before the fixes (f6ec322) they were 0.074 and 0.895, 2 questions apart. - Not executed here: the EAGLE3 verify, check-mode and fault tests. Their EAGLE3 draft's rotary fails on this host's sgl_kernel, which lacks main's f32 cos/sin cache, before the monokernels run. - Not in this PR: the unreachable context-parallel index path (about 220 lines across five kernel files) and the runner helpers duplicated from DeepSeek-V4.1. ## 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. --------- Co-authored-by: Zijie Chen <300606707+zijiecode@users.noreply.github.com> Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This was referenced Oct 18, 2026
Closed
kevin-mii
added a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 19, 2026
## Motivation FlyDSL's disk cache keys a launcher on a narrow set of inputs: - its own source; - plain helper functions in the launcher file's own directory; - the scalar values it closes over (dicts, Structs and dataclasses are dropped); - the FlyDSL install. It misses `common/`, `module.func` calls (even within the same directory), functions defined inside a kernel body, helper classes and the builder code. So each kernel closes over a sources digest as `keyed`, and that digest had three problems: - **It was copied per model.** Each model had its own `sources.py`. Nothing checked that a new launcher actually reads the digest. A launcher that misses it keeps serving kernels built from old code. - **It was coarse.** It hashed every file under `common/` and `<model>/gfx950/`, so a `runner.py` or `weights.py` edit rebuilt every kernel of the model. It also missed modules that aiter's helpers themselves import. - **GLM-5.3 and MiniMax-M3 keyed only the digest.** Their build arguments reached the key only as closure scalars. FlyDSL has no supported way to add to one kernel's key: `@flyc.jit` takes no arguments. ROCm/FlyDSL#1206, still open, asks for an opt-in key protocol for captured traits; source dependencies would fit the same hook. The two other routes were rejected: - `EXTRA_SOURCE_DIRS` applies to the whole process. It changes every kernel's key, so aiter's bundled AOT cache misses. - Keying on the traced IR needs a trace on every start, and a trace costs more than a compile. ## Modifications `common/plan.py`: - `kernel_sources(module)` replaces `source_digest` and the four `sources.py` files. - It hashes `module` and every `sgl_amd_mono` and `aiter.ops.flydsl.kernels` module it imports, transitively. Imports inside functions count. - It hashes parsed code without comments or docstrings. FlyDSL still hashes each jit and kernel function's own source, comments included, so this only spares edits outside those functions. - Modules are found on disk under each package's directory, with no imports. Each is parsed once per process. - `key_tuple` also takes a builder's `locals()`, allows `None`, and flattens nested Structs and dataclasses. Kernel modules: - Each one takes `SOURCES = kernel_sources(__name__)` at import, so the digest is of the code that runs. A digest read at the first build would pair an edit made since import with the code already loaded. - `deepseek_v41/gfx950/layer.py` - `deepseek_v41/gfx950/indexer/k2i.py` - `glm53/gfx950/kernel.py` - `glm53_flash/gfx950/kernel.py` - `minimax_m3/gfx950/post_attn.py` - GLM-5.3 and MiniMax-M3 key their launch on every builder argument: `keyed = key_tuple(locals(), SOURCES)`. DeepSeek-V4.1 and GLM-5.3-Flash already did this through their build Structs. `test/unit/test_source_digest.py` has two tests: - Every `@flyc.jit` launcher reads a `keyed` built from its module's import-time `SOURCES` and every build parameter, directly or through a `@flyc.kernel` it calls. - The digest changes when imported code changes. That includes an aiter helper, an attribute import, `import a.b`, and a level-2 relative import inside a function. It stays the same for an unimported module or a comment. What a kernel's digest covers, on this tree: | Kernel module | Plugin modules | aiter modules | |---|---|---| | `deepseek_v41.gfx950.layer` | 22 | 4 | | `deepseek_v41.gfx950.indexer.k2i` | 11 | 4 | | `glm53.gfx950.kernel` | 11 | 4 | | `glm53_flash.gfx950.kernel` | 11 | 4 | | `minimax_m3.gfx950.post_attn` | 15 | 6 (adds `act` and its import `tensor_shim`) | Effects: - An attention edit no longer rebuilds K2i, and an indexer edit no longer rebuilds the layer kernels. - `runner.py`, `weights.py`, `packing.py` and `peer.py` are in no digest. The build arguments they choose are keyed directly. - Every key changes once, so the first start after this PR rebuilds every kernel. ## Accuracy Tests MI350X, TP4, the plugin's serving tests on the final commit 397b87f. SGLang main da414adda7, except MiniMax-M3, which ran on main 6c3db0f520 + sgl-project/sglang#44601 + sgl-project/sglang#41998: - DeepSeek-V4.1 (`test/e2e/test_deepseek_v41_server.py`): 9 passed. - GLM-5.3-Flash (`test/e2e/test_glm53_flash_server.py`): 4 passed. - MiniMax-M3 decode (`test/e2e/test_minimax_m3_server.py -k TestMonokernelDecode`): 3 passed. - GLM-5.3 needs 8 GPUs to serve. Instead, its launch was compiled on one GPU with `runner.py`'s arguments, for a dense and an indexer layer. Cold compiles took 6.0 s and 9.0 s; a second build loaded from the disk cache in 0.2 s. - `test/unit`: 120 passed. Mutation checks; each one fails a unit test: - GLM-5.3 keyed on the digest only; - the digest taken in a builder instead of at import; - K2i dropping `tokens` from its key; - a launcher dropping `_ = keyed`; - walking only top-level imports; - hashing raw bytes; - leaving aiter out; - not following a `from m import x` module. ## Speed Tests and Profiling No kernel code changes. Only the cache key changes. The digests take milliseconds at import: each shared module is parsed once per process. ## Follow-ups (not in this PR) - Upstream: an opt-in key on `@flyc.jit`, the protocol ROCm/FlyDSL#1206 asks for, extended to source dependencies. Once it ships, the `_ = keyed` lines and the guard test can go. - The four runners each carry a copy of `ensure_writable_flydsl_cache`. On a read-only aiter install it repoints `FLYDSL_RUNTIME_CACHE_DIR` for the whole process, so aiter kernels first called after it miss aiter's bundled cache. It could be one helper in `common/`, scoped to the plugin's compiles. ## 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. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Kevin Mi <kevin.mi@radixark.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
4 of 5 tasks
This was referenced Oct 19, 2026
This was referenced Oct 28, 2026
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
A monokernel forward may never read some of the attention metadata a backend builds every step. sgl-amd-mono's DeepSeek-V4.1 forward scores every index source with K2i, so on each step it runs, the HIP radix backend builds the aiter decode schedule of each low ratio and nothing reads it. The build is recorded in the decode CUDA graph, so it runs on every replay.
Co-authored with @kkHuang-amd: this is kevin-mii#19 with the plugin contract stated as data.
Modifications
monokernels/spec.py: aStepMetadataenum andMonoSpec.unread_metadata, the per-step metadata a spec's forward never reads. It defaults to empty, so existing plugins are unchanged and the API version stays 1.monokernels/runtime.py:is_metadata_unread(model, forward_batch, metadata), true only on a step the forward switch routes, for metadata the spec lists.deepseek_v4_backend_hip_radix.py: the low-ratio decode schedule build skips a ratio whenis_metadata_unread, next to the identity-rows filter. Without--monokernel-backend, every schedule is built as before.CPU test
test_only_routed_steps_skip_the_metadata_a_spec_never_reads.pyproject_other.toml: the plugin pin moves into onemonokernels_hipextra thatsrt_hipandsrt_hip_rocm724both include, at sgl-amd-mono e07665c ([Perf] deepseek_v41: Declare the aiter decode schedule a routed step never reads sgl-amd-mono#5), whose DeepSeek-V4.1 spec declares the schedule unread. A bump is one line, and the two extras cannot disagree. Typo: rename image_url to image_file #5 merged first; until this PR merges, plugin main's DeepSeek-V4.1 spec needs it (it importsStepMetadata).docker/rocm.Dockerfile: the plugin install step from [Feat] [AMD] docker: Install the sgl-amd-mono monokernel plugin in the ROCm image #44621 checks out the commit pip installed throughmonokernels_hip, read from the plugin'sdirect_url.json, instead of its own pin;SGL_AMD_MONO_COMMITstays as an optional override, and the step skips SGLang versions that do not install the plugin. The image and pip then install the same commit. The step's text, run with git andpip installstubbed but the realpip showand metadata read, covers the installed, not-installed, override, gfx942 and opt-out cases; the image itself was not built.Accuracy Tests
MI350X, TP4, sgl-amd-mono's
test/e2e(9 tests) on DSpark (verify steps) and decode servers. KL is against the standard path's prefill (limits 0.03, long prompts 0.06):test/registered/unit/monokernels: 18 passed. The measurements ran at 9fd4430; 212f36d after it changes only a docstring, a comment and a test variable name, and a3a83ff pins sgl-amd-mono e07665c, the same plugin code as ae2f020. 31e3031 merges main (with #44405 and #44630), and 958d06d pins the plugin once and has the image check out what pip installed; 29 CPU tests (test/registered/unit/monokernelsand the CUDA graph config tests) pass on it. The e2e on b3e0ba6 with sgl-amd-mono e07665c: 9 passed (DSpark GSM8K 0.910, KL 0.0123, long-prompt KL 0.0319; decode GSM8K 0.910, KL 0.0103).Speed Tests and Profiling
Decode time per output token (ms), MI350X, TP4, DSpark with simulated acceptance 3.51, median of 3 rounds of 1024 tokens, one server per arm. Main is ec7ab9d with sgl-amd-mono 3970534; kevin-mii#19 is ec43db1bff with 3970534; this PR is 9fd4430 with sgl-amd-mono#5 (ae2f020):
The two skip runs take the same path and differ by up to 2% at 4 requests, which is the noise. At 1 request both are faster: 0.4-0.5% at 8k and 1.8-2.9% at 128k.
Not in this PR
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ✅ Run #39108114682
Latest PR Test (Extra): ❌ Run #39108114639
Latest PR Test (AMD ROCm 10): ❌ Run #39108114675