[ROCm][Perf] Use AITER top-k for the DeepSeek-V3.2/GLM-5 sparse indexer by whx-sjtu · Pull Request #61943 · vllm-project/vllm

whx-sjtu · 2026-10-16T02:38:37Z

Overview

On gfx950 with AITER enabled, the uncompressed sparse indexer of DeepSeek-V3.2 and GLM-5.x (index_topk 2048) now takes AITER's top-k for decode and for prefill chunks whose longest row reaches 32K tokens, through the adapters #52882 added for DeepSeek-V4. Related: #55327, #50470, #46172, #60012.

Claims

GLM-5.2-MXFP4, MI355X, TP4, --max-model-len 133120: 128K/1K TPOT -7.0% / -4.6% / -3.0% at c1 / c8 / c16, MTP3 c8 -3.0%; 8K/1K within noise; gsm8k and passkey unchanged.
DeepSeek-V3.2, TP4, default max_model_len (163840): 128K/1K c1 TPOT -15.6%, 4K/1K c16 -1.1%, 4K/1K c1 +0.7% (+0.085 ms); gsm8k unchanged.
Compressed indexers, other index_topk, GLM-5.3-Flash kpool, gfx942, AITER off and non-ROCm are unchanged.

Validation

test_top_k_per_row.py, MI355X: 43 passed (8 new), 246 skipped, 1 failed as on main (pre-existing, see below). glm5next dispatch tests pass. pre-commit clean.
AITER vs in-tree on FP8 scorer output: same selections on 29,934 rows (by value where scores tie).
Serving, mean TPOT (ms), GLM-5.2-MXFP4 unless noted; main 2 runs / PR 3 runs (4 each with MTP3, 2 each for DeepSeek-V3.2), no failed requests:

workload	main	PR	change
128K/1K c1 / c8 / c16	14.73 / 45.05 / 78.94	13.70 / 42.97 / 76.60	-7.0% / -4.6% / -3.0%
MTP3 128K/1K c8	34.84	33.79	-3.0%
8K/1K c1 / c8 / c16, MTP3 c8	13.48 / 17.37 / 22.66 / 9.06	13.45 / 17.39 / 22.54 / 9.16	within noise
DeepSeek-V3.2 128K/1K c1	16.23	13.70	-15.6%
DeepSeek-V3.2 4K/1K c1 / c16	12.81 / 22.86	12.90 / 22.61	+0.7% / -1.1%

gsm8k 5-shot strict: GLM-5.2 0.9287 main, 0.9348 PR (stderr 0.007); DeepSeek-V3.2 0.9522 main, 0.9507 PR (stderr 0.006). MTP3 acceptance length 3.29 / 3.30.

Per-run numbers, kernel data and reproduction

Image vllm/vllm-openai-rocm:nightly (AITER 0.1.25) running main at 7f730b4 and this branch from source, MI355X.

GLM-5.2-MXFP4, mean TPOT per run (ms). Order: main, PR, PR, main, PR. A further main run between the two PR sessions was dropped: another process held about 88 GiB per GPU while it started, so its KV cache was 1.23M tokens instead of 3.2M. MTP3 order: main, PR, PR, main, PR, PR, main, main.

workload	main	PR
8K/1K c1	13.43, 13.53	13.44, 13.45, 13.46
8K/1K c8	17.27, 17.47	17.27, 17.51, 17.38
8K/1K c16	22.46, 22.86	22.44, 22.66, 22.52
128K/1K c1	14.72, 14.75	13.66, 13.67, 13.78
128K/1K c8	44.55, 45.56	42.95, 42.98, 42.98
128K/1K c16	78.89, 79.00	76.58, 76.62, 76.60
MTP3 8K/1K c8	8.91, 9.16, 9.08, 9.10	9.14, 9.34, 9.08, 9.10
MTP3 128K/1K c8	34.27, 35.52, 34.45, 35.12	33.96, 33.67, 34.17, 33.35

