You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
GLM-5 decode at low concurrency is dominated by launch gaps and small kernels. The MonoKernel runs a whole decoder layer in one persistent launch per GPU, but today it only exists inside ATOM. This PR adds it to AITER so vLLM and other engines can use it.
What changes
Adds the GLM-5 fused decode-layer MonoKernel to AITER. Device code is from ROCm/FlyDSL#1204 (@coderfeli); the host wrapper and the TP4, FP8-KV and native-FP4 paths are from ROCm/ATOM#2435 (@XiaobingSuper), at ATOM bdc157af. The first commit carries the kernel sources over unchanged apart from their import paths; the later commits fix the original fused indexer and add a paged-cache indexer for vLLM. The last commit removes the Kimi-K3 configuration and the helpers that came with that copy but are unused here; the GLM-5 kernel and its launch path are unchanged.
One launch per MoE layer and rank covers input RMSNorm, the MLA projections, RoPE and KV-cache insert, sparse attention over precomputed indices, o_proj with an in-kernel all-reduce, the router, and the MXFP4 experts with an in-kernel all-reduce.
glm5_mono_launch_rows() pads a one-row decode step to two rows, because a one-row launch never completes on gfx950 (seen at TP4 with both FP4 MFMA modes).
Storage contract (what a caller must provide)
Input
Contract
Attention
Block-FP8 (128x128) or per-channel FP8, repacked by prepare_glm5_weights
Experts
AITER-shuffled MXFP4 (shuffle_weight(layout=(16, 16)) + e8m0_shuffle), shared expert fused as the last physical expert
KV cache
FP8, one 576-byte row per global slot, no scale
Per row
positions int64, slot_mapping int64 (pad rows -1), CSR sparse_kv_indptr int32 and global sparse indices
Tests
MI355X (gfx950), TP4, op_tests/multigpu_tests/test_flydsl_glm5_mono.py, 4 layers in one HIP graph:
Case
Result
1 row (launched as 2)
finite, real row unchanged when the pad row changes
2, 4, 8 rows
finite, bit-identical across replays
16 rows (two 8-row chunks)
finite, bit-identical across replays
2 real + 2 pad rows (slot -1, empty sparse range)
finite, real rows unchanged when pad rows change
Per-layer latency, TP4, 75 layers with distinct random weights in one graph, block-FP8 attention: 65.8 us (2 rows), 84.8 us (4 rows), 144.7 us (8 rows). Per-channel FP8 attention: 64.9, 92.4, 157.5 us.
Fused indexer on a vLLM paged index cache (opt-in)
Glm5MonoKernel(..., with_indexer=True, index_paged=True, ...) runs the sparse indexer inside the launch:
Index keys live in vLLM's paged FP8 cache: E4M3 values (plain or 16x16-shuffled) plus one power-of-two scale per token.
Each row is its own request. The row's block_table entry locates its keys, and the new key is quantized with vLLM's ue8m0 rule.
The exact top-2048 is published as global slots to out_indices at the row's sparse_kv_indptr offset, so later layers can reuse them.
No LDS context limit: radix keys past the LDS capacity are re-read from the score mailbox.
index_k_bf16=True reads a bf16 index-K projection.
Kernels without the indexer only gain three unused pointer arguments.
Also fixes the original fused indexer: the picks are now stored to LDS (TP8 memory fault), and the TP4 q_b split passes the layout check.
op_tests/multigpu_tests/test_flydsl_glm5_mono_paged_indexer.py, MI355X, TP4 and TP8, all cases pass:
Case
Contexts
4 rows, shuffled
1000, 8192, 30000, 70001
4 rows, plain layout
5, 2047, 2049, 120000
3 rows + 1 pad row
3000, 65536, 100
8 rows
1000 ... 200000
2 rows / 1 row + 1 pad row
8192, 60000 / 40000
Per case: new-key bytes and scale equal the ue8m0 reference, picks match a PyTorch top-2048 (overlap 1.0, one row 0.9995), and the layer output is bit-identical to the external-indices launch fed the same slots.
Top-k selection: one pass builds a 4096-bin histogram of the score keys; a second pass publishes every key above the threshold bin and moves the bin's keys into LDS, where the remaining radix digits and the tie gather run. A full radix over all keys is the fallback when the bin does not fit. The test also covers clustered keys and heavy ties.
Extra time per layer for the in-kernel indexer (75 layers, 2 / 4 / 8 rows): TP4 8K +26 / +38 / +57 us, 60K +82 / +109 / +169 us; TP8 8K +29 / +37 / +48 us, 60K +85 / +106 / +159 us.
8x MI355X, amd/GLM-5.3-Quark-MXFP4-AttnFP8, TP=8, fp8 KV. Random dataset (random-range-ratio 0.8), 8192/1024 at concurrency 4 to 256 and 60000/600 at 4 to 32.
The two builds use different AITER commits. TTFT and every point at concurrency 16 and above, where the fused path is off, agree within 0.7%, so the decode gains come from the fused path.
GPQA Diamond, thinking on, temperature 1 (198 questions)
0.879
0.899
GPQA runs at temperature 1, so one question is 0.5 points.
Dense-MLP layers (dense_experts)
GLM-5's leading dense layers have a 12288-wide MLP, six times the expert width, so the kernel can run them too:
pack_dense_mlp(gate_up, gate_up_scale, down, down_scale, slices) stores a row-major MXFP4 dense MLP as slices expert-shaped slices in the ATOM expert layout (same shuffle_weight / e8m0_shuffle as MoE experts). Slice j holds intermediates j * I / slices onward, so the sum over slices is the dense MLP exactly.
Glm5MonoKernel(..., dense_experts=N) routes every row to slices 0..N-1 with weight 1 and skips the router GEMV; the router stage still produces the activation blocks the experts read. Spare slots have weight 0. It is a build-time option, so MoE layers compile as before.
op_tests/multigpu_tests/test_flydsl_glm5_mono_dense.py, MI355X, TP4 (TP8 gives 0.0010 to 0.0011 in the last column):
Rows
Finite, replay bit-identical
Rel. error vs float32 MLP
Rel. error vs reference with the kernel's MXFP8 and BF16 roundings
2
yes
0.047
0.0006
4
yes
0.046
0.0012
8
yes
0.046
0.0009
After this change test_flydsl_glm5_mono.py (1/2/4/8 rows and pad rows) and test_flydsl_glm5_mono_paged_indexer.py still pass at TP4 and TP8. In vLLM (amd-sriram/vllm#27), all 78 GLM-5.3 layers in the kernel give TPOT -6.1% / -2.9% at concurrency 1 / 4 on TP4 and -8.8% / -5.0% on TP8 (8K), with NIAH 15/15 at both.
After the Kimi-K3 cleanup, the three op tests above give the same results as before at TP4 (MI355X, same build for both runs; 14 cases, identical output lines).
The last commit runs the down-stage routing only for one-row launches; multi-row launches already route before up/gate, so the second pass repeated it. Output is bit-identical (SHA-256 of every rank's output, TP4 and TP8, 2/4/8 rows); 75-layer graph latency drops 1.4 / 0.9 us per layer at TP4 with 2 / 4 rows and is unchanged elsewhere. In vLLM, TP4, 8192/256, both servers with the fused path on, two rounds with the GPUs swapped: TPOT 6.37 -> 6.28 ms (-1.3%) at concurrency 1 and 9.72 -> 9.65 ms (-0.7%) at concurrency 4.
Port of the GLM-5 MonoKernel from ROCm/ATOM#2435 (device code from
ROCm/FlyDSL#1204), with a host API and a 4-GPU op test.
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
Per-row requests via the block table, ue8m0 key insert, top-2048
published as global slots; no LDS context limit.
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
12-bit histogram, then winners and threshold-bin candidates in LDS;
full radix only on overflow. Test covers clustered and tied keys.
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
All standard extended tests (excludes ci:atom_full)
Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 6173 --add-label <label>
One backend per PR:
A PR changes one kernel backend: [Triton/Gluon] (Triton and Gluon count as one), [HIP], [ASM], [CK], [OPUS] or [FlyDSL]. If the title ends up with two backend tags, split the PR -- as stacked pull requests when one part cannot merge without the other.
PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to stop the title rewrites; labels stay in sync either way.
Adds a gfx950 GLM-5 fused decode-layer MonoKernel for low-concurrency inference, including sparse attention, paged indexing, MoE execution, and TP reductions.
Changes:
Adds the persistent FlyDSL kernel, public wrapper, layouts, packing, and IPC runtime.
dense_experts=N routes every row to N expert-shaped slices of a dense MLP
with weight 1 and skips the router GEMV; pack_dense_mlp builds the slices.
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
shuffle_weight restores the input dtype (aiter/ops/shuffle.py:205-206), so production MXFP4 weights can have dtype torch.float4_e2m1fn_x2, not just uint8. Those valid one-byte FP4 tensors are currently classified as W8A8 and the advertised ATOM path is rejected. Recognize the typed FP4 dtype as MXFP4 too.
This issue also appears on line 191 of the same file.
Include launches_per_step in runtime compatibility
Runtime compatibility ignores launches_per_step, even though mailbox tags are computed as step * launches_per_step + layer + 1. Sharing storage between otherwise layout-identical instances with different values can reuse a stale tag (both begin at tag 1), allowing consumers to observe old mailbox data. Include this value in the compatibility check.
The paged-cache check verifies the block and token strides but not the innermost byte stride. An as_strided cache with stride(2) != 1 passes, while the kernel performs contiguous byte loads and writes, so values and scales land at the wrong addresses. Validate stride(2) == 1.
Infer physical experts including the shared expert
When physical_experts is omitted, this defaults to only the routed experts, but the GLM kernel always addresses the shared expert at config.n_experts and the documented storage contract fuses it as the final physical expert. Consequently a valid 257-expert ATOM tensor fails the byte-count check under the dataclass default. Infer routed plus shared experts.
Removes the Kimi-K3 config and branches, the unused MXFP8/MXFP4 helpers
and layouts, and formats.py. Kernel, launch path and layouts unchanged.
Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
The kernel reads SGLang's FP8 projections and experts in place as AITER
(16, 16) tiles and caches Hadamard-rotated index keys. Row-major projections
(for example --fp8-gemm-backend triton), experts off the AITER MoE runner and
an indexer other than SGLang's gfx950 fused one gave wrong tokens without a
refusal; the spec now lists each as a reason. The rotary check compares every
indexer, since the first covered layer reuses the previous top-k and has none.
SGLANG_AMD_MONO_CHECK is read once at build, an always-present SGLang field
is read directly, the vendored config uses msgspec, and every adapted file and
NOTICE name ROCm/aiter#6173.
kevin-mii
pushed a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 18, 2026
## Motivation
On MI355X at TP4 with one request, a GLM-5.3-Flash decode step on
SGLang's path runs 1,301 kernels (torch profiler). Its median TPOT is
7.75 ms. vllm-project/vllm#61662 runs each MoE decoder layer of this
model as one persistent launch on gfx950.
This PR ports that kernel to the `amd` monokernel backend of SGLang (RFC
sgl-project/sglang#43949). With `--monokernel-backend amd`, each covered
layer of a decode step runs as one launch.
Result on MI355X:
- At TP4 with one request, the median TPOT decreases from 7.75 ms to
4.40 ms (1.76x).
- From 1 to 8 requests (1 to 4 at TP2), the speedup is 1.31x to 2.09x.
- Accuracy does not decrease. On AIME26 (30 problems x 8), pass@1 is
89.17% (off) and 91.67% (on). On GSM8K, it is 97.35% (off) and 97.27%
(on).
Scope:
- GPU: MI355X (gfx950), at TP 2, 4 or 8.
- Steps: decode steps of 1 to 8 rows (1 to 4 at TP2). All other steps
use SGLang's path.
- Layers: the MoE decoder layers 3 to 43 (41 launches: 30 KDA, 11 sparse
attention). The dense layers 0 to 2 and the last layer (44) use SGLang's
path.
- SGLang: this PR needs sgl-project/sglang#44353 (merged as ec7ab9d0c2).
It changes no SGLang file.
## Modifications
### How a decode step runs
```mermaid
flowchart TB
F["SGLang monokernel framework<br/>(sglang PR 44353)"] --> Q{"Decode step of 1 to 8 rows?<br/>(1 to 4 at TP2)"}
Q -- "no" --> N["SGLang's path"]
Q -- "yes" --> R["glm53_flash spec forward"]
R --> A["SGLang: embedding and<br/>dense layers 0 to 2"]
A --> L["41 launches, layers 3 to 43<br/>(30 KDA, 11 sparse attention)"]
L --> Z["SGLang: layer 44,<br/>final norm and logits"]
L -.->|reads and writes in place| M[("SGLang: KV, index and KDA state pools,<br/>page table, loaded weights")]
classDef sgl fill:#eef1f4,stroke:#8c959f,color:#1f2328
classDef amd fill:#dafbe1,stroke:#1a7f37,color:#1f2328
class F,Q,N,A,Z,M sgl
class R,L amd
```
Grey boxes are SGLang code, and this PR changes none of it. Green boxes
are new in this PR.
One launch does these operations for one layer, in this sequence:
1. The attention-side mHC.
2. The attention: KDA, or sparse MLA with its kpool indexer and its
cache writes.
3. `o_proj` and its TP all-reduce.
4. The FFN-side mHC and norm.
5. The router.
6. The shared and routed MXFP4 experts.
7. The MoE TP all-reduce.
### Changes to SGLang (`python/sglang/srt`)
- None. This PR adds and changes files in sgl-amd-mono only.
- Without `--monokernel-backend amd`, SGLang runs as before. The spec
registers for gfx950 only, so CUDA and other ROCm GPUs do not change.
- With the flag, SGLang calls the spec's forward, not
`Glm5NextForConditionalGeneration.forward`, for each decode step that
the spec covers. SGLang still does the scheduling, the cache management,
the sampling and all other steps.
- The plugin reads SGLang internals by name. Each module lists these
names in `SGLANG_READS`. `test/unit/test_contract.py` checks them on CPU
against the installed SGLang.
- At startup, the plugin uses more GPU memory for launch buffers and
weight copies (see "Startup").
### Files
- `sgl_amd_mono/glm53_flash/gfx950/`: the decoder-layer kernel from
vllm-project/vllm#61662 at ca328e297b. vLLM adapted it from
ROCm/aiter#6173. The first GLM-5.3-Flash commit copies the files
verbatim. The next commit adapts them:
- Every wait runs under `common/guard.py` (deadline, peer flags,
fail-stop).
- The kernel reads the KDA conv and recurrent state pool of SGLang in
place. It also reads the DSA MLA and index caches, the compress tails
and the 64-token page table in place.
- The page table is a launch argument, so each captured graph reads its
own table.
- The kernel reads the FP8 attention projections in the AITER (16, 16)
layout that SGLang's gfx950 block-FP8 linear method makes. It reads the
MXFP4 experts in the layout that SGLang's Quark path makes.
- Like the vLLM kernel, this kernel quantizes the expert inputs to
MXFP8. SGLang's Quark path quantizes them to MXFP4 (W4A4, as the
checkpoint declares). Thus the two paths are not bit-exact.
- `sgl_amd_mono/glm53_flash/spec.py`: the spec and its forward. The
forward is `Glm5NextForConditionalGeneration.forward` for a covered
decode step.
- A step with fewer rows than a kernel width (2, 4 or 8) runs at the
next width. The added rows are pad rows, and pad rows write no cache.
- At TP2, an 8-row sparse-attention launch does not fit in the LDS, so
steps of 5 to 8 rows at TP2 use SGLang's path.
- The framework sets `DeviceLimits.widths` for each ISA, not for each TP
size. Thus at TP2, the framework still lists 5 to 8 rows as covered.
- The spec refuses unsupported setups at startup. The config and model
checks give their reasons: page size, KV dtype, MoE runner,
`--enable-linear-replayssm`, rotary layout, and others.
- If the kernel cannot read a weight layout (for example a bf16 shared
expert), the build stops and gives the reason.
- `sgl_amd_mono/glm53_flash/check.py`: with `SGLANG_AMD_MONO_CHECK=1`
(eager mode only), each covered layer also runs SGLang's layer on the
same input and caches. The check compares the written streams, each
state row and each cache row with the bounds in `NUMERICS`. It logs each
value that is outside its bound.
- `sgl_amd_mono/common/`: three fixes for problems that the
GLM-5.3-Flash launches found:
- The guard reads its peer flags in loads of 4 dwords or fewer. At TP8,
the 8-dword load failed instruction selection.
- Peer buffers use whole 2 MiB granules. Before this fix, a peer faulted
on guard flags after the last whole granule.
- The peer handshake fails on all ranks when the buffer sizes are
different. Before this fix, ranks that built buffers in different orders
waited until the deadline.
- Registration in `sgl_amd_mono/__init__.py`, a README row, a NOTICE
entry, and these tests:
- `test/unit/test_glm53_flash_spec.py` (CPU)
- `test/kernels/test_glm53_flash_kernels.py` (8 GPUs)
- `test/e2e/test_glm53_flash_server.py` (TP4)
### Not in this PR
- An FP8 KV cache.
- Speculative decoding (MTP).
- Decode steps of more than 8 rows.
## Accuracy Tests
Setup:
- MI355X, TP4, amd/GLM-5.3-Flash-Quark-MXFP4 at b5688f25, sgl-eval
0.1.2.
- The two arms ran on the same GPUs. Each arm used a fresh server with
CUDA graphs and 8 or fewer requests at a time. Thus all decode steps of
the "on" arm ran the launches.
| Eval | Settings | Monokernels off | Monokernels on |
| --- | --- | --- | --- |
| GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 97.35% | 97.27%
(13 vs 14 discordant, exact McNemar p = 1.00) |
| AIME25 (30 x 4 seeds) | thinking, T 0.6, top-p 0.95, 65,536 tokens |
pass@1 82.50% | pass@1 91.67% |
| RULER2 NIAH at 1M tokens (4 subtasks x 2) | no thinking, greedy, 256
tokens | 40.62% | 44.12% |
- With the monokernels, AIME25 is 9.2 points higher. The paired
bootstrap 95% CI over the 30 problems is [+1.7, +17.5] points. A sign
test over the problems gives p = 0.09 (10 problems better, 3 worse).
- 14 of 120 samples reached the token limit in the "off" arm, and 6 in
the "on" arm. The median output was 7,348 tokens (off) and 5,051 tokens
(on).
- A possible cause of this difference is the MXFP8 expert inputs. I did
not isolate it.
Per layer, against SGLang's own layer on the same input and caches:
- Setup: `SGLANG_AMD_MONO_CHECK=1`, eager mode. Batches of 1, 3 and 8
rows (4, not 8, at TP2). Positions up to 2.4k. 779 checks for each TP
size.
- Result: no value is outside its bound.
| Output | TP2 | TP4 | TP8 |
| --- | --- | --- | --- |
| written streams, cosine | >= 0.974 | >= 0.992 | >= 0.980 |
| MLA KV rows, cosine | >= 0.9991 | >= 0.9992 | >= 0.9993 |
| index cache values, cosine (scales bit-exact) | >= 0.99998 | >=
0.99998 | >= 0.99998 |
| KDA recurrent state, cosine | >= 0.999999 | >= 0.999999 | >= 0.999998
|
The streams differ by approximately the error of the MoE: the kernel
uses MXFP8 expert inputs, and SGLang's path uses MXFP4.
Decode against prefill (`sglang.test.kl_test_utils`, 8 prompts x 512
tokens, mean KL):
- SGLang's path: 0.038.
- Monokernels on: 0.009 to 0.042 for each batch, over two e2e runs.
- The e2e test sets the bound to 0.06. It runs one graph for each batch
of 2, 3, 7 and 8 rows.
- With an injected fault (one page-table pointer for all graphs), the
3-row batch measured 0.146.
GPQA-Diamond did not run, because the dataset is gated on the Hub.
### AIME26 on daf4739
GSM8K is near its ceiling for this model, so it finds only large errors.
The SGLang contribution guide asks for a harder eval, and the GLM-5.3
cookbook reports AIME25 and AIME26. Thus this PR also reports AIME26 on
its final head.
- Setup: MI355X, TP4, SGLang main 4cca7fcd3e, sgl-amd-mono daf4739,
sgl-eval 0.1.3 (the version that SGLang pins). Each arm used a fresh
server.
- 8 or fewer requests ran at a time, so each decode step of the "on" arm
ran the launches.
- Generation: thinking, T 1.0 and top-p 0.95 (the recommended settings
of the model), 65,536 max tokens.
| AIME26 (30 problems x 8) | Monokernels off | Monokernels on |
| --- | --- | --- |
| pass@1 | 89.17% | 91.67% |
| majority@8 | 96.67% | 96.67% |
| samples at the token limit | 14 of 240 | 12 of 240 |
| median output tokens | 8,463 | 6,516 |
- Paired over the 30 problems, the "on" arm is 2.50 points higher
(bootstrap 95% CI [+0.0, +5.4]). 9 problems are better, 3 are worse and
18 are tied (sign test p = 0.15).
- The direction is the same as on AIME25: higher accuracy and shorter
outputs with the monokernels. A possible cause is the MXFP8 expert
inputs.
## Speed Tests and Profiling
Setup:
- Hardware: one node with 8 x MI355X (gfx950).
- Image: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`
(`sha256:e4498f5d595f`), with ROCm 10.0.0, torch 2.11.0, FlyDSL 0.3.4.1
and AITER e7d2453f (= `AITER_COMMIT_DEFAULT`).
- SGLang: sgl-project/sglang#44353 at 04fcb04e34. Its monokernel
framework is the same as the merged ec7ab9d0c2.
- sgl-amd-mono: this PR on 3970534. The rebase onto 815b104 changed no
code. The later commits change no launched code:
- One commit moves an import below the TP2 fallback.
- daf4739 refuses `--enable-linear-replayssm` and corrects comments and
docs.
- The e2e tests ran again on daf4739 (see the end).
- Server: the cookbook's MI355X high-throughput cell (no speculative
decoding), with these differences:
- This checkpoint. The cookbook's OneNexus/GLM-5.3-Flash-MXFP4 keeps the
shared expert in bf16, and the spec refuses that.
- `--kv-cache-dtype bfloat16`.
- `--ep-size 1`.
- TileLang DSA prefill and decode.
- `--moe-runner-backend aiter`.
- The only difference between the arms is `--monokernel-backend amd`.
Method:
- Workload: `sglang.bench_serving`, random dataset with fixed lengths
(`--random-range-ratio 1.0`), max(16, 4c) prompts after a warmup.
- Each arm ran on a fresh server, in two rounds. The second round
swapped the arm order and the GPU set (TP8: the order only).
- At load, each server checked that no other process held memory on its
GPUs.
- The table gives the median TPOT per request in ms over both rounds,
with a bootstrap 95% CI.
| TP | In/out | Conc. | Monokernels off | Monokernels on | Speedup |
| --- | --- | --- | --- | --- | --- |
| 2 | 8192/256 | 1 | 9.31 [9.28, 9.31] | 5.68 [5.68, 5.68] | 1.64x |
| 2 | 8192/256 | 2 | 9.58 [9.57, 9.60] | 5.73 [5.72, 5.75] | 1.67x |
| 2 | 8192/256 | 4 | 10.25 [10.23, 10.27] | 7.15 [7.13, 7.18] | 1.43x |
| 2 | 512/128 | 1 | 9.17 [9.14, 9.18] | 5.48 [5.46, 5.48] | 1.67x |
| 2 | 512/128 | 4 | 10.31 [10.26, 10.33] | 6.88 [6.86, 6.91] | 1.50x |
| 4 | 8192/256 | 1 | 7.75 [7.74, 7.75] | 4.40 [4.40, 4.40] | 1.76x |
| 4 | 8192/256 | 2 | 7.98 [7.98, 8.00] | 4.43 [4.42, 4.49] | 1.80x |
| 4 | 8192/256 | 4 | 8.39 [8.38, 8.40] | 5.03 [5.02, 5.03] | 1.67x |
| 4 | 8192/256 | 8 | 9.23 [9.22, 9.49] | 7.03 [7.01, 7.39] | 1.31x |
| 4 | 512/128 | 1 | 7.66 [7.64, 7.66] | 4.16 [4.16, 4.17] | 1.84x |
| 4 | 512/128 | 2 | 7.91 [7.91, 7.92] | 4.18 [4.18, 4.19] | 1.89x |
| 4 | 512/128 | 4 | 8.35 [8.33, 8.36] | 4.83 [4.82, 4.83] | 1.73x |
| 4 | 512/128 | 8 | 9.39 [9.35, 9.41] | 6.76 [6.74, 6.77] | 1.39x |
| 8 | 8192/256 | 1 | 7.30 [7.30, 7.30] | 3.77 [3.77, 3.78] | 1.93x |
| 8 | 8192/256 | 2 | 7.51 [7.50, 7.53] | 3.79 [3.79, 3.80] | 1.98x |
| 8 | 8192/256 | 4 | 7.90 [7.89, 7.90] | 4.25 [4.24, 4.25] | 1.86x |
| 8 | 8192/256 | 8 | 8.43 [8.42, 8.62] | 5.32 [5.31, 5.51] | 1.58x |
| 8 | 512/128 | 1 | 7.21 [7.20, 7.21] | 3.52 [3.52, 3.53] | 2.05x |
| 8 | 512/128 | 2 | 7.41 [7.41, 7.42] | 3.55 [3.54, 3.56] | 2.09x |
| 8 | 512/128 | 4 | 7.85 [7.84, 7.86] | 4.04 [4.03, 4.04] | 1.95x |
| 8 | 512/128 | 8 | 8.40 [8.39, 8.41] | 5.06 [5.05, 5.07] | 1.66x |
After the largest width, both arms use SGLang's path. The control rows
stay within 1%:
- each TP at 16 requests: 0.999x to 1.004x;
- TP2 at 8 requests: 0.998x and 1.004x.
### Cookbook shape on daf4739
The GLM-5.3 cookbook reports speed at random 8192 in / 1024 out. This
table uses that shape at TP4 on daf4739.
- Each arm ran on a fresh server, in two rounds. The second round
swapped the arm order and the GPU set.
- TPOT is the median per request over both rounds, with a bootstrap 95%
CI. TTFT is the median per request. Throughput is the mean over the
rounds.
| Conc. | TPOT off (ms) | TPOT on (ms) | Speedup | TTFT off / on (ms) |
Output tok/s per GPU off / on |
| --- | --- | --- | --- | --- | --- |
| 1 | 7.75 [7.74, 7.75] | 4.39 [4.39, 4.39] | 1.76x | 235 / 232 | 31.3 /
53.8 |
| 2 | 8.01 [8.00, 8.03] | 4.42 [4.41, 4.43] | 1.81x | 321 / 317 | 59.6 /
104.3 |
| 4 | 8.40 [8.40, 8.41] | 5.02 [5.01, 5.02] | 1.67x | 211 / 199 | 115.8
/ 190.9 |
| 8 | 9.22 [9.21, 9.38] | 7.02 [7.01, 7.17] | 1.31x | 1,116 / 1,112 |
176.6 / 220.3 |
| 16 | 10.08 [9.81, 10.30] | 10.07 [9.80, 10.31] | 1.00x | 1,082 / 1,104
| 337.5 / 337.7 |
At 16 requests, a step has more than 8 rows, so both arms use SGLang's
path.
Profile (TP4, one request, torch profiler, rank 0):
- Monokernels on: a step runs 184 kernels, and the GPU is busy for 4.28
ms. Of this time, the 41 layer launches use 3.44 ms (84 us each). The
trace shows each launch as `down`.
- Monokernels off: a step runs 1,301 kernels.
Startup:
- In a new container, the build of all widths takes 25 to 32 s (FlyDSL
and comgr compilation). When the caches hold the kernels, the build
takes approximately 2 s.
- Extra memory for each GPU:
- TP2: 490 MiB of launch buffers and 455 MiB of weight copies. 261 MiB
of the copies is the packed shared expert, which SGLang keeps unfused at
TP2.
- TP4: 214 MiB of launch buffers and 150 MiB of weight copies.
- TP8: 184 MiB of launch buffers and 128 MiB of weight copies.
Tests on MI355X in the image above, with SGLang main at 4cca7fcd3e:
- unit (CPU): 75 passed. They also passed on 04fcb04e34, and on daf4739
with SGLang main 4cca7fcd3e and 2948662026.
- kernels: 24 passed on 5a8b90f:
- Each GLM-5.3-Flash launch kind at TP 2, 4 and 8, on random weights in
SGLang's layouts. All ranks return the same output and write only their
valid rows. AITER-shuffled FP8 projections give the row-major result
bitwise.
- The guard tests.
- The DeepSeek-V4.1 kernel tests with the `common` fixes.
- e2e (TP4) has four tests: GSM8K, KL for each batch, the check mode,
and an injected wait fault that must stop all ranks.
- Before this review: 4 passed (GSM8K 0.99, KL 0.009 to 0.018 for each
batch).
- On daf4739, in one run: 4 passed. GSM8K-200 is 0.99. KL is 0.018,
0.015, 0.010 and 0.027 for batches of 2, 3, 7 and 8 rows. The 779 checks
are within `NUMERICS` (lowest stream cosine 0.9904). After the injected
stall, the faults are only at address 0x0.
The kernel comes from vllm-project/vllm#61662 (Apache-2.0). vLLM adapted
it from ROCm/aiter#6173 (FlyDSL GLM monokernel).
### Review fixes (on top of daf4739)
- New startup refusals:
- an index cache that SGLang keeps row-major (no AITER preshuffle paged
MQA);
- `--enable-hisparse`;
- `--expert-distribution-recorder-mode`;
- SGLang's benchmark routing overrides;
- `moe_layer_freq` other than 1;
- a TP2 server whose decode graphs all exceed the 4-row cap.
- `SGLANG_READS` declares every SGLang name that the spec and the weight
readers read.
- The stage loop no longer overwrites the kernel's symbol. Before, every
variant was named `down`.
- Cleanup:
- startup errors say what they need;
- the docs match the code (the README row, the copy lists, NOTICE's
provenance);
- `common/peer.py`'s 2 MiB unit is named `_GRANULE`.
- The README has one GPU test command per model.
- CODEOWNERS adds `**/*glm53_flash* @jhinpan`.
- Tests on 7944894:
- Setup: MI350X, TP4, SGLang main 4eb606b2ae with
sgl-project/sglang#44601.
- unit: 75 passed.
- e2e: 4 passed. KL was 0.009, 0.014, 0.009 and 0.029. GSM8K-200 was
0.97. The lowest stream cosine was 0.986. The injected fault stopped all
ranks.
- e254332 adds only a startup refusal.
## Checklist
- [x] `pre-commit run --all-files` passes.
- [x] Tests under `test/` cover the change.
- [x] The SGLang commit or PR this was tested against is named above.
- [x] Accuracy and speed results are attached when kernel or forward
code changes.
## Review and Merge Process
1. The `lint` and `Check PR title` checks must pass.
2. Ask a code owner from `.github/CODEOWNERS` to review.
3. Merge with "Squash and merge": the PR title becomes the commit title
on `main`.
---------
Co-authored-by: Kevin Mi <kevin.mi@radixark.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
kevin-mii
pushed a commit
to sgl-project/sgl-amd-mono
that referenced
this pull request
Oct 18, 2026
## Motivation
On MI355X at TP8 with one request, a GLM-5.3 decode step on SGLang's
path has a median TPOT of 11.2 ms. The model is the official
`zai-org/GLM-5.3` (FP8 block-128). ROCm/aiter#6173 runs each GLM-5
decoder layer as one persistent FlyDSL launch. vllm-project/vllm#60831
integrates a variant of this kernel that keeps only MXFP4 experts.
This PR ports the aiter kernel to the `amd` monokernel backend of SGLang
(RFC sgl-project/sglang#43949). It keeps the FP8 block-128 expert path
that the official checkpoint needs. With `--monokernel-backend amd`,
each MoE decoder layer of a covered decode step runs as one launch.
Result on MI355X:
- At TP8 with one request, the median TPOT decreases from 11.24 ms to
5.60 ms (2.01x).
- From 1 to 8 requests at TP8, the speedup is 1.38x to 2.18x. From 1 to
4 requests at TP4, it is 1.45x to 1.82x.
- Accuracy does not decrease. On AIME26 (30 problems x 8), pass@1 is
82.92% (off) and 86.67% (on). On GSM8K, it is 97.35% (off) and 97.27%
(on).
Scope:
- GPU: MI355X (gfx950), at TP 4 or 8.
- Steps: decode steps of 1 to 8 rows (1 to 4 at TP4). All other steps
use SGLang's path.
- Layers: the 75 MoE decoder layers 3 to 77. 18 launches run the
indexer, and 57 reuse the previous top-k. The dense layers 0 to 2 use
SGLang's path.
- SGLang: this PR needs sgl-project/sglang#44353 (merged as ec7ab9d0c2).
It changes no SGLang file.
- Stack: this PR uses the three `common` fixes of #6, which has merged.
Its diff now shows only GLM-5.3.
## Modifications
### How a decode step runs
```mermaid
flowchart TB
F["SGLang monokernel framework<br/>(sglang PR 44353)"] --> Q{"Decode step of 1 to 8 rows?<br/>(1 to 4 at TP4)"}
Q -- "no" --> N["SGLang's path"]
Q -- "yes" --> R["glm53 spec forward"]
R --> A["SGLang: embedding and<br/>dense layers 0 to 2"]
A -- "top-k of layer 2" --> L["75 launches, layers 3 to 77<br/>(18 run the indexer,<br/>57 reuse the previous top-k)"]
L --> Z["SGLang: final norm and logits"]
L -.->|reads and writes in place| M[("SGLang: MLA cache, paged FP8 index cache,<br/>page table, FP8 weights")]
classDef sgl fill:#eef1f4,stroke:#8c959f,color:#1f2328
classDef amd fill:#dafbe1,stroke:#1a7f37,color:#1f2328
class F,Q,N,A,Z,M sgl
class R,L amd
```
Grey boxes are SGLang code, and this PR changes none of it. Green boxes
are new in this PR.
One launch does these operations for one layer, in this sequence:
1. The input RMSNorm, the q_a and kv_a projections, the q_a RMSNorm, and
q_b with RoPE.
2. The KV RMSNorm, the k_pe RoPE and the MLA cache write.
3. The indexer, on an indexer layer only:
- the index K, Q and weight projections;
- the index K norm, RoPE, Hadamard transform and cache write;
- the index scores and the exact top-2048.
A layer without an indexer reads the top-k of the layer before it.
4. The absorbed query (W_UK), the sparse MLA, W_UV and `o_proj`.
5. The attention TP all-reduce and the residual add.
6. The post-attention RMSNorm, the router in f32, and the FP8
quantization of the activations.
7. The top-8 selection and the experts: 1 shared and 8 routed, in FP8
block-128.
8. The MoE TP all-reduce and the residual add.
### Changes to SGLang (`python/sglang/srt`)
- None. This PR adds and changes files in sgl-amd-mono only.
- Without `--monokernel-backend amd`, SGLang runs as before. The spec
registers for gfx950 only, so CUDA and other ROCm GPUs do not change.
- With the flag, SGLang calls the spec's forward, not
`DeepseekV2ForCausalLM.forward`, for each decode step that the spec
covers. SGLang still does the scheduling, the cache management, the
sampling and all other steps.
- The plugin reads SGLang internals by name. Each module lists these
names in `SGLANG_READS`. `test/unit/test_contract.py` checks them on CPU
against the installed SGLang.
- At startup, the plugin uses more GPU memory. Thus the KV pool is 0.7%
smaller at TP8 and 3.4% smaller at TP4 (see "Startup").
### Files
- `sgl_amd_mono/glm53/gfx950/`: the decoder-layer kernel from
ROCm/aiter#6173 at 56d6390def. The first GLM-5.3 commit copies the files
verbatim. The next commits adapt them:
- Every wait runs under `common/guard.py` (deadline, peer flags,
fail-stop).
- The kernel reads SGLang's FP8 weights in place, in AITER's (16, 16)
layout. These are the attention projections, the indexer projections and
the experts. They include the two gate/up paths that load expert weights
outside the shared unit.
- The kernel reads SGLang's bf16 MLA cache, paged FP8 index cache and
64-token page table in place.
- The page table and its stride are launch arguments, so each captured
graph reads its own table.
- Index keys and queries get the 128-point Hadamard transform after
RoPE. SGLang's gfx950 fused indexer caches rotated keys, so the cache
keeps one format. The spec refuses all other indexers, because they
cache unrotated keys.
- `sgl_amd_mono/glm53/spec.py`: the spec and its forward. The forward is
`DeepseekV2ForCausalLM.forward` for a covered decode step.
- Layer 3 reuses the top-k of layer 2. The forward copies SGLang's
indices into the compact buffer that the launch reads.
- A step with fewer rows than a kernel width (2, 4 or 8) runs at the
next width. One row runs at width 2, because a one-row launch does not
complete on gfx950. The added rows are pad rows, and pad rows write no
cache.
- At TP4, a step of more than 4 rows uses SGLang's path. At TP4, an
8-row launch is slower than SGLang's path: TPOT 18.15 ms against 16.95
ms at 8192 / 256 and c8.
- The spec refuses unsupported setups at startup and gives the reasons.
Some examples:
- another checkpoint format, page size, KV dtype or MoE runner;
- an unfused shared expert;
- FP8 projections outside AITER's (16, 16) layout, for example with
`--fp8-gemm-backend triton`;
- experts that do not use the AITER MoE runner;
- an indexer other than SGLang's gfx950 fused indexer.
- Weight copies:
- W_UK and W_UV, packed from SGLang's per-tensor FP8 `w_kc` and `w_vc`
(the same values that SGLang multiplies by);
- the router and the indexer's `weights_proj`, packed for the MFMA
tiles;
- the indexer's k LayerNorm in f32;
- one bf16 rotary table.
- `sgl_amd_mono/glm53/check.py`: with `SGLANG_AMD_MONO_CHECK=1` (eager
mode only), each covered layer also runs SGLang's layer on the same
input and caches. The check compares the stream, the MLA KV row and the
index keys (after dequantization) with the bounds in `NUMERICS`.
- Registration in `sgl_amd_mono/__init__.py`, a README row, a NOTICE
entry, and these tests:
- `test/unit/test_glm53_spec.py` (CPU)
- `test/e2e/test_glm53_server.py` (TP8)
- There is no random-weight launch test like the one for glm53_flash.
The e2e check compares each covered layer with SGLang's layer on the
real weights. The cost is that the test needs the checkpoint.
### Routing numerics
- The kernel routes in f32, as the model defines it: sigmoid plus
correction bias for the choice, and sigmoid for the weights.
- SGLang's AITER top-k casts the correction bias of GLM-5.3
(approximately 34) to bf16. Thus SGLang's choice and weights are
different from f32 routing.
- For this reason, the check mode routes SGLang's reference layer in
f32.
### Not in this PR
- An FP8 KV cache.
- Speculative decoding (MTP).
- MXFP4 checkpoints.
- Decode steps of more than 8 rows (4 at TP4).
## Accuracy Tests
Setup (also for the speed tests):
- Hardware: one node with 8 x MI355X (gfx950).
- Image: `rocm/sgl-dev:v0.5.21-rocm10-mi35x-20261014`
(`sha256:e4498f5d595f`), with ROCm 10.0.0, torch 2.11.0, FlyDSL 0.3.4.1
and AITER e7d2453f.
- SGLang: main at 4cca7fcd3e.
- sgl-amd-mono: this PR. The TP8 runs are older than two commits, and
neither commit changes what TP8 launches:
- cbc8f15 adds the TP4 row cap.
- 02d42fa adds the review fixes: startup refusals, the rotary check over
each indexer, msgspec config containers, comments and notices.
- The e2e tests ran again on 02d42fa (see the end).
- Model: `zai-org/GLM-5.3` at ba32f8c19c (FP8 block-128).
- Server: the cookbook's AMD recipe (TP8, TileLang DSA prefill and
decode, `--chunked-prefill-size 131072`, no MTP), with
`--mem-fraction-static 0.80`, `--max-running-requests 96` and
`--cuda-graph-max-bs-decode 96`. SGLang selects the bf16 KV cache,
64-token pages and the AITER MoE runner. TP4 is not a cookbook cell, and
it uses the same flags.
- The only difference between the arms is `--monokernel-backend amd`.
Each arm runs on a fresh server.
- Evals: sgl-eval 0.1.2 at TP8, with 8 or fewer requests at a time.
| Eval | Settings | Monokernels off | Monokernels on |
| --- | --- | --- | --- |
| GSM8K (1,319) | thinking, T 1.0, top-p 0.95, seed 42 | 97.35% | 97.27%
(10 vs 11 discordant, exact McNemar p = 1.00) |
| AIME25 (30 x 4 seeds) | thinking, T 0.6, top-p 0.95, 65,536 tokens |
90.00% | 89.17% (-0.83 points, bootstrap 95% CI [-5.8, +4.2] over
problems; 2 better, 4 worse; sign test p = 0.69) |
| RULER2 NIAH at 1M tokens (4 subtasks x 2) | no thinking, greedy, 256
tokens | 31.56% | 25.60% |
- AIME25: 11 answers (off) and 13 answers (on) reached the token limit.
The median length was 7,630 tokens (off) and 9,351 tokens (on). For each
problem, the monokernels ran longer on 17 problems and shorter on 12
(sign test p = 0.46).
- NIAH: `mk_niah_basic` scores 100 in both arms. Most other answers
reached the 256-token limit (4 of 8 samples off, 6 of 8 on). The
reasoning lists the needles, and the limit stops the answer. Thus the
multi-value subtasks mostly measure where the answer stops.
- The arms also use different routing precision: f32 (on) and SGLang's
bf16 correction bias (off). See "Routing numerics".
Per layer, against SGLang's own layer on the same input and caches:
- Setup: `SGLANG_AMD_MONO_CHECK=1`, eager mode. Batches of 1, 3 and 8
rows. Positions up to 2.4k. 1,425 checks for each TP size.
- Result: no value is outside its bound.
- The TP4 run is older than the TP4 row cap, so it also covers 8-row
launches.
| Output | TP4 | TP8 |
| --- | --- | --- |
| written stream, cosine | >= 0.9916 | >= 0.9938 |
| MLA KV row, cosine | >= 0.9996 | >= 0.9995 |
| index keys after dequantization, cosine | >= 0.99993 | >= 0.99994 |
Decode against prefill (`sglang.test.kl_test_utils`, 512 tokens, mean KL
for each batch):
- 0.009, 0.017, 0.016 and 0.030 for batches of 2, 3, 7 and 8 rows.
- The bound is 0.06.
### AIME26 on 02d42fa
GSM8K is near its ceiling for this model, so it finds only large errors.
The SGLang contribution guide asks for a harder eval, and the GLM-5.3
cookbook reports AIME25 and AIME26. Thus this PR also reports AIME26 on
its final head.
- Setup: MI355X, TP8, SGLang main 4cca7fcd3e, sgl-amd-mono 02d42fa,
sgl-eval 0.1.3 (the version that SGLang pins). Each arm used a fresh
server.
- 8 or fewer requests ran at a time, so each decode step of the "on" arm
ran the launches.
- Generation: thinking, T 1.0 and top-p 0.95 (the recommended settings
of the model), 65,536 max tokens.
| AIME26 (30 problems x 8) | Monokernels off | Monokernels on |
| --- | --- | --- |
| pass@1 | 82.92% | 86.67% |
| majority@8 | 96.67% | 96.67% |
| samples at the token limit | 34 of 240 | 29 of 240 |
| median output tokens | 12,197 | 11,492 |
- Paired over the 30 problems, the "on" arm is 3.75 points higher
(bootstrap 95% CI [-1.2, +9.6]). 7 problems are better, 5 are worse and
18 are tied (sign test p = 0.77). The difference is within noise.
- 12% to 14% of the samples reach the 65,536-token limit. Thus these
scores are lower than scores with a larger limit, and you cannot compare
them with the cookbook numbers. Both arms use the same limit.
- The arms also use different routing precision (see "Routing
numerics").
## Speed Tests and Profiling
Method:
- The setup is the same as for the accuracy tests.
- Workload: `sglang.bench_serving` with the random dataset at fixed
lengths (8192 in and 256 out, 512 in and 128 out). Each concurrency c
uses max(16, 4c) prompts after a warmup.
- Each arm runs on a fresh server, in two rounds. The second round swaps
the arm order. TP4 also swaps the GPU set.
- At load, each server checks that no other process holds memory on its
GPUs.
- The tables give the median TPOT per request over both rounds, with a
bootstrap 95% CI.
TP8 (GPUs 0-7):
| Shape | c | Monokernels off (ms) | Monokernels on (ms) | Speedup |
| --- | --- | --- | --- | --- |
| 8192 / 256 | 1 | 11.24 [11.23, 11.24] | 5.60 [5.60, 5.60] | 2.01x |
| 8192 / 256 | 2 | 12.36 [12.35, 12.37] | 5.78 [5.77, 5.78] | 2.14x |
| 8192 / 256 | 4 | 13.30 [13.29, 13.31] | 7.33 [7.32, 7.34] | 1.81x |
| 8192 / 256 | 8 | 15.85 [15.81, 17.80] | 11.51 [11.49, 13.47] | 1.38x |
| 8192 / 256 | 16 | 18.65 [18.65, 22.51] | 18.94 [18.91, 22.53] | 0.98x
|
| 512 / 128 | 1 | 10.79 [10.79, 10.79] | 5.43 [5.43, 5.44] | 1.99x |
| 512 / 128 | 2 | 11.95 [11.93, 11.97] | 5.48 [5.47, 5.48] | 2.18x |
| 512 / 128 | 4 | 12.81 [12.81, 12.83] | 6.86 [6.85, 6.86] | 1.87x |
| 512 / 128 | 8 | 15.41 [15.40, 15.41] | 10.93 [10.91, 10.95] | 1.41x |
| 512 / 128 | 16 | 17.71 [17.68, 17.73] | 17.68 [17.66, 17.70] | 1.00x |
- At c16, a step has more than 8 rows, so both arms use SGLang's path.
The 16-row decode steps take the same time in both arms (median ITL 18.6
ms at 8192 / 256).
- The 8192 / 256 c16 TPOT is different because of the scheduler. It
admitted one burst of prefix-cached prompts as 1 + 10 + 5 requests, not
as 1 + 15. This stops approximately 10 running requests once, for
approximately 0.1 s.
TP4 (round 1 on GPUs 0-3, round 2 on GPUs 4-7):
| Shape | c | Monokernels off (ms) | Monokernels on (ms) | Speedup |
| --- | --- | --- | --- | --- |
| 8192 / 256 | 1 | 12.02 [11.99, 12.04] | 7.09 [7.09, 7.10] | 1.70x |
| 8192 / 256 | 2 | 13.05 [13.04, 13.06] | 7.25 [7.24, 7.26] | 1.80x |
| 8192 / 256 | 4 | 14.68 [14.65, 14.70] | 10.10 [10.08, 10.11] | 1.45x |
| 8192 / 256 | 8 | 16.96 [16.93, 19.52] | 16.95 [16.93, 19.51] | 1.00x |
| 8192 / 256 | 16 | 23.19 [23.17, 27.96] | 23.17 [23.06, 27.95] | 1.00x
|
| 512 / 128 | 1 | 11.61 [11.58, 11.64] | 6.95 [6.95, 6.95] | 1.67x |
| 512 / 128 | 2 | 12.63 [12.61, 12.65] | 6.94 [6.94, 6.95] | 1.82x |
| 512 / 128 | 4 | 14.33 [14.31, 14.34] | 9.70 [9.68, 9.71] | 1.48x |
| 512 / 128 | 8 | 16.74 [16.71, 16.75] | 16.73 [16.68, 16.75] | 1.00x |
| 512 / 128 | 16 | 22.26 [22.20, 22.29] | 22.22 [22.18, 22.23] | 1.00x |
- At c8 and c16, a TP4 step has more than 4 rows, so both arms use
SGLang's path.
- Before the TP4 row cap, the 8-row launches made c8 slower:
- 8192 / 256: 18.15 ms against 16.95 ms (0.93x).
- 512 / 128: 17.44 ms against 16.76 ms (0.96x).
### Cookbook shape on 02d42fa
The GLM-5.3 cookbook reports speed at random 8192 in / 1024 out. This
table uses that shape at TP8 on 02d42fa.
- Each arm ran on a fresh server, in two rounds. The second round
swapped the arm order.
- TPOT is the median per request over both rounds, with a bootstrap 95%
CI. TTFT is the median per request. Throughput is the mean over the
rounds.
| Conc. | TPOT off (ms) | TPOT on (ms) | Speedup | TTFT off / on (ms) |
Output tok/s per GPU off / on |
| --- | --- | --- | --- | --- | --- |
| 1 | 11.24 [11.24, 11.24] | 5.58 [5.58, 5.59] | 2.01x | 541 / 540 |
10.6 / 20.2 |
| 2 | 12.34 [12.33, 12.35] | 5.75 [5.75, 5.76] | 2.15x | 287 / 285 |
19.7 / 40.9 |
| 4 | 13.27 [13.25, 13.28] | 7.28 [7.28, 7.29] | 1.82x | 365 / 353 |
31.3 / 49.7 |
| 8 | 15.80 [15.79, 16.26] | 11.44 [11.42, 11.89] | 1.38x | 3,067 /
3,059 | 49.1 / 62.5 |
| 16 | 18.62 [18.60, 19.51] | 18.60 [18.60, 19.51] | 1.00x | 1,080 /
1,037 | 86.7 / 87.0 |
At 16 requests, a step has more than 8 rows, so both arms use SGLang's
path.
Profile (torch profiler, rank 0, TP8, one request, mean of 10 decode
steps):
- SGLang's path: a decode step runs 1,821 GPU kernels in 11.70 ms of GPU
time.
- Monokernels on: a decode step runs 212 kernels in 5.59 ms. The 75
launches use 4.76 ms (63.5 us each). The other 137 kernels (0.83 ms) are
the embedding, the dense layers 0 to 2, the final norm and the logits.
Startup:
- In a new container, the kernels build in approximately 40 s at TP8 (6
kernels) and 30 s at TP4 (4 kernels).
- A later server in the same container builds them in approximately 2 s,
because the container keeps the compiled code objects
(`~/.cache/comgr`).
- During the build, PyTorch memory increases by 511.8 MiB for each GPU
at TP8 and by 573.1 MiB at TP4.
- The weight copies are 491 MiB at TP8 and 622 MiB at TP4:
- W_UK / W_UV: 131 or 263 MiB;
- the router: 225 MiB;
- the rotary tables: 128 MiB;
- `weights_proj`: 7 MiB.
- The peer buffers are outside PyTorch.
- In total, the KV pool decreases by 0.7% at TP8 (from 1,533,824 to
1,522,624 tokens). At TP4, it decreases by 3.4% (from 539,584 to 521,024
tokens).
Tests on MI355X in the image above, with SGLang main at 4cca7fcd3e:
- unit (CPU): 88 passed on 02d42fa, with SGLang main 4cca7fcd3e and
2948662026. Two of the cases fail on cbc8f15: the layout and indexer
refusals, and the rotary check over each indexer.
- e2e (TP8): 4 passed on 02d42fa in one run:
- GSM8K-200 through the launches: 0.995.
- KL: 0.036, 0.025, 0.018 and 0.021 for batches of 2, 3, 7 and 8 rows.
- The check mode: all values are within `NUMERICS`.
- An injected wait fault stops all ranks, and the faults are only at
address 0x0.
- The lowest stream cosine of the check mode changes between runs, on
one row. It was 0.9938 and 0.9939 in the earlier runs, and 0.9895 (layer
31, position 30) on 02d42fa. All other checks are 0.994 or more.
The kernel comes from ROCm/aiter#6173 at 56d6390def (Apache-2.0, FlyDSL
Project Contributors). The file headers are kept.
vllm-project/vllm#60831 integrates the MXFP4 variant.
### Review fixes (on top of 02d42fa)
- main is merged in (#6 with its review fixes, and #10).
- New startup refusals:
- unfused DSA top-k (`SGLANG_DSA_FUSE_TOPK=0`). With it, layer 2 hands
layer 3 request-local positions instead of KV slots, and layers 3 to 5
would attend other requests' rows.
- `--enable-hisparse`, whose logical KV slots the kernels write raw.
- `--expert-distribution-recorder-mode`.
- SGLang's benchmark routing overrides.
- a scaled RoPE such as YaRN. SGLang scales the softmax by mscale
squared, while the kernel's scale is fixed.
- an f32 router gate that bf16 cannot hold exactly.
- a served context longer than the copied rotary table.
- a TP4 server whose decode graphs all exceed the 4-row cap.
- The kernel cache key holds the sources digest, as GLM-5.3-Flash's and
DeepSeek-V4.1's do. So an edit to `common/`, or to an aiter helper the
kernel imports, rebuilds it.
- `SGLANG_READS` declares every SGLang name that the spec reads.
- Cleanup:
- the lineage comments are gone;
- the stale kernel and runner docstrings match the code;
- the guard flags size uses `common.guard.FLAG_BYTES`;
- a few names are clearer, and `gfx950/` has a package `__init__`.
- CODEOWNERS adds `**/*glm53* @jhinpan`, above the Flash line, because
the last match wins.
- The README has the GLM-5.3 row and its own test command.
- Tests on a470273:
- `test/unit`: 89 passed, with SGLang main 4eb606b2ae and
sgl-project/sglang#44601.
- pre-commit passes.
- Not executed: the TP8 e2e. It needs 8 x MI355X and the zai-org/GLM-5.3
checkpoint.
## Checklist
- [x] `pre-commit run --all-files` passes.
- [x] Tests under `test/` cover the change.
- [x] The SGLang commit or PR this was tested against is named above.
- [x] Accuracy and speed results are attached when kernel or forward
code changes.
## Review and Merge Process
1. The `lint` and `Check PR title` checks must pass.
2. Ask a code owner from `.github/CODEOWNERS` to review.
3. Merge with "Squash and merge": the PR title becomes the commit title
on `main`.
---------
Co-authored-by: Kevin Mi <kevin.mi@radixark.ai>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
maeehart
added a commit
to maeehart/sgl-amd-mono
that referenced
this pull request
Oct 22, 2026
…the Quark MXFP4 checkpoint
This adds a GLM-5.3 spec for the Quark MXFP4 checkpoint (amd/GLM-5.3-Quark-MXFP4-AttnFP8) with an fp8 KV cache on gfx950. The FlyDSL kernel and its helpers are copied from vllm-project/vllm#60831 at 6deee07c52, which adapted ROCm/aiter#6173, and keep their license headers. The SGLang side reads the weights, the fp8 KV cache, the index-key cache and the page table that SGLang already has, so SGLang itself needs no change.
glm_moe_dsa picks the spec from the checkpoint: glm53 for the block-FP8 checkpoint and glm53_mxfp4 for the Quark MXFP4 checkpoint. The kernel runs a decode step with one launch per MoE layer (layers 3 to 77). At TP4 it takes steps of 1 to 4 requests, because its 8-row launch was slower than SGLang's own path there. At TP8, which is not tested yet, it takes steps of 1 to 8 requests. Every other step runs SGLang's normal path. SGLang's startup line still lists batch sizes 1 to 8 at TP4, because the spec declares the widths of TP8. The spec refuses expert parallelism, so the server needs --ep-size 1, and it is switched on with --monokernel-backend amd.
Tested on 4x MI355X at TP4 with SGLang 8e3e2980db: GSM8K 5-shot gives the same accuracy with and without the kernel, and greedy replies drift between batch widths no more than on SGLang's normal path. test/unit/test_glm53_mxfp4_spec.py checks the spec selection and the refusals on CPU, and test/kernels/glm53_mxfp4/test_paged_indexer.py checks the paged indexer on the GPU.
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
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
GLM-5 decode at low concurrency is dominated by launch gaps and small kernels. The MonoKernel runs a whole decoder layer in one persistent launch per GPU, but today it only exists inside ATOM. This PR adds it to AITER so vLLM and other engines can use it.
What changes
Adds the GLM-5 fused decode-layer MonoKernel to AITER. Device code is from ROCm/FlyDSL#1204 (@coderfeli); the host wrapper and the TP4, FP8-KV and native-FP4 paths are from ROCm/ATOM#2435 (@XiaobingSuper), at ATOM
bdc157af. The first commit carries the kernel sources over unchanged apart from their import paths; the later commits fix the original fused indexer and add a paged-cache indexer for vLLM. The last commit removes the Kimi-K3 configuration and the helpers that came with that copy but are unused here; the GLM-5 kernel and its launch path are unchanged.One launch per MoE layer and rank covers input RMSNorm, the MLA projections, RoPE and KV-cache insert, sparse attention over precomputed indices, o_proj with an in-kernel all-reduce, the router, and the MXFP4 experts with an in-kernel all-reduce.
aiter/ops/flydsl/kernels/glm5_mono/aiter/ops/flydsl/glm5_mono.pyglm5_mono_launch_rows()op_tests/multigpu_tests/test_flydsl_glm5_mono.pyglm5_mono_launch_rows()pads a one-row decode step to two rows, because a one-row launch never completes on gfx950 (seen at TP4 with both FP4 MFMA modes).Storage contract (what a caller must provide)
prepare_glm5_weightsshuffle_weight(layout=(16, 16))+e8m0_shuffle), shared expert fused as the last physical expertpositionsint64,slot_mappingint64 (pad rows -1), CSRsparse_kv_indptrint32 and global sparse indicesTests
MI355X (gfx950), TP4,
op_tests/multigpu_tests/test_flydsl_glm5_mono.py, 4 layers in one HIP graph:Per-layer latency, TP4, 75 layers with distinct random weights in one graph, block-FP8 attention: 65.8 us (2 rows), 84.8 us (4 rows), 144.7 us (8 rows). Per-channel FP8 attention: 64.9, 92.4, 157.5 us.
Fused indexer on a vLLM paged index cache (opt-in)
Glm5MonoKernel(..., with_indexer=True, index_paged=True, ...)runs the sparse indexer inside the launch:block_tableentry locates its keys, and the new key is quantized with vLLM's ue8m0 rule.out_indicesat the row'ssparse_kv_indptroffset, so later layers can reuse them.index_k_bf16=Truereads a bf16 index-K projection.Also fixes the original fused indexer: the picks are now stored to LDS (TP8 memory fault), and the TP4 q_b split passes the layout check.
op_tests/multigpu_tests/test_flydsl_glm5_mono_paged_indexer.py, MI355X, TP4 and TP8, all cases pass:Per case: new-key bytes and scale equal the ue8m0 reference, picks match a PyTorch top-2048 (overlap 1.0, one row 0.9995), and the layer output is bit-identical to the external-indices launch fed the same slots.
Top-k selection: one pass builds a 4096-bin histogram of the score keys; a second pass publishes every key above the threshold bin and moves the bin's keys into LDS, where the remaining radix digits and the tie gather run. A full radix over all keys is the fallback when the bin does not fit. The test also covers clustered keys and heavy ties.
Extra time per layer for the in-kernel indexer (75 layers, 2 / 4 / 8 rows): TP4 8K +26 / +38 / +57 us, 60K +82 / +109 / +169 us; TP8 8K +29 / +37 / +48 us, 60K +85 / +106 / +159 us.
End-to-end results in vLLM (TP8)
Served through vllm-project/vllm#60216 (measured from its fork copy, amd-sriram/vllm#27), which uses the fused indexer on every indexer layer.
8x MI355X,
amd/GLM-5.3-Quark-MXFP4-AttnFP8, TP=8, fp8 KV. Random dataset (random-range-ratio0.8), 8192/1024 at concurrency 4 to 256 and 60000/600 at 4 to 32.VLLM_ROCM_GLM_MONO_DECODE=0ac68c308+ vllm-project/vllm#50535, #51309, #51314, #59280, #59420, #59421, #59422, amd-sriram/vllm#17, #18, #277bfe279d+ amd-sriram#7VLLM_ROCM_GLM_MONO_DECODE=15b2f7d1d+ amd-sriram#7 + amd-sriram#10The two builds use different AITER commits. TTFT and every point at concurrency 16 and above, where the fused path is off, agree within 0.7%, so the decode gains come from the fused path.
Geomean over the whole sweep: 8192/1024 TPOT -11.7%, output throughput +13.2%, TTFT +0.0%; 60000/600 TPOT -7.0%, output throughput +6.4%, TTFT -0.4%.
niah_single_2, 64K (lm_eval)GPQA runs at temperature 1, so one question is 0.5 points.
Dense-MLP layers (
dense_experts)GLM-5's leading dense layers have a 12288-wide MLP, six times the expert width, so the kernel can run them too:
pack_dense_mlp(gate_up, gate_up_scale, down, down_scale, slices)stores a row-major MXFP4 dense MLP asslicesexpert-shaped slices in the ATOM expert layout (sameshuffle_weight/e8m0_shuffleas MoE experts). Slice j holds intermediatesj * I / slicesonward, so the sum over slices is the dense MLP exactly.Glm5MonoKernel(..., dense_experts=N)routes every row to slices 0..N-1 with weight 1 and skips the router GEMV; the router stage still produces the activation blocks the experts read. Spare slots have weight 0. It is a build-time option, so MoE layers compile as before.op_tests/multigpu_tests/test_flydsl_glm5_mono_dense.py, MI355X, TP4 (TP8 gives 0.0010 to 0.0011 in the last column):After this change
test_flydsl_glm5_mono.py(1/2/4/8 rows and pad rows) andtest_flydsl_glm5_mono_paged_indexer.pystill pass at TP4 and TP8. In vLLM (amd-sriram/vllm#27), all 78 GLM-5.3 layers in the kernel give TPOT -6.1% / -2.9% at concurrency 1 / 4 on TP4 and -8.8% / -5.0% on TP8 (8K), with NIAH 15/15 at both.After the Kimi-K3 cleanup, the three op tests above give the same results as before at TP4 (MI355X, same build for both runs; 14 cases, identical output lines).
The last commit runs the down-stage routing only for one-row launches; multi-row launches already route before up/gate, so the second pass repeated it. Output is bit-identical (SHA-256 of every rank's output, TP4 and TP8, 2/4/8 rows); 75-layer graph latency drops 1.4 / 0.9 us per layer at TP4 with 2 / 4 rows and is unchanged elsewhere. In vLLM, TP4, 8192/256, both servers with the fused path on, two rounds with the GPUs swapped: TPOT 6.37 -> 6.28 ms (-1.3%) at concurrency 1 and 9.72 -> 9.65 ms (-0.7%) at concurrency 4.