[ROCm][Perf][GLM-5] Fused sparse-layer decode via AITER MonoKernel

[amd-sriram]

Overview

Opt-in decode path for GLM-5 (glm_moe_dsa) on ROCm gfx950. With VLLM_ROCM_GLM_MONO_DECODE=1, each decoder layer of a pure decode step (3 dense, 75 MoE) runs as one AITER MonoKernel launch per GPU. That launch includes the sparse indexer and both tensor-parallel all-reduces.

Depends on ROCm/aiter#6173, which adds the kernel (aiter.ops.flydsl.glm5_mono). The kernel comes from ROCm/FlyDSL#1204 and ROCm/ATOM#2435. The structure follows the MiniMax-M3 fused decode in #59653.

Claims

• 8x MI355X, GLM-5.3 MXFP4, TP8, fp8 KV, serving benchmark:
  • TPOT -40.9% at concurrency 4 and -27.4% at 8 (8K input / 1K output).
  • TPOT -18.4% at concurrency 4 and -8.1% at 8 (60K input / 600 output).
• 4x MI355X, TP4: TPOT 1.41x faster at concurrency 1 (8K input / 256 output).
• Accuracy matches the regular path: NIAH 1.000 vs 1.000, GPQA Diamond 0.899 vs 0.879, GSM8K within run-to-run noise.
• The last two commits run the dense layers in the kernel and advance its epoch once per step: TPOT -6.1% / -2.9% at concurrency 1 / 4 on TP4 and -8.8% / -5.0% on TP8, on top of the above (8K).
• Default off. When it is on, steps the kernel does not cover (prefill, mixed, more than 8 tokens) keep the regular layers and their CUDA graphs.

Validation

Base used for these runs. All measurements were taken on ac68c308 plus #50535, #51309, #51314, #59280, #59420, #59421 and #59422. AITER was 5b2f7d1d plus ROCm/aiter#6173 (and 7bfe279d for the baseline). This branch is the same 4 files rebased onto current main:

• No conflicts.
• No upstream change under vllm/models/deepseek_v32/amd/.
• Every API the path uses is unchanged.
• The rebased branch has not been re-run on GPU yet; that will be done before marking ready.

Serving, TP8. vllm bench serve-style random dataset, random-range-ratio 0.8. 8192/1024 swept at concurrency 4 to 256, 60000/600 at 4 to 32. Baseline is the same stack with the flag off.

ISL/OSL	Concurrency	TPOT off	TPOT on	Change	Output tok/s/GPU off	on	Change
8192/1024	4	13.10 ms	7.74 ms	-40.9%	36.67	61.02	+66.4%
8192/1024	8	15.05 ms	10.92 ms	-27.4%	62.31	85.72	+37.6%
8192/1024	16 (path off)	17.89 ms	17.84 ms	-0.3%	105.16	106.11	+0.9%
60000/600	4	24.59 ms	20.07 ms	-18.4%	16.62	19.62	+18.1%
60000/600	8	36.41 ms	33.46 ms	-8.1%	23.13	24.88	+7.6%
60000/600	16 (path off)	62.53 ms	62.51 ms	0.0%	27.22	27.34	+0.4%

• Geomean over the full sweep: 8K TPOT -11.7% and output throughput +13.2%; 60K TPOT -7.0% and output throughput +6.4%. TTFT is unchanged (+0.0% / -0.4%).
• The two builds used different AITER commits. TTFT and every point at concurrency 16 and above, where the path is off, agree within 0.7%.

Serving, TP4. Same build for both runs, only the flag differs. 8K input / 256 output:

Concurrency	TPOT on	TPOT off	Speedup
1	8.11 ms	11.43 ms	1.41x
2	8.98 ms	12.56 ms	1.40x
4	12.08 ms	15.98 ms	1.32x
8	25.98 ms	29.21 ms	1.12x
16	30.14 ms	29.95 ms	path off

This TP4 table was measured with the indexer still outside the kernel. With the in-kernel indexer, a single request's step time drops further at TP4: from 8.44 to 6.92 ms at 8K, and from 8.70 to 7.80 ms at 60K.

Accuracy.

Eval	Flag off	Flag on
TP8, NIAH niah_single_2 at 64K (lm_eval)	1.000	1.000
TP8, GPQA Diamond, thinking on, temperature 1, 198 questions	0.879	0.899
TP4, NIAH at 8K/32K/60K, 15 prompts, 4 concurrent	15/15	15/15
TP4, GSM8K 5-shot strict, 250 questions, 8 concurrent	0.900	0.896
TP8, GSM8K, same setup	0.916	0.912 and 0.884 (two runs)

GSM8K on 250 questions moves by about 3 points between runs. GPQA at temperature 1 moves by 0.5 points per question.

• No NaN in a probe at 1, 2, 4 and 8 concurrent requests, at TP4 or TP8.

Dense layers in the kernel, one epoch advance per step (last two commits).