Median ITL (ms): 128K c1 14.83 to 13.76, c8 18.22 to 17.01, c16 21.52 to 20.49, MTP3 128K c8 27.04 to 25.39. Output throughput at 128K: +5.8% / +3.7% / +2.8%. KV cache size and startup time match across arms.
The 128K c1 ITL drop (1.07 ms) matches the kernel estimate of 21 indexer layers x 53 us (1.11 ms).
gsm8k flexible: 0.9295 main, 0.9363 PR. MTP3 gsm8k, first 400 questions: 0.935 main, 0.9325 PR (stderr 0.012). Passkey at 32K / 64K / 120K: 12/12, 8/8, 6/6 in both arms.
Operator check: decode logits from AITER's deepgemm_fp8_paged_mqa_logits (buffer poisoned past each row), prefill logits from rocm_fp8_mqa_logits, calls mirroring rocm_aiter_sparse_attn_indexer; widths 133120 and 1M, next_n 1 (1-D), 2 (native 2-D) and 4 (one row per token), packed multi-sequence prefill, and graph replay with new lengths. Every row has min(L, k) distinct indices in [0, L) followed by -1.

DeepSeek-V3.2 (FP8 checkpoint), TP4, same env and flags but no --max-model-len (163840), mean TPOT per run (ms), order main, PR, PR, main; median ITL and output throughput follow TPOT:

workload	main	PR
4K/1K c1	12.83, 12.80	12.91, 12.89
4K/1K c16	22.84, 22.88	22.62, 22.60
128K/1K c1	16.24, 16.22	13.72, 13.69

4K/1K c16 mean TTFT: 2333, 2326 ms main, 2465, 2327 ms PR. 128K/1K c1 output throughput +10.0%. KV cache 2.07M tokens in every session. gsm8k flexible: 0.9530 main, 0.9507 PR.
A static cut-off for 131072-200000 columns would put this deployment back on the in-tree kernel: it would recover the 0.085 ms at 4K c1 and lose the 128K gain (2.53 ms) and the 4K c16 gain. GLM-5.2 at 133120 would lose its 128K gains as well.

Decode top-k per indexer layer (us, in-tree to AITER), GLM-5.2 indexer (32 heads, k=2048); scoring and top-k for 21 layers captured in one graph, LLC flushed before each replay, median of 20:

width	context	1 row	8 rows	16 rows
133120	4K	17.5 to 20.4	17.3 to 22.3	17.2 to 22.7
133120	32K	34.0 to 22.2	34.9 to 22.8	34.7 to 23.5
133120	128K	80.8 to 27.9	81.8 to 28.8	81.3 to 28.5
1M	4K	37.0 to 25.9	35.3 to 29.3	35.3 to 35.3
1M	32K	42.8 to 26.8	42.3 to 29.6	43.7 to 32.3
1M	128K	46.3 to 27.1	50.0 to 31.2	50.2 to 34.4

MTP1 (2-D bounds) and MTP3 (one row per token) follow the same pattern. At 262144 columns AITER wins at every measured context.

Prefill top-k GPU time per 16K-row chunk (us, in-tree to AITER): context 16K 304 to 370, 20K 426 to 469, 24K 568 to 558, 32K 757 to 680, 64K 1681 to 1320, 128K 4026 to 2635. 512-row chunks cross over near 20K. Host enqueue per call: in-tree 5-12 us, AITER 27-36 us.

VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_ROCM_USE_AITER_MLA=1 \
VLLM_ROCM_USE_AITER_RMSNORM=1 VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0 \
VLLM_ROCM_USE_AITER_FP8BMM=0 VLLM_ROCM_USE_AITER_FP4BMM=0 \
vllm serve amd/GLM-5.2-MXFP4 --tensor-parallel-size 4 --kv-cache-dtype fp8 \
  --block-size 64 --no-enable-prefix-caching --max-model-len 133120 \
  --gpu-memory-utilization 0.9
# MTP3: --speculative-config '{"method":"mtp","num_speculative_tokens":3}'
# DeepSeek-V3.2: deepseek-ai/DeepSeek-V3.2, same env and flags without --max-model-len.

