[MI355X][ROCm][AMD] Fuse the DSA indexer QK prologue for GLM-5.x

[phucnguyen-ht]

Purpose

Every DeepSeek-Sparse-Attention layer (DeepSeek-V3.2, GLM-5.x, and their MTP drafts) runs a lightweight indexer whose pre-processing costs five kernel launches per layer per step:

1. LayerNorm on k
2. RoPE on the leading rope dims of q and k
3. per-token-group FP8 quantization of q
4. a pointwise fold of the q scale into the indexer weights
5. FP8 quantization of k plus the paged indexer K-cache write

At decode sizes this is launch-bound: the cost barely moves with token count. AITER already ships a kernel that does all five in one launch (indexer_qk_rope_quant_and_cache). This PR wires it in behind a new, default-off env flag VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION.

The kernel writes the indexer K cache itself, so the sparse-attention indexer op skips its own insert (skip_k_cache_insert). Two data-movement optimizations ride along: the rope cos_sin_cache halves are registered once as strided views on the indexer's rotary embedding instead of being re-split every layer every step, and one zero-initialized q_fp8/weights output pair is shared by the indexer layers of a model (allocated next to topk_indices_buffer and sized the same way), with only the first indexer of a forward pass zero-filling it. The change is Python-only: the kernel is already in the AITER version the tree pins.

The fused path engages only when all of the following hold, and otherwise falls back silently to today's code:

• VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION=1
• gfx942 / gfx950 (on_mi3xx()): AITER ships this kernel in its CK build, and the RDNA paths that SparseAttnIndexer.forward_hip now also serves have Triton-only AITER
• the in-place-RoPE indexer path, i.e. the rotary_embedding custom op is enabled
• DSA indexer shapes index_head_dim == 128 and qk_rope_head_dim == 64
• no context parallelism (decode_context_parallel_size == 1 and prefill_context_parallel_size == 1). The fused kernel is driven by slot_mapping, which is PAD_SLOT_ID on ranks that do not own a token, so those ranks would skip the row and never produce its quantized query — while every CP rank still needs the query to score its own KV shard.

Test Plan

• Image: vllm/vllm-openai-rocm:nightly-5a4c8d99242e9e069b604d0e9b969e77f7dd501d.
• 8× MI355X (gfx950) — GLM-5.2-FP8 and GLM-5.2-MXFP4: op test, kernel benchmark, coherence check, GSM8K (speculative decoding off and on), GPQA-Diamond, RULER/NIAH, serving throughput.
• Serve configuration: the official recipe for GLM-5.2, AMD selection (https://recipes.vllm.ai/zai-org/GLM-5.2), plus --no-enable-prefix-caching for the throughput sweep. Accuracy leaves prefix caching at its default, since GSM8K's few-shot prompts share prefixes; both arms are identical per workload.

Serve — GLM-5.2-FP8.

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION=1   # the flag this PR adds

vllm serve /models/GLM-5.2-FP8 --port $PORT \
  --kv-cache-dtype fp8_e4m3 --tensor-parallel-size 8 \
  --linear-backend aiter --moe-backend aiter \
  --tool-call-parser glm47 --enable-auto-tool-choice --reasoning-parser glm45 \
  --no-enable-prefix-caching        # throughput sweep only; accuracy keeps the default

Serve — GLM-5.2-MXFP4.

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_USE_AITER_INDEXER_QK_FUSION=1
export VLLM_USE_V2_MODEL_RUNNER=1
export VLLM_ROCM_USE_AITER_FP8BMM=0
export VLLM_ROCM_USE_AITER_FP4BMM=0

vllm serve /models/GLM-5.2-MXFP4 --port $PORT \
  --kv-cache-dtype fp8_e4m3 --tensor-parallel-size 8 \
  --linear-backend aiter --moe-backend aiter \
  --tool-call-parser glm47 --enable-auto-tool-choice --reasoning-parser glm45 \
  --no-enable-prefix-caching --trust-remote-code

Speculative decoding. Append this to either serve command:

--speculative-config '{"method":"mtp","num_speculative_tokens":5}'

Op test

pytest tests/kernels/attention/test_rocm_aiter_indexer_qk_fusion.py -v

Kernel benchmark

python benchmarks/kernels/benchmark_indexer_qk_fusion.py --repeat 5

Accuracy — coherence check. One chat completion, to confirm the model is not producing garbage before spending hours on the rest.

curl -s http://127.0.0.1:$PORT/v1/chat/completions -H 'Content-Type: application/json' -d '{
  "model": "<model>", "messages": [{"role": "user", "content": "What is 17 times 23?"}],
  "temperature": 0.6, "top_p": 0.95, "max_tokens": 8192}'

Accuracy — GSM8K

lm_eval --model local-completions --tasks gsm8k \
  --model_args "model=<model>,base_url=http://127.0.0.1:$PORT/v1/completions,num_concurrent=64,max_retries=3,tokenized_requests=False,timeout=600"

Accuracy — GPQA-Diamond, 198 questions, zero-shot, against the running server:

python3 -m sglang.test.run_eval --eval-name gpqa --model <model> \
  --host 127.0.0.1 --port $PORT --num-threads 32 --max-tokens 100000 \
  --temperature 1.0 --top-p 0.95 \
  --chat-template-kwargs '{"enable_thinking": true, "think_end_token": "</think>"}'

