Repository navigation
[Feat] [MI355X] glm53_flash: Add GLM-5.3-Flash decode monokernels - #6
Merged
Merged
Conversation
A buffer load carries at most four dwords, so at TP8 the single 8-wide load of every rank's error flags failed instruction selection.
A peer maps an exported uncached allocation in 2 MiB granules. With 4 KiB rounding, the guard flags past the last whole granule (2 MiB + 256 bytes into a 4-row GLM-5.3-Flash buffer) faulted on the peer that raised them.
Each buffer's handshake is a collective. Ranks that build several buffers in different orders map another buffer's memory and wait until the deadline. The handshake now exchanges each rank's size with its handle and fails on every rank when they differ.
Verbatim copies of vllm-project/vllm#61662 at ca328e297b (vllm/models/ glm5next/amd/ops/mono and mono/weights.py), themselves adapted from ROCm/aiter#6173 at c39b56c36 (Apache-2.0). The next commit adapts them to SGLang and this package.
With --monokernel-backend amd, each MoE decoder layer of GLM-5.3-Flash but the last runs as one persistent launch on decode steps of 1 to 8 rows (1 to 4 at TP2), on TP 2, 4 or 8: the attention-side mHC, the attention (KDA, or sparse MLA with its kpool indexer and cache writes), o_proj and its TP reduce, the FFN mHC and norm, the router, the shared and routed MXFP4 experts and the MoE TP reduce. A step runs at the next kernel width (2, 4 or 8), the rest pad rows that write no cache. The spec's forward runs SGLang's embedding, dense layers, last layer, final norm and logits processor around the launches. The kernels read SGLang's state in place: the KDA conv and recurrent state pool, the DSA MLA and index caches and compress tails, and the 64-token page table (a launch argument, so each captured graph reads its own). They read the FP8 projections in the AITER (16, 16) layout SGLang's gfx950 block-FP8 linear method leaves them in, and the MXFP4 experts as SGLang's Quark path shuffles them. SGLANG_AMD_MONO_CHECK=1 (eager) runs the normal layer next to each covered layer and checks the written streams and every state or cache row it writes against the spec's NUMERICS.
- unit (CPU): the covered layers, the registry entry, the row routing and the TP2 row cap, and the config and model refusal reasons - kernels (8 GPUs): each launch kind on random weights in SGLang's layouts at TP 2, 4 and 8 finishes on every rank with the same output and writes exactly its valid rows' state or cache rows; AITER-shuffled FP8 projections equal row-major ones bitwise - e2e (TP4): each covered layer within NUMERICS against the normal layer (1, 3 and 8 rows, positions past 2048); GSM8K while requests join and leave a batch that steps past 8 rows; decode against prefill (KL) on one graph per batch of 2, 3, 7 and 8 rows; and a rank stalled by fault injection stops the server before any token of the failed step is sent
On SGLang main, importing layer_boundary reaches fp8_kernel's gfx1250 probe, so the CPU unit test of the TP2 fallback failed without a GPU. The fallback path does not use the import.
This was referenced Oct 16, 2026
4 tasks done
The launch updates the KDA state in place, while --enable-linear-replayssm defers SGLang's state updates to a ring, so the spec now requires it off. SGLANG_AMD_MONO_CHECK is read once at build instead of per layer and step. Two comments, the README design summary and NOTICE now describe the GLM-5.3-Flash code as it is.
jhinpan
marked this pull request as ready for review
October 16, 2026 23:04
Collaborator
Author
|
@kevin-mii This PR is ready for review, and #7 (GLM-5.3) is stacked on it.
|
This was referenced Oct 17, 2026
Without AITER's preshuffle paged MQA (Triton < 3.5 and no AOT gluon kernels, SGLANG_DSA_HIP_DISABLE_PRESHUFFLE=1, or SGLANG_USE_AITER=0), SGLang's k-pool writer keeps the index cache row-major, while the kernels write and read it as AITER's 16-pool tiles. A routed step then scores a misread cache; SGLang's own k-pool decode refuses that runtime, and the spec now does too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The contract test pins only declared names, so these reads (the TP group's CPU group, BuildContext's fields, forward_mode, the decoder layer's norms and modules, the shared expert's projections and the experts' w13 / w2 weights) would fail at a GPU server's startup instead of on CPU after an SGLang rename. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The README row counted the covered layers ambiguously, the copy lists missed the f32 router bias, weights.py said the caller records the copies, NOTICE covered only the kernel of the files adapted from vLLM, the e2e docstring skipped the check-mode test, and the kernel docstring still described vLLM's FFN-only launch. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The cache-binding reasons, the o_norm check and the page-table check named only what they found (or misreported a wrong shape as a missing weight), so a layout change in SGLang left the reader to find the expected layout in the code. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The constant changed from a 4 KiB page to the 2 MiB IPC granule its comment describes, so _PAGE no longer named what it holds. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
In weights.py, idx held the layer number next to w_idx and idx_mod, where idx meant the indexer; spec.py already calls the layer number i. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
spec.py declared the decoder layer's mlp, which only weights.py reads (and declares). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
SGLang accepts --enable-hisparse for GLM-5.3-Flash (a DSA model): its pool hands out logical KV slots and translates them to device rows on every write, while the kernels write and read the raw slots, past the device buffer or into other requests' rows. With --expert-distribution-recorder-mode (also set for EPLB and expert balancedness reports) the covered layers record nothing, and layer 44 records outside SGLang's per-layer context, which fails decode graph capture. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The forward hands each covered layer's owed FFN post to the next covered layer, so it needs the MoE layers back to back; a config with moe_layer_freq 2 passed the checks and gave a dense layer the owed struct at the first routed step. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The stage-base loop rebound name, so every variant (KDA or DSA, each width) was emitted as the last stage's name, "down": profiles could not tell the launches apart and IR dumps of all variants shared one directory. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ow cap At TP2 only steps of up to 4 rows run on the kernels; when every captured width was wider (say --cuda-graph-bs-decode 8 16), no runtime was built and every step ran SGLang's forward while startup reported the widths as built. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The serving tests launch their model from the hub id when its path variable is unset, so the combined command on a DeepSeek-V4.1 box also launched (and downloaded) GLM-5.3-Flash at TP4. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
SGLANG_SIMULATE_UNIFORM_EXPERTS and SGLANG_SIMULATE_ROUND_ROBIN_EXPERTS replace the router's choice in SGLang's top-k, which routed steps never call: the kernel route with the model's router, so a benchmark with either set compared prefill under the override against decode under real routing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii
pushed a commit
to jhinpan/sgl-amd-mono
that referenced
this pull request
Oct 18, 2026
main has sgl-project#6 (GLM-5.3-Flash with its review fixes) and sgl-project#10. Every GLM-5.3-Flash and common/ file takes main's version, since this branch carried sgl-project#6's pre-review copies; NOTICE keeps both provenance paragraphs, the README both model rows (the GLM-5.3 row in the Flash row's wording) and one GPU test command per model, and the registry both GLM specs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii
pushed a commit
that referenced
this pull request
Oct 18, 2026
main now has GLM-5.3-Flash (#6) and the DeepSeek-V4.1 fixes (#10). CODEOWNERS, the registry and the README keep both models; the README's launch-count bullet takes main's wording with MiniMax-M3's one-launch sentence appended, and the GPU test commands are one per model. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii
pushed a commit
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
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>
1 of 2 tasks
5 tasks done
This was referenced Oct 21, 2026
Open
Closed
4 tasks done
jhinpan
pushed a commit
that referenced
this pull request
Oct 23, 2026
## Motivation The GLM-5.3-Flash layer launch (#6) follows vLLM in two places where SGLang's own decode does something else, so on some inputs it picks other experts, or attends to other keys, than SGLang's path would on the same input: - **Router ties.** SGLang's decode under Quark with AITER (`biased_grouped_topk`, the path `biased_grouped_topk_gpu` takes while `SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK` is off, its default) rounds the router logits and the correction bias to bf16, and keeps equal keys in AITER's grouped top-k lane order. The launch kept f32 keys and broke ties by the lower expert id. - **Keys of long rows.** SGLang's kpool top-k (`fast_kpool_topk_transform_fused`) gives a row with more than 512 complete pools the selected pools' 2,048 tokens and then its tail (the 1 to 3 tokens past its last complete pool), up to 2,051 keys; any shorter row (up to 2,051 tokens) attends to every token. The launch scored rows from 2,049 tokens on (`pos >= 2048`) and kept 2,048 keys, the tail replacing the last selected pool's last tokens. This PR makes the launch select as SGLang does. It was the first commit of #22 and is split out so that the routing change lands, and can be bisected or reverted, on its own. **Merge it before #22**, which is rebased on it and already contains these commits. ## Modifications - `gfx950/kernel.py`: - Router: the GEMM's logits and the bias are rounded to bf16 before the sigmoid and the add, and equal keys go in AITER's order (`tie_rank`, the lane priority of AITER's one-wave grouped top-k). - Sparse attention: `row_keys(pos)` gives each row's complete pools, whether it is scored, and its key count. A scored row's selection is the 512 pools' tokens followed by its tail; the split attention takes a 33rd task for the tail's keys, and the merge adds it to the second half. - `gfx950/layout.py`: `MAX_KEYS`, `BASE_SPLITS`, `N_SPLIT = 33`, `SEL_STRIDE`; the selection mailbox holds `SEL_STRIDE` tokens a row. - `spec.py`: `requires.env` adds `SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK: False` next to the two `SGLANG_SIMULATE_*` overrides, and `SGLANG_READS` declares it. With it on, ROCm routes through the Triton gate, where the lower expert id wins ties, so the launch would no longer match SGLang there. - Tests: - `check_layer.py --check-routing`: router weights whose keys tie exactly in bf16 and differ in f32. The launch's top-8, in order, must equal AITER's `biased_grouped_topk` on the launch's own FFN input, every draw. The same run counts the rows where lower-id ties or f32 keys would pick differently (both must be nonzero, so the inputs exercise the change). The host-side order with the kernel's `tie_rank` is printed as a diagnostic only. - `check_layer.py --check-keys`: rows of 4,102, 5,003 and 7,001 tokens (tails of 2, 3 and 1) and one of 2,051. A scored row's selection must be 512 whole, distinct, complete pools followed by its tail tokens in order; perturbing the MLA rows of the tokens a 2,048-key selection would drop for the tail must change the row's output; perturbing an unselected pool must not. The 2,051-token row attends to every token. - `test_glm53_flash_kernels.py` runs both at TP 2, 4 and 8. - `test_glm53_flash_spec.py`: the spec declares `SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK: False` in `requires.env`. - Registers, on main's kernel: at TP4, 8 rows the KDA build spills 12 VGPRs (main 10) and the sparse-attention build 13 (main 17); TP2 is unchanged. ## Accuracy Tests All results in this section were measured on this PR alone (its head over main's kernel), on `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`, SGLang ab580b2e06, `amd/GLM-5.3-Flash-Quark-MXFP4@b5688f25`. Reruns by @jhinpan on 21309bc (this head) unless marked. - **Kernel checks at TP4 and TP2:** `--writes` (FFN, KDA, sparse attention), `--compare-shuffled`, `--check-routing`, `--check-keys` all pass. - Routing: top-8 == AITER `biased_grouped_topk` in 18/18 rank-draws; the negative controls (lower-id ties, f32 keys) differ in 12/12 rows each. - Keys: every scored row is 512 whole distinct pools followed by its tail; the tokens a 2,048-key selection drops change the output, and an unselected pool does not. - **Both checks fail on main** (bug regressions): main's top-8 differs from AITER's in every draw (TP8: 24/24 rank-draws), and its scored rows fail the whole-pool and tail checks. - **Same code as reviewed on #22:** outputs and every state and cache write are bitwise equal to `a524ebf`. - **Env gate:** a server started with `SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK=1` is refused at startup: `--monokernel-backend amd cannot run this setup: SGLANG_OPT_USE_JIT_KERNEL_GROUPED_TOPK True is not False`. - **Unit tests:** 119 passed; the 2 failures are deepseek_v41's contract cases, as on main (they need an SGLang newer than ab580b2e06 since #17). - **E2E, `test_glm53_flash_server.py` at TP4:** 4/4 pass. GSM8K past the largest width 0.985; KL decode against prefill 0.005 to 0.014; check mode 779 checks, 0 violations (minimum cosine 0.979 sparse-attention streams, 0.992 KDA streams, 0.9993 MLA KV rows, 0.999999 KDA state); the injected fault stops the server with no token returned. - **TP8:** `test_glm53_flash_kernels.py` and `test_guard_hip.py` on 8 GPUs (TP 2, 4 and 8): 8 passed. Check mode at TP8: 779 checks, 0 violations (minimum cosine 0.989 sparse-attention streams, 0.990 KDA streams, 0.9992 MLA KV rows, 0.999998 KDA state). - **Check mode at TP2** (ours, this PR's kernel over main): 779 checks, 0 violations; sparse attention past 2,048 tokens exercised (minimum cosine 0.993 sparse-attention streams, 0.992 KDA streams, 0.9991 MLA KV rows, 0.999999 KDA state). - **GSM8K** (sgl-eval 0.1.3; 1,319 questions, thinking, T 1.0, top-p 0.95, seed 42, 8 at a time, TP4): 97.04% with this PR, 96.97% off; 10 questions right only with this PR and 9 only with off (exact McNemar p = 1.0). ## Speed Tests and Profiling On main's kernel this change costs a little time. Measured by @jhinpan on 21309bc against main, TP4: - Serving TPOT, 2048 / 256, fresh servers, two rounds each: +0.6% at one request (4.37 → 4.40 ms), +0.7% at 2, +1.6% at 4 (5.03 → 5.11 ms), +1.1% at 8, and −0.2% at 16 (SGLang's path). - Kernel-only, model order, per launch: +0.7 / +0.8 / +1.5% at 2 / 4 / 8 rows with the sparse-attention rows at position 2,047 (no tail), and +1.1 / +1.9 / +1.4% at position 8,190 (a tail). - A single sparse-attention launch at position 8,190: +2.5 / +5.0 / +3.9%. - At position 2,049 the launch is 15 / 10 / 9% faster, because rows of 2,049 to 2,051 tokens are no longer scored. The cost comes from the 33rd split task for the tail's keys: it does work on 3 of every 4 positions past 2,051 tokens, and even idle it adds a task to every (row, head group) of the split stage and moves the task placement of the stages after it, which is why the step costs time without a tail too. The selection also writes the tail after the 512 pools, the merge adds the tail's partial, and the router computes tie ranks for its keys. #22 replaces this kernel, so these numbers describe main only between the two merges. ## 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 (main ab580b2e06). - [x] Accuracy and speed results are attached when kernel or forward code changes. --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
4 tasks done
jhinpan
pushed a commit
that referenced
this pull request
Oct 23, 2026
<!-- Branch fix/glm53-flash-index-check, one commit on main (992a3b2). --> ## Motivation Under `SGLANG_AMD_MONO_CHECK`, the check compares a sparse-attention layer's completed kpool index pages against SGLang's own. Today it compares the FP8 values and the per-pool power-of-two scales separately (`dsa.index.cosine` 0.999, `dsa.index_scale.rel` 0.01). The indexer reads neither the values nor the scales on their own. It reads `FP8 value × pool scale`. A pool whose amax sits at a power of two can legitimately take the next scale in the launch and the previous one in SGLang, with its values requantized. The keys are then the same to FP8 precision, but the values and the scale each differ. This happened once in #23's TP2 serving check: layer 15, a pool completed at position 31, 3 rows. The check read `dsa.index.cosine` 0.986 and `dsa.index_scale.rel` 0.10, two violations on a correct launch. On the dequantized keys the same case reads 0.999996. @jhinpan agreed in the review of #23 that comparing dequantized keys is the right invariant, and asked for it as its own PR because it changes #6's check for every sparse-attention layer. ## Modifications - `check.py`: the index-cache output is now compared on its dequantized keys (each tiled FP8 value times its pool's f32 scale) against `dsa.index.cosine`. - `spec.py`: `dsa.index_scale.rel` goes from `NUMERICS`. - `test/unit/test_glm53_flash_check.py` (CPU) covers three cases: - the tiled page decode (16-pool × 16-byte tiles); - a pool requantized at the next scale (passes); - a pool's scale doubled with its values kept (fails). This negative control stays: for pools of similar norm, one doubled scale drops the page's cosine to about 0.993, against the 0.999 bound. No launch, spec forward or weight changes. Only the check's numerics change. ## Accuracy Tests No serving numerics change (the check runs only under `SGLANG_AMD_MONO_CHECK=1`). - Check mode at TP4 on this branch: 779 checks, 0 violations; the minimum `dsa.index.cosine` (dequantized) is 0.999984. Unit tests pass, apart from the two DeepSeek-V4.1 contract cases that main's #17 needs a newer SGLang for. - The same check on #23's head: 874 checks, 0 violations; the minimum `dsa.index.cosine` is 0.999983 (TP4) and 0.999981 (TP2, including the layer-15 case). ## Speed Tests and Profiling None: check mode only. ## 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 (SGLang main ab580b2e06). - [x] Accuracy and speed results are attached when kernel or forward code changes (no kernel or forward change). Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
jhinpan
pushed a commit
that referenced
this pull request
Oct 23, 2026
## Motivation The GLM-5.3-Flash monokernels (#6) run each MoE decoder layer 3 to 43 as one persistent launch. At TP4 with one request, a decode step takes 4.40 ms (8K/256), of which the 41 launches take about 3.44 ms. This PR makes those launches about 1.6 to 1.7 times faster without changing what they cover or how SGLang drives them. Result on MI355X, 2048 in / 256 out, median TPOT, SGLang's path → main's monokernels → this PR: - TP4, one request: 7.63 → 4.34 → 2.87 ms (1.52x over main, 2.66x over SGLang's path). - TP4, 8 requests: 9.33 → 7.08 → 4.59 ms (1.54x over main). - TP2, one request: 9.22 → 5.62 → 3.87 ms (1.45x over main). - TP8, one request: 7.21 → 3.73 → 2.62 ms (1.43x over main, 2.75x over SGLang's path). - Past the largest width all arms run SGLang's path and match within 0.3%. - Kernel-only, the 41 launches in model order: 1.6 to 1.7 times faster at TP2 and TP4, 1.46 to 1.59 times at TP8. Accuracy does not decrease: GSM8K (1,319, thinking) 97.12% off, 97.19% on (11 vs 12 discordant); AIME25 (30) 83.3% off, 86.7% on. **On top of #27** (`[Fix] glm53_flash: Route ties and long rows' keys as SGLang does`, merged), which aligns the router's ties and the keys of rows past 2,048 tokens with SGLang's own decode. This PR's diff is the kernel and its tests; the measurements below include #27. Scope (unchanged from #6): - GPU: MI355X (gfx950), TP 2, 4 or 8. - Steps: decode steps of 1 to 8 rows (1 to 4 at TP2). - Layers: 3 to 43 (41 launches: 30 KDA, 11 sparse attention). - Spec, weights, `check.py`, `NUMERICS`, `SGLANG_READS`, `common/`: unchanged (beyond #27). - SGLang: no change; tested on main ab580b2e06. ## Modifications ### 1. Faster decode-layer launches (`c6ee629`, with the review's `e27297d`, `ebe1443`, `d6192f4`) The layer kernel is replaced with a faster one, on the same interfaces: `FlashBuild`, `runner.py`, the 64-token page table as a launch argument, `common/guard.py`, `common/peer.py`, `common/ops.py`, FP8 projections read in place. `kernel.py` and `layout.py` change; nothing else does. The changes, by kind (the commit message lists them): - **Placement.** Every CU gets a tile in the projection stages; tail tasks (the tail split; the beta, f_a and g_a rows) run first; a stage's tasks stay off CTAs whose earlier role ends late. - **Prefetch.** A task issues its weights, scales, states and table entries before its wait, sized to the registers the wait leaves free, and the next task's before the current one reduces. The per-step words (epoch, peer, the argument and row tables) load at launch start with the first stages' wait-free loads. - **Hand-offs.** Broadcast rows carry a ready flag polled from one lane; many-producer broadcasts are sampled from one wave; releases drain device-scope stores instead of fencing; inputs are polled in the order their producers finish. - **Fusion and reuse.** The gate publishes the layer input row before its Sinkhorn (16 lanes, DPP) and quantizes the FFN row to MXFP8 for the experts. An up/gate task covers 32 intermediate values for half the rows and publishes mid as MXFP8, so down no longer quantizes it. The shared expert runs before the router scores. A split task covers 16 heads, so each key row is gathered once. - **Registers.** Late stages recompute thread indices; the guard reads its deadline and its stop words (`guard.ERROR_WORDS`) wave-uniform. No build spills VGPRs at TP 2, 4 or 8, any width, with FP8 preshuffled (what serving uses) or row-major. (Before the review round the row-major sparse-attention build at TP2, 4 rows spilled one VGPR; removing the code warming freed it.) Main spills at TP2, 4 rows (KDA 4; sparse attention 4 preshuffled, 8 row-major) and at TP4, 8 rows (KDA 10, sparse attention 17). The guard, its error words, `SGLANG_AMD_MONO_FAULT_INJECT` (sequential sites, armed by steps), the kernel symbols (`kernel_symbol`, one per build) and the two output sets by layer parity are as on main. **Numerics.** This commit also changes numerics slightly, through the rewritten gate (Sinkhorn on 16 lanes with DPP, the seam's mix partials summed as per-stream chains: another fp32 summation order). Measured on review, one launch of each kind on the same weights and inputs, before and after this commit (TP4, 4 and 8 rows): the state and cache writes are bitwise equal; the FFN-only launch differs in post / comb by up to 4e-7 absolute; through the attention-side mixes the difference grows to about 6e-5, and 24 to 30% of the KDA launch's bf16 `x` values differ (cosine 0.999991). The per-layer checks against SGLang below are unchanged within noise. ### 2. Tests and tools (`a749c35`, `e102e1c`) `check_layer.py --replays N --layers K`: a step of K launches of one kind on one runtime (each layer its own state or caches, inputs chained), captured after an eager step; N replays from the same state give the eager outputs and writes bitwise, with clean error words and every rank the same output. `test_glm53_flash_kernels.py` runs it for KDA and sparse attention at every width each TP size builds. `test/kernels/glm53_flash/time_layers.py`: the kernel-only timing driver behind the table below (one graph of a whole step, each layer its own random weights, state and caches; model order or one kind; any width and TP size). The README's testing section gives the command. ## Accuracy Tests Setup: MI355X, TP4, `amd/GLM-5.3-Flash-Quark-MXFP4` at b5688f25, sgl-eval 0.1.3 (the version SGLang pins), the e2e test's server flags (`test/e2e/glm53_flash_server.py`). Each arm a fresh server on the same GPUs, 8 or fewer requests at a time (every decode step of the "on" arm runs the launches). | Eval | Settings | Off | On (this PR) | Paired | | --- | --- | --- | --- | --- | | GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 97.12% | 97.19% | 11 vs 12 discordant; one-sided McNemar p = 0.66 (on worse), two-sided 1.00 | | AIME25 (30 x 1) | thinking, T 0.6, top-p 0.95, 65,536 tokens | 83.3% | 86.7% | 1 vs 2 discordant; 5 / 4 samples at the token limit | | RULER2 NIAH at 128K (4 subtasks x 2) | thinking off in the request, greedy, 4,096 tokens | 96.9% | 96.9% | identical per subtask (mk basic / easy 1.0, mv basic 1.0, mv easy 0.875) | - RULER2: #6 used 256 output tokens at 1M. GLM-5.3-Flash reasons even with thinking off in the request (sgl-eval's `--no-thinking` sets `thinking`; the template reads `enable_thinking`, and the model still writes reasoning with both false), so at 256 tokens the multi-value subtasks hit the token limit in both arms (both 50.0%, identical). The table uses 4,096 tokens; 128K keeps the run short. - GSM8K: one-sided McNemar is the no-regression test (fails only when "on" is worse, p < 0.025). Per layer, against SGLang's own layer on the same input and caches (`SGLANG_AMD_MONO_CHECK=1`, eager; batches of 1, 3 and 8 rows, 4 at TP2; positions up to about 2.6K): no value outside its bound. | Output | TP2 | TP4 | TP8 | | --- | --- | --- | --- | | written streams, cosine: min / p05 / median, this PR | 0.970 / 0.994 / 0.9986 | 0.983 / 0.994 / 0.9986 | 0.991 / 0.994 / 0.9987 | | the same, main, same requests | 0.988 / 0.994 / 0.9987 | 0.992 / 0.994 / 0.9987 | 0.985 / 0.994 / 0.9987 | | MLA KV rows, cosine (min), this PR / main | 0.9991 / 0.9992 | 0.9987 / 0.9992 | 0.9991 / 0.9991 | | KDA recurrent state, cosine (min) | 0.999999 | 0.999999 | 0.999998 | | checks, violations | 779, 0 | 779, 0 | 779, 0 | The minimum is one row of one layer (TP2: layer 23, sparse attention; at TP8 main's is the lower); the 5th percentile and the median match main's at every TP size. The streams differ from SGLang's mostly through the MoE (MXFP8 expert inputs against MXFP4, as on main). ## Speed Tests and Profiling Setup: - Hardware: one node with 8 x MI355X. - Image: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014` (`sha256:e4498f5d595f`, #6's image): ROCm 10.0.0, torch 2.11.0, FlyDSL 0.3.4.1, AITER e7d2453f. - SGLang main ab580b2e06 (contains #44353 and #44601). sgl-amd-mono main 0bce188. - Server: the e2e test's flags; the arms differ only in `--monokernel-backend amd` and the plugin tree (main or this PR). - Method: `sglang.benchmark.serving`, random dataset with fixed lengths, max(16, 4c) prompts after c warmup requests; two rounds, the second in reverse arm order; the tables give the mean of the two rounds' median TPOT (ITL in parentheses). TP4, 2048 in / 256 out (8192 / 256 in the second table): median TPOT in ms, ITL in parentheses. | Conc. | Off | Main | This PR | vs main | vs off | | --- | --- | --- | --- | --- | --- | | 1 | 7.63 (7.63) | 4.34 (4.34) | 2.87 (2.87) | 1.52x | 2.66x | | 2 | 7.96 (7.95) | 4.38 (4.37) | 2.91 (2.90) | 1.51x | 2.74x | | 4 | 8.36 (8.35) | 5.02 (5.01) | 3.43 (3.42) | 1.47x | 2.44x | | 8 | 9.33 (9.23) | 7.08 (6.98) | 4.59 (4.51) | 1.54x | 2.03x | | 16 (SGLang's path in every arm) | 9.80 (9.59) | 9.77 (9.57) | 9.78 (9.58) | 1.00x | 1.00x | | Conc. | Off | Main | This PR | vs main | vs off | | --- | --- | --- | --- | --- | --- | | 1 | 7.65 (7.65) | 4.35 (4.35) | 2.90 (2.90) | 1.50x | 2.64x | | 2 | 7.98 (7.96) | 4.38 (4.37) | 2.95 (2.93) | 1.49x | 2.71x | | 4 | 8.36 (8.34) | 5.01 (5.01) | 3.45 (3.44) | 1.45x | 2.42x | | 8 | 9.30 (9.17) | 7.05 (6.98) | 4.59 (4.53) | 1.54x | 2.03x | | 16 (SGLang's path) | 10.34 (9.40) | 10.32 (9.37) | 10.34 (9.39) | 1.00x | 1.00x | TP2, 2048 / 256 (4-row cap: 8 and 16 requests run SGLang's path in every arm): | Conc. | Off | Main | This PR | vs main | vs off | | --- | --- | --- | --- | --- | --- | | 1 | 9.22 (9.22) | 5.62 (5.62) | 3.87 (3.88) | 1.45x | 2.38x | | 2 | 9.53 (9.52) | 5.72 (5.71) | 3.95 (3.94) | 1.45x | 2.41x | | 4 | 10.35 (10.33) | 7.23 (7.21) | 4.94 (4.93) | 1.46x | 2.10x | | 8 | 11.98 (11.78) | 11.98 (11.78) | 11.99 (11.80) | 1.00x | 1.00x | | 16 | 13.06 (12.63) | 13.06 (12.62) | 13.09 (12.62) | 1.00x | 1.00x | TP8, 2048 / 256: | Conc. | Off | Main | This PR | vs main | vs off | | --- | --- | --- | --- | --- | --- | | 1 | 7.21 (7.21) | 3.73 (3.73) | 2.62 (2.62) | 1.43x | 2.75x | | 2 | 7.42 (7.41) | 3.76 (3.75) | 2.64 (2.64) | 1.42x | 2.81x | | 4 | 7.87 (7.85) | 4.21 (4.21) | 3.13 (3.13) | 1.35x | 2.51x | | 8 | 8.36 (8.32) | 5.26 (5.22) | 3.86 (3.84) | 1.36x | 2.17x | | 16 (SGLang's path) | 8.04 (7.95) | 8.04 (7.95) | 8.04 (7.95) | 1.00x | 1.00x | Both rounds agree within 0.6% on every "on" point. The TP2 runs used an earlier build of this PR whose code objects are byte-identical to the head's (two comments differ). ### Kernel-only `test/kernels/glm53_flash/time_layers.py`: one HIP graph of a whole step of launches, each layer its own random weights (so they stream from HBM), no profiler; µs per launch. Model order is layers 3 to 43 (30 KDA, 11 sparse attention at position 2,047) on two runtimes, as the spec builds them. The table was measured with the same method before the review round; after it (warming removed), the committed driver gives 50.1 / 61.8 / 88.7 µs per launch in model order at TP4, 2 / 4 / 8 rows (table: 50.0 / 61.6 / 88.3). | TP | Rows | KDA | Sparse attention, pos 640 | Sparse attention, pos 8191 | Model order (pos 2047) | | --- | --- | --- | --- | --- | --- | | 2 | 2 | 107.0 → 69.0 | 109.4 → 66.7 | 131.4 → 83.6 | 111.2 → 69.5 (1.60x) | | 2 | 4 | 135.7 → 91.9 | 155.4 → 93.5 | 186.4 → 109.0 | 144.0 → 93.0 (1.55x) | | 4 | 2 | 78.4 → 48.1 | 86.8 → 53.2 | 108.7 → 68.0 | 84.3 → 50.0 (1.69x) | | 4 | 4 | 92.0 → 59.1 | 109.8 → 66.2 | 129.9 → 81.6 | 99.3 → 61.6 (1.61x) | | 4 | 8 | 133.5 → 86.5 | 163.4 → 92.0 | 198.1 → 110.4 | 145.3 → 88.3 (1.65x) | | 8 | 2 | 64.3 → 41.4 | 74.7 → 48.0 | 97.6 → 63.4 | 69.4 → 43.8 (1.59x) | | 8 | 4 | 73.3 → 52.7 | 89.7 → 58.1 | 111.3 → 73.7 | 80.0 → 55.0 (1.46x) | | 8 | 8 | 94.6 → 68.7 | 120.1 → 78.9 | 144.3 → 94.2 | 105.3 → 72.2 (1.46x) | Main → this PR, µs per layer, medians of 7 trials of 200 replays, two passes in alternating order (they agree within 0.3 µs). - FP8 projections read in place (main's choice, kept) against row-major copies, the sparse-attention launch: in place is 1.4 to 2.4 µs faster at TP2, 0.4 to 0.9 µs faster at TP4, and within 0.5 µs either way at TP8. - TP8 gains less than TP4 at 4 and 8 rows: the kernel was tuned at TP4, and at TP8 the per-rank work is smaller, so the launch's chain of cross-CU and cross-GPU hand-offs is a larger share. ### Tests On the head after the review, rebased onto main after #27 merged (`e102e1c`). The PR's diff is unchanged by the rebase (same patch as before it), and all 24 code objects (TP 2, 4 and 8; KDA and sparse attention; every width; preshuffled and row-major FP8) are byte-identical to the pre-rebase head's; main's `common/` changes since are comments only. SGLang ab580b2e06: - unit (CPU): 121 passed on SGLang main f91b81b82a (main: 120). On ab580b2e06, 119 passed; the 2 failures are deepseek_v41's contract cases, which need an SGLang newer than ab580b2e06 since #17, and fail the same way on main. - kernels: `check_layer.py` in every mode (`--writes` with a pad row, `--compare-shuffled`, `--check-routing`, `--check-keys`, `--replays 20 --layers 4` at every width) at TP2 and TP4: all pass, on this PR's head and on #27 alone. Every build's kernel symbol is distinct (`glm53_flash_layer_s<rows>_tp<tp>_k<kda>_d<dsa>_q1_h<shuffled>`, TP 2, 4 and 8, KDA and sparse attention, every width). - `SGLANG_AMD_MONO_CHECK` at TP4: 779 checks, 0 violations (minimum cosine 0.989 sparse-attention streams, 0.986 KDA streams, 0.9992 MLA KV rows, 0.999999 KDA state). - Speed, the head against the previous head (`720c4dc`, with the code warming), same server flags, two rounds, TP4 2048 / 256, median TPOT: 2.87 → 2.88 ms at one request, 3.43 → 3.45 at 4, 4.59 → 4.61 at 8 (+0.3 to +0.6%; ITL likewise). The speed tables above were measured on the previous head; with the warming removed, the head is that much slower. Kernel-only, model order at TP4 (3 alternating passes, µs per launch, previous head → head): sparse attention at position 640, 49.4 → 49.7 / 61.2 → 61.4 / 88.1 → 88.1 at 2 / 4 / 8 rows; at 8191, 53.5 → 54.0 / 65.4 → 65.7 / 93.0 → 93.4. Sparse attention alone at 8191: 68.1 → 67.6 / 81.5 → 80.9 / 110.4 → 110.1. On the previous head (`720c4dc`, the same kernel with the code warming): `test_glm53_flash_kernels.py` over TP 2, 4 and 8 (8 GPUs) 5 passed, `test_guard_hip.py` 4 passed, and `test_glm53_flash_server.py` at TP4 4 passed (GSM8K past the largest width, decode against prefill KL, the check mode, the injected fault). ### Startup - Launch buffers: 490.2 MiB at TP2 (main 490), 214.0 MiB at TP4 (main 213.7), 184.0 MiB at TP8 (main 184); weight copies unchanged (`weights.py` is not touched). - Build with the FlyDSL cache warm: 1.8 s at TP2, 2.3 to 2.5 s at TP4, 2.6 s at TP8. - Server ready after 156 to 166 s in every arm (weights from the page cache). ## 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: Hassan Eslami <hassan@marlowe.ai>
jhinpan
pushed a commit
that referenced
this pull request
Oct 25, 2026
## Motivation After #22, a routed decode step of GLM-5.3-Flash runs layers 3 to 43 as launches (41 of them). Everything else in the step is still SGLang's: the embedding and its all-reduce, the dense layers 0 to 2, layer 44, the final norm, the LM head and the logits all-gather. Together these take about 0.35 ms of the 2.86 ms step at TP4 with one request. This PR runs every layer and the tail as launches. That makes 46 launches per step: layer 0 from the tokens, layers 1 and 2 (dense), layers 3 to 44, and one tail launch. Result on MI355X, 2048 in / 256 out, median TPOT, #22 → this PR: - TP4, one request: #22 2.88 → **2.53 ms** (1.14x). - TP4, 8 requests: #22 4.61 → **4.25 ms** (1.08x). - On the heads before #22's review round, TP2 and TP8 at one request gained 1.12x and 1.17x over #22, and the TP4 ratios were the same as above. - Past the largest width, every arm runs SGLang's path and they match within 0.4%. Accuracy does not decrease (TP4, same server flags, on vs off): GSM8K (1,319, thinking) 96.82% off, 97.12% on; AIME25 (30) 86.7% in both arms; RULER2 NIAH at 128K 87.5% off, 93.75% on. Scope: - GPU: MI355X (gfx950), TP 2, 4 or 8. - Steps: decode steps of 1 to 8 rows (1 to 4 at TP2), as in #22. - Coverage: all 45 layers, the embedding, the final norm, the LM head and the logits all-gather. Output assembly and sampling stay SGLang's. - SGLang: no change; tested on main ab580b2e06. ## Modifications ### 1. Dense, token and tail launches (`kernel.py`, `layout.py`, `config.py`, `runner.py`) There are three new builds of the layer launch, each selected by a `FlashBuild` field that is off by default. With those fields off, the KDA, sparse-attention and FFN launches compile to the same code objects as #22, byte for byte, at every TP size and width. Each build has its own kernel symbol: `glm53_flash_tail_s<rows>_tp<tp>` for the tail, and `kernel_symbol`'s layer name with `n1` (dense), `e1` (tokens) and `h1` (AITER-shuffled dense weights) for the others. A compile-only build of every key the spec and the tests build (TP 2, 4 and 8; every width; preshuffled and row-major FP8) gives 72 distinct symbols. The new builds use #22's guard as is (`_Guard`, whose stop words come from `guard.ERROR_WORDS`), and none of #22's removed warming. - **`dense`** (with `kda`): a dense layer (0 to 2). - KDA attention, as the KDA launch runs it. - The dense SwiGLU MLP on its FP8 gate | up and down weights with 128 x 128 block scales. - The weights are read in place, row-major or AITER-shuffled (`dense_shuffled`; the two layouts are tested bitwise equal). - bf16 activations, then the MLP's TP reduce. - **`tokens`** (with `dense`): layer 0, the step's first launch. - The attention seam looks up this rank's rows of the vocab-parallel embedding for the step's tokens. - It pushes them to every rank and sums the ranks' rows in rank order, as SGLang's embedding all-reduce does (one rank holds each token, so the sum is the row). - One CTA writes the step's KDA state indices and sparse-attention row table for the later launches, in exactly the spec's own format. - The width's rows past the step's are pads. - **`head`**: the tail. - The last layer's owed mHC post, then the mean of the four streams, the final RMSNorm, this rank's LM-head rows, and the logits' all-gather into f32 [rows, vocab]. - The gathered logits use two symmetric slots by launch parity, so a rank one tail ahead never overwrites logits a slower peer is still reading. Off TP4, the dense launch needed two geometry fixes. At TP2 it stages its intermediate in a larger LDS region. At TP8 it runs the down projection's 12 K blocks two per wave, with the spare blocks scaled by 0. ### 2. Kernel tests and timing (`check_layer.py`, `test_glm53_flash_kernels.py`, `time_layers.py`) `time_layers.py` (#22's kernel-only driver) times the new kinds too: `--kind dense`, `embed` and `head`, and `--kind full` for the whole served step (layer 0 from the tokens, the dense layers, layers 3 to 44 and the tail). The kernel tests all run at TP 2, 4 and 8: - **Each new kind:** rank agreement, clean error words, and only its own rows written. - **Graph replays:** 20 replays of each new kind equal the eager step bitwise, including 4 tails back to back. - **Dense layout:** shuffled FP8 equals row-major bitwise. - **`--check-dense`:** the MLP output against an f32 MLP on the launch's own FFN input stays within the spec's bound; a negative control must fail it. - **`--check-embed`:** the token launch equals the dense launch fed the gathered rows (outputs and state writes, bitwise), at 2 of 2, 3 of 4 and 5 of 8 rows, and writes the spec's tables. - **`--check-head`:** the logits match a torch tail and an f32 tail; a negative control must fail. ### 3. The spec (`spec.py`, `weights.py`, `check.py`) - **Layer 0** takes the step's tokens. SGLang's embedding is not run, and the spec's torch row and index tables go away (layer 0's launch writes them). Layers 1 and 2 are dense launches. - **Layer 44** is covered. - After each KDA launch, the spec runs SGLang's decode track-state copy (`_track_mamba_state_decode`). At the last mamba layer, that copy saves every KDA layer's state for the radix cache, exactly as SGLang's own decode does. - The refusal of a covered last KDA layer goes. - **The tail** runs after layer 44. - SGLang's final norm is not run. - The logits processor returns the tail's logits through `_get_logits`, so the output assembly and sampling stay SGLang's (see Reviewer notes). - A routed step with input embeddings runs the model's own forward. - **Refusals:** - a dense MLP that is not block FP8 (e.g. bf16); - an embedding or LM head that is quantized, padded, has added rows, or is not sharded over the TP group; - an f32 LM head; - logits gathered other than over the TP group (DP attention, DP LM head); - a logit scale or soft-capping. - **Hidden states.** When the step captures hidden states (fixed per decode graph), the logits processor is handed the final norm's rows (`normed_rows`), which it stores. Otherwise the last layer's x stands in. The launches stay on either way. - **Step tables.** The token launch's step-input tables are kept by content. A graph reads the table its eager warm-up step made, and an eager step on other tensors takes its own table, so it never rewrites a captured one. - **Weights:** `weights.py` reads the dense MLP, the embedding and the LM head in place. There are no new copies. - **`SGLANG_READS`** gains the logits processor's `_get_logits` and the flags above, `vocab_parallel_embedding`'s shard indices, `_track_mamba_state_decode`, `residual.batch.clear` and `capture_hidden_mode`. - **`SGLANG_AMD_MONO_CHECK`:** - layer 0 is checked against SGLang's embedding plus layer 0; - the dense layers are checked, including their MLP against an f32 MLP; - layer 44 is checked on its contracted output; - the tail's logits are checked against SGLang's final norm and logits, and against an f32 tail. ### 4. Bounds (`NUMERICS`) Bounds are relative, at 3x the largest value `SGLANG_AMD_MONO_CHECK` measured in serving at TP2, TP4 and TP8 (the server check's prompts, 1, 3 and 8 or 4 rows, positions past 2048): | Bound | Value | Serving check, max (TP2 / TP4 / TP8) | What it catches | | --- | --- | --- | --- | | `dense.x_vs_fp32.rel` | 1.2e-2 | 3.1e-3 / 3.0e-3 / 2.9e-3 | Errors in the dense MLP: against an f32 MLP on the launch's own FFN input. `check_layer.py --check-dense`: one rank's down scales x2 gives 0.33 to 0.70 | | `dense.mlp_in.rel` | 2e-2 | 6.2e-3 / 5.3e-3 / – | The launch's FFN input against the input of SGLang's MLP (a forward pre-hook): attention, o_proj, the mixes and `post_attention_layernorm`. Control, the dense launches' `post_attention_layernorm` x1.1: 0.099 to 0.101, failing on all three dense layers (57 of 57 checks) | | `dense.streams.rel` | 0.11 | 3.4e-2 / 3.4e-2 / 3.4e-2 | Gross errors only; see below | | `dense.conv.rel` | 3e-3 | 9.3e-4 / 9.4e-4 / 8.8e-4 | The conv state rows the dense layers write | | `dense.state.rel` | 3.5e-3 | 9.1e-4 / 7.7e-4 / 1.1e-3 | The recurrent state rows | | `head.logits.rel` | 1.8e-3 | 2.4e-4 / 3.7e-4 / 3.1e-4 | Against SGLang's own norm, LM head and gather. `check_layer.py --check-head`: final norm x1.01 gives 9e-3 | | `head.logits_vs_fp32.rel` | 0.06 | 9e-5 / 1.3e-4 / 2.2e-4 | Errors in the tail: each row's distance to an f32 tail in excess of SGLang's | - **Why `dense.streams.rel` is loose.** On layers 0 to 2 the streams are the embedding's, which are small. There, the difference from SGLang is SGLang's own FP8-activation MLP error: 3.6e-2 from an f32 MLP, against the launch's 3.7e-3. So this bound measures stock's error. The bounds referenced to f32 (`dense.x_vs_fp32`, `head.logits_vs_fp32`) and to SGLang's MLP input (`dense.mlp_in`) catch real errors, together with the state bounds. - **Negative controls.** The kernel tests' controls above are caught at TP 2, 4 and 8. - **Every bound is checked.** `check.REPORTED` lists each kind's outputs. `_report` rejects any other output, and a unit test keeps every bound to a listed output. - **No `head.logits.max_abs`.** A single bf16 ulp of a logit of 8 or more (0.0625) exceeded a fixed absolute bound. The two relative logits bounds do not depend on the logits' scale. ## Accuracy Tests Setup: MI355X, TP4, `amd/GLM-5.3-Flash-Quark-MXFP4` at b5688f25, sgl-eval 0.1.3, the e2e test's server flags. Each arm is a fresh server on the same GPUs, at most 8 requests at a time. Measured on this PR's head before #22's review round, whose launches are unchanged here. @jhinpan's rerun on `d9cd3c4` gives GSM8K 97.19% on against 97.04% off (11 vs 9 discordant, McNemar two-sided p = 0.82). | Eval | Settings | Off | On (this PR) | Paired | | --- | --- | --- | --- | --- | | GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 96.82% | 97.12% | 9 vs 13 discordant; one-sided McNemar p = 0.86 (on worse), two-sided 0.52 | | AIME25 (30 x 1) | thinking, T 0.6, top-p 0.95, 65,536 tokens | 86.7% | 86.7% | 2 vs 2 discordant | | RULER2 NIAH at 128K (4 subtasks x 2) | thinking off in the request, greedy, 4,096 tokens | 87.5% | 93.75% | mk basic / easy 1.0 both, mv basic 1.0 both, mv easy 0.5 / 0.75 | | KL decode vs prefill (2, 3, 7, 8 prompts, 512 tokens) | greedy | 0.0256 | 0.0162 | #22: 0.0133 | All five cases of upstream's e2e test pass at TP4 on this PR's head (`88e1aed`), with `HSA_DISABLE_COREDUMP_ON_EXCEPTION=1` for the fault case: GSM8K past the largest width; decode against prefill (KL); the check test (874 checks, 0 violations; also 0 at TP2); the returned hidden states (each decode row times the LM head gives the sampled token: 21 of 21 rows); and the fault stop. ## Speed Tests and Profiling Setup as #22: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`, SGLang main ab580b2e06, a fresh server per arm with the e2e test's flags, `sglang.benchmark.serving` random 2048 / 256, max(16, 4c) prompts after c warm-up requests. Two rounds, the second in reverse order; the table gives the mean of the two rounds' median TPOT, with ITL in parentheses. The rounds agree within 0.2%. TP4, 2048 / 256, #22 at `e102e1c` against this PR at `d6ecd5e`. Rebased on main, this PR's code objects are byte-identical to `d6ecd5e`'s, and main's launches are #22's: | Conc. | #22 | This PR | vs #22 | | --- | --- | --- | --- | | 1 | 2.88 (2.88) | 2.53 (2.53) | 1.14x | | 2 | 2.92 (2.92) | 2.56 (2.55) | 1.14x | | 4 | 3.44 (3.43) | 3.09 (3.08) | 1.11x | | 8 | 4.61 (4.52) | 4.25 (4.16) | 1.08x | | 16 (SGLang's path) | 9.79 (9.58) | 9.79 (9.56) | 1.00x | @jhinpan's rerun on the previous heads (#22 `34898f4`, this PR `d9cd3c4`, TP4) gives the same ratios: 2.92 → 2.57, 2.96 → 2.60, 3.47 → 3.12 and 4.64 → 4.29 ms at 1 / 2 / 4 / 8 requests, and 7.72 / 8.00 / 8.41 / 9.36 ms with the monokernels off. ### Kernel-only (`time_layers.py`) `torchrun --nproc-per-node 4 test/kernels/glm53_flash/time_layers.py --kind <kind>`: µs per launch at 2 / 4 / 8 rows (random weights, one HIP graph a step). TP4 is measured on this PR's head (`88e1aed`); TP2 on `d9cd3c4`, whose code objects are the same. | TP | `model` (#22's 41 launches) | `full` (this PR's 46) | `dense` | `embed` | `head` | | --- | --- | --- | --- | --- | --- | | 4 | 50.2 / 61.5 / 88.7 | 49.9 / 61.0 / 86.6 | 39.6 / 45.2 / 55.3 | 40.9 / 47.0 / 59.0 | 58.4 / 61.7 / 66.6 | | 2 | 69.7 / 93.0 | 70.1 / 92.0 | | | | - `model` runs on code objects that are byte-identical to #22's. - **Spills:** no build that the spec or the kernel tests make spills VGPRs or uses scratch, at any TP size, width or FP8 layout. #22's 32 builds are byte-identical in this PR, and all 72 kernel symbols are distinct. ### Startup and memory - **Launch buffers:** 508.3 MiB at TP2 (#22 490.2), 236.7 MiB at TP4 (#22 214.0) and 206.5 MiB at TP8 (#22 184.0). The new runtimes are layer 0, the dense layers and the tail, with the tail's logits slots. - **KV pool:** - TP4: 9,010,176 tokens against #22's 9,011,456 (-0.01%) and off's 9,016,320 (-0.07%). - TP8: 9,970,688 against 9,973,760 (-0.03%). - TP2: 7,022,848 against 7,018,240 (page granularity). - **Build time:** 3.7 s (TP2), 5.6 s (TP4) and 6.3 s (TP8) with the FlyDSL cache warm. A tree's first server boot builds in about 24 s. - **Server ready:** 131 s at TP4 in every arm; 146 to 261 s at TP2 and 161 to 196 s at TP8 across arms (SGLang's radix init wait varies). ## Reviewer notes - **`_get_logits` is the one SGLang method the spec replaces** (a hook in SGLang's monokernel interface could remove it later). - The spec binds the logits processor's `_get_logits` to its own at build time. The binding is active only inside a routed step and otherwise calls SGLang's. - It is needed because the tail launch already produces the gathered f32 logits [rows, vocab], exactly what `_get_logits` returns after its LM head and all-gather. - Replacing just that method keeps `LogitsProcessor.forward`'s pruning, output assembly, logprob paths and SGLang's sampler unchanged. The alternative would re-implement the processor's forward in the plugin. - The stored hidden states are the final norm's rows whenever the step captures them, as SGLang's own forward stores them (see Hidden states). - Everything else the spec does is its own forward (as in #6), calling SGLang only through declared reads. - **The dense and tail bounds** come from the serving check (section 4). - **Numerics.** Only the new launches change numerics. Commit 1 leaves #22's KDA, sparse-attention and FFN code objects byte-identical, and the token launch writes the same tables as the spec's torch ops did (tested). Commit 4 changes only the bounds. ## Checklist - [x] `pre-commit run --all-files` passes. - [x] Tests under `test/` cover the change. - [x] The SGLang commit or PR this was tested against is named above. - [x] Accuracy and speed results are attached when kernel or forward code changes. - [ ] A new model adds one line to `.github/CODEOWNERS` that names its owner. (No new model.) --------- Co-authored-by: Hassan Eslami <hassan@marlowe.ai>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
On 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
amdmonokernel 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:
Scope:
Modifications
How a decode step runs
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 amdGrey 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:
o_projand its TP all-reduce.Changes to SGLang (
python/sglang/srt)--monokernel-backend amd, SGLang runs as before. The spec registers for gfx950 only, so CUDA and other ROCm GPUs do not change.Glm5NextForConditionalGeneration.forward, for each decode step that the spec covers. SGLang still does the scheduling, the cache management, the sampling and all other steps.SGLANG_READS.test/unit/test_contract.pychecks them on CPU against the installed SGLang.Files
sgl_amd_mono/glm53_flash/gfx950/: the decoder-layer kernel from [ROCm][Perf][GLM-5.3-Flash] Mono decode layer for GLM-5.3-Flash on MI355X vllm-project/vllm#61662 at ca328e297b. vLLM adapted it from [FlyDSL] GLM-5 fused decode-layer MonoKernel ROCm/aiter#6173. The first GLM-5.3-Flash commit copies the files verbatim. The next commit adapts them:common/guard.py(deadline, peer flags, fail-stop).sgl_amd_mono/glm53_flash/spec.py: the spec and its forward. The forward isGlm5NextForConditionalGeneration.forwardfor a covered decode step.DeviceLimits.widthsfor each ISA, not for each TP size. Thus at TP2, the framework still lists 5 to 8 rows as covered.--enable-linear-replayssm, rotary layout, and others.sgl_amd_mono/glm53_flash/check.py: withSGLANG_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 inNUMERICS. 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: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
Accuracy Tests
Setup:
Per layer, against SGLang's own layer on the same input and caches:
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.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):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.
Speed Tests and Profiling
Setup:
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).--enable-linear-replayssmand corrects comments and docs.--kv-cache-dtype bfloat16.--ep-size 1.--moe-runner-backend aiter.--monokernel-backend amd.Method:
sglang.bench_serving, random dataset with fixed lengths (--random-range-ratio 1.0), max(16, 4c) prompts after a warmup.After the largest width, both arms use SGLang's path. The control rows stay within 1%:
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.
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):
down.Startup:
Tests on MI355X in the image above, with SGLang main at 4cca7fcd3e:
commonfixes.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)
--enable-hisparse;--expert-distribution-recorder-mode;moe_layer_freqother than 1;SGLANG_READSdeclares every SGLang name that the spec and the weight readers read.down.common/peer.py's 2 MiB unit is named_GRANULE.**/*glm53_flash* @jhinpan.Checklist
pre-commit run --all-filespasses.test/cover the change.Review and Merge Process
lintandCheck PR titlechecks must pass..github/CODEOWNERSto review.main.