# 128K/1K: num-prompts max(C, 3). 8K/1K: --random-input-len 8192,
# num-prompts max(3C, 6), --seed 42. MTP3 c8: 8 prompts at 128K, 24 at 8K.
# DeepSeek-V3.2 4K/1K: --random-input-len 4096, 6 prompts at c1, 48 at c16, --seed 42.
vllm bench serve --backend vllm --model amd/GLM-5.2-MXFP4 --dataset-name random \
  --random-input-len 131072 --random-output-len 1024 --ignore-eos --temperature 0 \
  --num-prompts $N --max-concurrency $C --seed 43

lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --batch_size 1 \
  --model_args model=amd/GLM-5.2-MXFP4,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=64,tokenized_requests=False

VLLM_ROCM_USE_AITER=1 pytest -v tests/kernels/test_top_k_per_row.py \
  tests/models/glm5next/test_sparse_indexer_topk_dispatch.py

Details

Gate: compress ratio 1 and k=2048 only. Decode always takes AITER (FULL graphs are not keyed by context length); prefill when chunk.max_seqlen_k, a host int, reaches 32K.
Cost: with max_model_len between 131072 and 200000 and up to 48 rows, AITER loses 3-6 us per layer below 8-16K context. DeepSeek-V3.2 (default 163840, 61 indexer layers) measures +0.085 ms TPOT at 4K/1K c1. Keeping the in-tree kernel for that width band would also give up the 2.53 ms gain at 128K on the same deployment, so the gate does not exclude it.
Not measured end to end: GLM-5/5.1, TP8. CI (MI300, NVIDIA) runs only the new selection test; the GPU tests need gfx950.

Design notes, pre-existing failure and related PRs