Accuracy — RULER / NIAH, 500 samples per context length, offline engine:

lm_eval --model vllm \
  --model_args 'pretrained=<model>,tensor_parallel_size=8,max_model_len=262144,gpu_memory_utilization=0.85,kv_cache_dtype=fp8,trust_remote_code=True,enable_thinking=True,think_end_token=</think>,max_gen_toks=8192,block_size=64' \
  --tasks niah_single_2 --metadata '{"max_seq_lengths":[65536,131072]}' \
  --batch_size auto --seed 1234 --log_samples --output_path ./niah

Serving throughput, for CONC in 4 8 16 32 64 128 256:

vllm bench serve --port $PORT --model <model> --dataset-name random \
  --random-input-len 8192 --random-output-len 1024 \
  --num-prompts $((3*CONC)) --max-concurrency $CONC \
  --num-warmups 8 --request-rate inf --ignore-eos \
  --percentile-metrics ttft,tpot,itl,e2el

Lint

pre-commit run --files $(git diff --name-only HEAD~1 HEAD)

The op test scores the fused op and the unfused flow against an fp64 golden of the same math over 20 cases (num_tokens ∈ {1,7,32,257,1023} × block_size ∈ {1,64} × is_neox ∈ {True,False}): the fused path must be no less accurate than the unfused one, and the two must agree to within one fp8 code on every element above 1e-3 of the tensor maximum, for both q_fp8 * weights_out and the dequantized indexer K cache. Two further cases pin the rows the kernel skips to zero and opcheck the op's schema, for 22 in total.

Test Result

Op-level equivalence — 22 passed on 8× MI355X (gfx950).

Kernel speed — one indexer layer's pre-processing, median of --repeat 5 at block_size=64, n_head=32 (index_n_heads for both GLM-5.2 checkpoints) and the interleaved RoPE layout GLM-5.x selects, on 8× MI355X (gfx950):

tokens	unfused	fused	speedup
1	70.29 µs	11.14 µs	6.31×
8	74.10 µs	11.31 µs	6.55×
32	74.17 µs	11.16 µs	6.65×
64	74.51 µs	11.17 µs	6.67×
256	74.93 µs	11.36 µs	6.59×
1024	73.00 µs	25.06 µs	2.91×

pre-commit — all hooks pass (ruff check, ruff format, typos, markdownlint, mypy 3.10, SPDX headers, root lazy imports, forbidden imports, torch.cuda API check, configuration validation, …), exit code 0.

Coherence check — both arms answer 391 with a non-empty content, i.e. the model closed its thinking block and answered:

Arm	reasoning	content	Answer
Before	454 chars	23 chars	17 times 23 is **391**.
After	575 chars	19 chars	17 times 23 is 391.

GSM8K (1319 questions, lm_eval gsm8k)

8× MI355X, TP8. Reported as strict-match with lm_eval's standard error, flexible-extract in brackets.

Model	MTP	Before	After
GLM-5.2-FP8	off	94.54% ±0.63 (94.54%)	94.16% ±0.65 (94.31%)
GLM-5.2-FP8	5 tokens	94.39% ±0.63 (94.39%)	94.39% ±0.63 (94.39%)
GLM-5.2-MXFP4	off	93.10% ±0.70 (93.10%)	93.40% ±0.68 (93.40%)
GLM-5.2-MXFP4	5 tokens	94.01% ±0.65 (94.01%)	93.63% ±0.67 (93.48%)

GPQA-Diamond (198 questions, zero-shot)

Model	Before	After
GLM-5.2-MXFP4	89.40%	91.90%

RULER / NIAH (niah_single_2, 500 samples per length)

Model	Context	Before	After
GLM-5.2-MXFP4	65,536	100% (500/500)	100% (500/500)
GLM-5.2-MXFP4	131,072	100% (500/500)	100% (500/500)

Serving throughput, ISL 8192 / OSL 1024

GLM-5.2-FP8, MTP off, 8× MI355X TP8:

Conc	Out tok/s before	Out tok/s after	Δ	Mean TPOT ms
4	154.6	157.2	+1.66%	24.8 → 24.3
8	273.2	278.3	+1.84%	27.4 → 26.8
16	461.2	469.9	+1.88%	31.5 → 30.9
32	700.2	705.2	+0.72%	41.0 → 40.8
64	989.9	995.2	+0.54%	57.9 → 57.1
128	1308.4	1311.7	+0.25%	85.8 → 86.0
256	1547.1	1547.8	+0.05%	144.2 → 144.0

GLM-5.2-MXFP4, MTP off, 8× MI355X TP8:

Conc	Out tok/s before	Out tok/s after	Δ	Mean TPOT ms
4	187.5	193.9	+3.41%	20.3 → 19.7
8	324.8	330.3	+1.71%	23.0 → 22.7
16	562.4	564.5	+0.37%	26.0 → 25.9
32	853.8	864.1	+1.21%	33.9 → 33.4
64	1249.8	1259.7	+0.79%	45.7 → 45.6
128	1642.5	1650.0	+0.46%	69.0 → 68.7
256	1975.0	1980.8	+0.30%	113.1 → 112.5

Note: This change includes AI-generated code.