Repository navigation
[Feat] [AMD] monokernels: one runtime switch for vendor monokernel plugins - #44353
Conversation
AMD monokernels for SGLang (RFC sgl-project/sglang#43949), moved out of sgl-project/sglang#44353. SGLang keeps the framework (sglang.experimental.monokernels); this package holds the spec, its glue and the kernels, registered under the "sglang.monokernels" entry points. - python/sgl_amd_mono/common: the HIP device runtime (guarded waits, error words, epochs, peer memory). - python/sgl_amd_mono/deepseek_v41: the spec and its glue, and the gfx950 kernels (K1 and K2 on all 40 layers, TP2 and TP4, up to 48 rows). - python/sgl_amd_mono/test: the kernel checks run under torchrun. - test/{unit,kernels,e2e}: the spec and SGLang contract on CPU, each launch against SGLang's ops, and serving with a fault on one rank. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
/rerun-test --changed test_resolution_is_reproducible.py test_compressed_tensors_wna16_moe_reload.py |
|
🚀 🚀 |
|
CI status on 0cb18ca, PR Test (Base):
Without The other red workflows don't come from this PR. On AMD, |
## 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>
## 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>
Motivation
Persistent "monokernels" run most of a decoder layer, including its TP all-reduce, in one or two launches, which removes the launch gaps and collective latency that dominate small-batch decode. RFC #43949 proposes one framework for them. This PR is that framework, with the kernels themselves living outside SGLang in vendor plugins: AMD's are in sgl-project/sgl-amd-mono (DeepSeek-V4.1-Flash and its DSpark draft on MI350X / MI355X).
Modifications
--monokernel-backend <name>inexec.kernel, an experimental feature, disabled by default; selecting a backend logs that its kernels, coverage and plugin API may change. General plugins (sglang.srt.plugins) register a backend's spec per architecture throughsglang.srt.monokernels.spec.register; specs load lazily and declare the API version they target (MONOKERNEL_API_VERSION, bumped on any breaking change, likeSERVE_BACKEND_API_VERSIONforsglang servebackends).guards.handle_monokernel) and align configured decode capture sizes with its supported widths. A spec states what it needs as data (Requirements: dtype, decode-phase attention backend, server arguments, HF config fields, SGLang environment variables), checked the same way for every spec; an optionalconfig_disable_reasonscovers the rest. Observe checkpoint tensors before model-specific conversion (runtime.monokernel_weights_iterator) so plugins can retain required weight copies, and refuse a spec that keeps copies when the load path skipped the observer.runtime.install_monokernel_forward), and install a full-forward dispatcher that preserves the model's signature. Every built width of the model's step (decode, or verify on a speculative server at any token count) uses the plugin's forward; other shapes use the original forward. Models with derived weight copies refuse online weight updates.BuiltMonokernelswithforward, the spec callbacks' signatures,Plan) and state the forward's obligations: the model forward's arguments and outputs, the KV cache and auxiliary hidden-state writes of the standard path for every row, padded rows included, and CUDA-graph capturability. Keep framework code inpython/sglang/srt/monokernels/; vendor kernels and model-specific forwards remain in their plugin packages.srt_hipextras ofpyproject_other.toml, so AMD installs (and SGLang's ROCm images) include the plugin;--monokernel-backend amdstays opt-in. Add the server-argument documentation row.Accuracy Tests
DeepSeek-V4.1-Flash, MI350X, TP4, 04fcb04 with sgl-amd-mono 8f1bff3 (on the DSpark server, target and draft both on monokernels). Each server steps past the largest covered width and back as its batch fills and drains:
A fault injected on one rank stops every rank, and no token of the failed step reaches the client. CPU tests:
test/registered/unit/monokernels17 passed; the plugin's serving tests (above) 9 passed. Not executed: MI355X, and AgentX at TP2 (in progress).Speed Tests and Profiling
Per-user decode tok/s, real DSpark acceptance (mean accept length 3.72 off, 3.69 on), monokernels off -> on: 1 request 444.1 -> 776.0 (+75%), 4 requests 362.8 -> 441.4 (+22%), 8 requests 267.4 -> 327.5 (+22%). Measured at 05e333d with sgl-amd-mono 905f61b; the later commits change startup checks and logging, not the step path.
AgentX (the published MI355X agentic recipe, run on MI350X, TP4, DSpark, 1 concurrent user for 900 s, forced acceptance 3.51): P90 410.3 -> 543.8 tok/s/user (+32.5%), P50 438.6 -> 591.5 (+34.9%), 117 -> 132 requests. Both arms ran with the simulated-acceptance fold, a separate PR that keeps sampled verify steps on the CUDA graph under forced acceptance. Measured with sgl-amd-mono 8faf602, before its K2i indexer kernel; not rerun since.
Checklist
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ❌ Run #38781471976
Latest PR Test (Extra): ❌ Run #38781471940
Latest PR Test (AMD ROCm 10): ⏳ Run #38781471831