Cause: the in-tree decode kernel has gfx950 policies only for k=512 and k=1024 ([ROCm] gfx950 top-k policy is gated on topK==1024, leaving index_topk=2048 models (GLM-5.x) on the generic 10-way split #55327). At k=2048 its cost grows with context; AITER's stays near flat.
Semantics come from [ROCm][Perf] Optimize DeepSeek V4 C4A top-k with AITER #52882's wrappers. ROCm DSA decode carries (B, next_n) bounds for native MTP1 and one row per token otherwise, so seq_lens.reshape(-1) with next_n=1 gives AITER each row's exact end. Prefill turns AITER's chunk-global indices into sequence-local ones and keeps -1.
Decode is not split by context length because a FULL graph replays one kernel for every length. Eager decode could use the length, but would then pick a different kernel for the same shape depending on graph mode. AITER's extra host time per call only shows in eager steps.
test_aiter_decode_topk_matches_per_row_under_mtp (from [ROCm][Perf][GLM-5.3-Flash] Add AITER topk backend for decodes #58167, requires_gfx950, so not in CI) fails identically on main: it passes (B, 1) seq_lens with next_n=4 to the in-tree k=512 gfx950 kernel, which reads any 2-D seq_lens per row, so its reference rows are wrong. It does not reach the code changed here.
[ROCm][Perf] gfx950 device length aware top-k split policy for k=2048 #60012 (open) tunes the in-tree k=2048 kernel above 200000 columns for [ROCm] gfx950 top-k policy is gated on topK==1024, leaving index_topk=2048 models (GLM-5.x) on the generic 10-way split #55327. With this PR it still serves AITER-off deployments and short prefill chunks; the gate should be re-measured if it lands. [ROCm] Calibrate gfx950 decode Top-K routing by rows × compressed-context #57844 (open) edits the same selector. [ROCm][Perf] Allow DSA indexer native decode above next_n=2 #59421 (draft, native MTP above next_n=2) is compatible because decode flattens the bounds.
[ROCm] Route sparse-indexer decode top-k through AITER dispatcher #50470 (draft) and [ROCm] [Perf] Optimize rocm aiter sparse mla indexer #46172 also move the k=2048 decode top-k to AITER. Both have merge conflicts and no activity since 2026-08-22, and both swap the call unconditionally, passing vLLM's next_n with its 2-D MTP bounds. On both, the author of [ROCm][Perf] Optimize DeepSeek V4 C4A top-k with AITER #52882 noted that AITER v0.1.19 needs those bounds adapted and asked to treat [ROCm][Perf] Optimize DeepSeek V4 C4A top-k with AITER #52882 as the consolidated direction; [ROCm][Perf] Optimize DeepSeek V4 C4A top-k with AITER #52882 then merged with AITER enabled for compressed indexers only. This PR applies its adapters to the k=2048 path, adds prefill and serving results, and leaves every other indexer unchanged. The aiter.top_k_per_row_decode entry point that [ROCm] Route sparse-indexer decode top-k through AITER dispatcher #50470 waited for ships in AITER 0.1.25 and is the function called here.
Duplicate search, gh pr list --state all --search: "aiter topk", "top-k ROCm", "topk ROCm", "indexer ROCm", "indexer aiter", "DeepSeek-V3.2 ROCm", "GLM-5 ROCm", "sparse indexer AMD", "top_k_per_row_decode", "top_k_per_row_prefill", "compress_ratio aiter topk".

Not a duplicate: #50470 and #46172 target the same decode call, have been inactive since 2026-08-22, and were marked superseded by #52882's author; #52882 left k=2048 on the in-tree kernel (details above). Thanks @JH-Leon-KIM-AMD for #50470.

AI assistance (Claude) was used for the implementation, the measurements and this description.

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.

The ROCm sparse indexer only routes compressed indexers (compress_ratio > 1, e.g. DeepSeek-V4 C4A) to AITER's top-k. The uncompressed DSA indexer of DeepSeek-V3.2 and GLM-5 (k=2048) stays on the in-tree kernel, which has no gfx950 path for k=2048. On MI355X with GLM-5.2's indexer (32 heads, k=2048), AITER cuts decode top-k from 81 to 28 us per layer at 128K context. It wins or ties every measured decode shape except short contexts when max_model_len is in (128K, 200K], where it costs up to ~6 us per layer. Decode graphs cannot see the context length, so decode always takes AITER. Prefill runs eagerly and takes AITER once the longest row of a chunk reaches 32K tokens (e.g. 4.0 -> 2.6 ms per 16K-row chunk at 128K); shorter chunks, where AITER is slower, stay on the in-tree kernel. GLM-5.2-MXFP4, TP4, FP8 KV: 128K/1K TPOT -7.0%/-4.6%/-3.0% at c1/c8/c16, MTP3 128K c8 -3.0%, 8K/1K within noise; gsm8k and passkey unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: whx-sjtu <56632993+whx-sjtu@users.noreply.github.com>

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

Fangzhou-Ai · 2026-10-16T19:52:07Z

/ci run

github-actions · 2026-10-16T19:52:18Z

✅ @whx-sjtu, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

github-actions · 2026-10-16T19:52:34Z

✅ Triggered Buildkite CI #95466 for commit 4e25d4038339.

ℹ️ This PR is 94 commits behind upstream main. Merged upstream main into this branch as 4e25d4038339 before starting CI. Pull the branch before pushing more changes.

whx-sjtu · 2026-10-17T02:05:37Z

/ci run

github-actions · 2026-10-17T02:05:57Z

✅ Triggered Buildkite CI #95519 for commit bcf74fd00abf.

whx-sjtu · 2026-10-19T06:53:59Z

/ci run

github-actions · 2026-10-19T06:55:57Z

✅ Triggered Buildkite CI #95846 for commit b352f023b758.

Let `--attention-config '{"indexer_kv_dtype": "mxfp4"}'` serve the DeepSeek-V3.2-style DSA models (GLM-5.x, DeepSeek-V3.2 and their MTP drafts) on gfx950, reusing the paged MXFP4 indexer that DeepSeek-V4.1 already runs on: the model writes the indexer K cache in the page order aiter's paged_mxfp4_mqa_logits reads, and the existing RocmSparseAttnIndexer / RocmMxfp4IndexerMetadataBuilder score it in place for prefill and decode, with the per-step decode schedule and native next_n rows for MTP. - fused_norm_rope gains an MXFP4 cache write: the indexer key, after its LayerNorm and RoPE, is quantized per 32 values (2^ceil(log2(amax/6)) e8m0 scales) and stored values-then-scales in RocmPagedMxfp4CacheLayout order. - fused_q can emit the index query as MXFP4 with its e8m0 bytes, leaving the query scale out of the weights. - The ROCm indexer cache sizes its row as 68 bytes and selects the new DeepseekV32RocmMxfp4IndexerBackend (64/128-token pages; context parallelism is rejected at startup). FP8 stays the default. A token's indexer key shrinks from 132 to 68 bytes. MI355X, TP4, FP8 KV, against the FP8 indexer with AITER top-k (vllm-project#61943): GLM-5.2-MXFP4 KV cache +3.0% tokens; random 128K/1K TPOT -5.4% (c8), -5.9% (c16), MTP3 c8 -6.8%, TTFT -7% to -8%; 128K c1 and 8K/1K within noise; GSM8K 20-shot 0.9325 vs 0.9272, passkey unchanged. DeepSeek-V3.2: KV cache +10.2% tokens, 128K/1K c8 TPOT -12.5%, c1 and 4K/1K c16 within noise, GSM8K 5-shot 0.9575 vs 0.9515. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: whx-sjtu <56632993+whx-sjtu@users.noreply.github.com>

…er (vllm-project#61943) Signed-off-by: whx-sjtu <56632993+whx-sjtu@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: vllm-agent <claw@inferact.ai>

mergify Bot added deepseek Related to DeepSeek models glm rocm Related to AMD ROCm labels Oct 16, 2026

github-project-automation Bot added this to AMD Oct 16, 2026

github-project-automation Bot moved this to Todo in AMD Oct 16, 2026

whx-sjtu marked this pull request as ready for review October 16, 2026 02:53

whx-sjtu requested review from AndreasKaratzas, WoosukKwon, mgoin, tjtanaa, tlrmchlsmth, yewentao256 and zyongye as code owners October 16, 2026 02:54

claude Bot reviewed Oct 16, 2026

View reviewed changes

whx-sjtu mentioned this pull request Oct 16, 2026

[ROCm][Perf] MXFP4 indexer K cache for GLM-5 / DeepSeek-V3.2 DSA #61985

Open

6 tasks done

Fangzhou-Ai approved these changes Oct 16, 2026

View reviewed changes

Fangzhou-Ai added the ready ONLY add when PR is ready to merge/full CI is needed label Oct 16, 2026

Merge branch 'main' into rocm-glm-fp8-indexer-aiter-topk

4e25d40

Merge branch 'main' into rocm-glm-fp8-indexer-aiter-topk

bcf74fd

Merge branch 'main' into rocm-glm-fp8-indexer-aiter-topk

b352f02

shen-shanshan approved these changes Oct 19, 2026

View reviewed changes

shen-shanshan merged commit fb61ac0 into vllm-project:main Oct 19, 2026
185 checks passed

github-project-automation Bot moved this from Todo to Done in AMD Oct 19, 2026

Navigation Menu

Repository navigation

[ROCm][Perf] Use AITER top-k for the DeepSeek-V3.2/GLM-5 sparse indexer - #61943

[ROCm][Perf] Use AITER top-k for the DeepSeek-V3.2/GLM-5 sparse indexer#61943
shen-shanshan merged 4 commits into
vllm-project:mainfrom
whx-sjtu:rocm-glm-fp8-indexer-aiter-topk

whx-sjtu commented Oct 16, 2026 •

edited by github-actions Bot

Loading

Uh oh!

claude Bot left a comment

Uh oh!

Fangzhou-Ai commented Oct 16, 2026

Uh oh!

github-actions Bot commented Oct 16, 2026

Uh oh!

github-actions Bot commented Oct 16, 2026

Uh oh!

whx-sjtu commented Oct 17, 2026

Uh oh!

github-actions Bot commented Oct 17, 2026

Uh oh!

whx-sjtu commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

whx-sjtu commented Oct 16, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Claims

Validation

Details

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

Fangzhou-Ai commented Oct 16, 2026

Uh oh!

github-actions Bot commented Oct 16, 2026

Uh oh!

github-actions Bot commented Oct 16, 2026

Uh oh!

whx-sjtu commented Oct 17, 2026

Uh oh!

github-actions Bot commented Oct 17, 2026

Uh oh!

whx-sjtu commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

whx-sjtu commented Oct 16, 2026 •

edited by github-actions Bot

Loading