Check (4x MI355X, TP4, 8192/256)	PR head before	With these commits
Layers in the kernel	75 (18 with the fused indexer)	78 (3 dense, 21 with the fused indexer)
TPOT, concurrency 1	6.70 ms	6.29 ms (-6.1%)
TPOT, concurrency 4	11.68 ms	11.34 ms (-2.9%)
NaN probe, 1/2/4/8 concurrent up to 61K tokens	none	none
NIAH 8K/32K/60K, 15 prompts	15/15	15/15
GSM8K 5-shot strict, 250 questions	0.896	0.904

Same image and GPUs for both columns, one run each.

TP8 (8x MI355X), same check: TPOT 5.67 to 5.17 ms (-8.8%) at concurrency 1 and 9.98 to 9.48 ms (-5.0%) at 4; no NaN, NIAH 15/15, GSM8K 0.880 (earlier TP8 runs of this path gave 0.884 to 0.912; 250 questions move by about 3 points).

• pytest tests/models/test_glm_mono_decode_weights.py: 4 passed. It checks that the W_UK/W_UV split keeps the checkpoint's FP8 blocks bit-exact. It is not wired into Buildkite: the path needs gfx950 and the GLM-5 checkpoint, and CI has neither.
• Not run yet: VLLM_GPU_SYNC_CHECK=error. The forward path has no .item(), .cpu() or .tolist() calls; eligibility is decided from host-side metadata.

Server command (TP8)

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 \
VLLM_ROCM_GLM_MONO_DECODE=1 VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4 \
VLLM_ROCM_USE_AITER_SPARSE_INDEXER=1 VLLM_USE_V2_MODEL_RUNNER=1 \
vllm serve amd/GLM-5.3-Quark-MXFP4-AttnFP8 \
  --tensor-parallel-size 8 --kv-cache-dtype fp8 --block-size 64 \
  --max-num-seqs 256 --max-num-batched-tokens 16384 --max-model-len 258048 \
  --no-enable-prefix-caching --async-scheduling --trust-remote-code \
  --compilation-config '{"pass_config": {"fuse_rope_kvcache_cat_mla": true}, "use_inductor_graph_partition": true}' \
  --hf-overrides '{"use_index_cache": true, "index_topk_freq": 4, "index_share_for_mtp_iteration": false}'

The baseline is the same command with VLLM_ROCM_GLM_MONO_DECODE=0.

Details

When the path runs.

• Pure decode steps of 1, 2, 4 or 8 tokens, one token per request. A step of 1 is launched as 2 with a pad row.
• Startup checks, each failing with a reason: gfx950, TP 4 or 8, --kv-cache-dtype fp8, VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1, no speculative decoding, indexer_rope_interleave, and indexer RoPE tables equal to the MLA ones.
• The decision depends only on the token count and the attention metadata, so every rank makes the same one.

Residual stream. The embedding feeds the first launch, and the kernel carries the full residual stream through all 78 layers; the final norm runs on it.

Weights.

• Experts are read in place: AITER-shuffled MXFP4, with the shared expert in the last slot.
• The dense layers' MLP (six times the expert width) is packed once as six expert-shaped MXFP4 slices; the kernel routes every row to all six with weight 1. This needs the row-major MXFP4 layout of the AITER Triton GEMM, the default.
• Block-FP8 attention projections are repacked once, about 4 GB per GPU.
• W_UK/W_UV are cut from kv_b_proj, keeping the checkpoint's FP8 blocks.
• The indexer's wk_weights_proj (bf16) and wq_b (block FP8) are bound as-is.

Indexer.

• On indexer layers the kernel reads and writes the layer's paged FP8 index cache through block_table and req_id_per_token.
• It writes the top-2048 global slots to paged_kv_indices, which the following shared-index layers reuse.
• Layer 2's fused indexer writes the slots that the following shared-index layers read, so nothing is converted outside the kernel.

CUDA graphs.

• Weights bind on the first eligible step, which runs eagerly before capture. Reaching the path inside a capture first raises.
• Pad rows added when a batch is padded to a captured size have slot < 0; the kernel writes no cache entry for them.
• Each kernel runtime advances its epoch once per step. Every layer has its own launch slot within the step, so the step counter is not bumped between layers.

Numerics. Greedy text diverges from the regular path after a few tokens. The kernel quantizes activations to MXFP8 for the MXFP4 experts, where the regular path quantizes them to MXFP4. That is why accuracy is checked with evals rather than token match.

Size. About 545 lines outside tests, almost all in the new mono_decode.py. The only change to an existing file is an 11-line hook in DeepseekV32Model plus one env var. No model runner, scheduler or attention backend changes.

Limitations.

• Only gfx950 and GLM-5 (glm_moe_dsa).
• Concurrency above 8 uses the regular path.
• Needs an AITER build that contains [FlyDSL] GLM-5 fused decode-layer MonoKernel ROCm/aiter#6173.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @amd-sriram.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork