[ROCm][Perf] Skip DSA indexer scoring for short dense-MHA prefills

[whx-sjtu]

Overview

On ROCm, the deepseek_v32 path (DeepSeek-V3.2, GLM-5.x) now skips the sparse indexer's scoring for pure short prefills that run on dense MHA, as CUDA already does (#48407, #51298). Builds on the ROCm dense-MHA prefill from #49263.

Claims

• GLM-5.1-FP8, MI355X, TP4: 1K-token prefill median TTFT -8.1% at c1, -1.6% at c32. gsm8k unchanged.
• GLM-5.2-MXFP4 (indexer on 21 of 78 layers): no significant change.
• Mixed batches, prefills longer than index_topk, MTP layers, FULL cudagraph and capture still score. Other platforms are unchanged.

Validation

• test_mla_short_prefill_indexer.py: 35 passed, 8 skipped (CUDA-only) on MI355X, 7 new. pre-commit clean.
• test_rocm_fused_v32_dense_prefill_matches_mqa (existing, AMD V1 Attention CI): passed, auto and fp8 KV. Its dense run skips one step whose keys a later 2049+-token decode scores.
• Serving, GLM-5.1-FP8, TP4, output 1 token, 2 runs per arm, ABBA, no failed requests:

workload	median TTFT (ms), main → PR	change
1K in, c1	164.5 → 151.2	-8.1%
1K in, c32	1948.6 → 1918.0	-1.6%
1900 in, c1	175.6 → 173.6	within noise
8K in, c1 (sparse path)	620.2 → 623.4	within noise

• gsm8k 5-shot strict, GLM-5.1-FP8: main 0.9469 / 0.9424, PR 0.9447 / 0.9462 (stderr 0.006).
• GLM-5.2-MXFP4, MTP3, 3 runs per arm: TTFT -0.6% to -1.9%, within noise. gsm8k strict 0.9356 → 0.9356, acceptance length 3.31 → 3.31.

Per-run numbers, probes, profile and reproduction

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

GLM-5.1-FP8 (78 indexer layers, no spec decode), median TTFT in ms:

run	arm	1K c1	1900 c1	1K c32	8K c1
F1	main	166.5	175.8	1950.1	622.9
F2	PR	149.2	174.6	1916.9	623.1
F3	PR	153.1	172.7	1919.1	623.6
F4	main	162.5	175.4	1947.0	617.4

gsm8k flexible: main 0.9462 / 0.9409, PR 0.9439 / 0.9484. 1K c32 total throughput 16616 → 17042 tok/s (+2.6%, within noise).

GLM-5.2-MXFP4 (21 indexer layers plus the MTP layer; TTFT runs use MTP3 with synthetic acceptance 3.0), median TTFT in ms, runs in order A B B A A B:

arm	1K c1	1900 c1	1K c32	8K c1
main	152.4 / 154.3 / 152.8	148.9 / 150.3 / 152.2	1501.8 / 1503.7 / 1513.2	471.6 / 475.7 / 473.0
PR	152.0 / 150.6 / 154.1	145.7 / 147.4 / 149.5	1489.8 / 1492.2 / 1492.2	471.8 / 470.5 / 457.9

GLM-5.2 gsm8k with real MTP3: strict 0.9356 → 0.9356, flexible 0.9356 → 0.9363, acceptance length 3.310 → 3.311. At concurrency 1 (250 questions, every prefill a pure short prefill): 0.952 → 0.944 (stderr 0.014). A probe on TP rank 0 counted indexer-layer calls on pure dense-MHA prefills: main scored 8382 of 8382, the PR 382 of 8404, all on the MTP layer.

Why 1K gains more than 1900 (GLM-5.1-FP8, TP rank 0). Prefill runs eagerly, and on main the host needs about 148 ms to issue a step at either length. A probe timed the forward without a profiler (median of 20, ms). Host times move by up to 10 ms between sessions with node load; a first main session under heavy load (1K host issue 187 ms) was repeated:

arm	1K: host issue / GPU tail after it	1900: host issue / GPU tail after it
main	147.6 / 4.4	149.8 / 12.3
PR, run 1	131.3 / 6.9	131.1 / 27.6
PR, run 2	142.0 / 0.4	143.4 / 15.3

At 1K the GPU finishes within a few ms of the last launch, so the step is host-bound and the indexer's host work shows up in TTFT. At 1900 the GPU runs 12-28 ms past the host; host + tail stays at 159-162 ms in every run, and only the indexer's GPU time is saved. A torch profiler trace per step: the PR removes 312 kernels (4 per layer), 1.43 ms of GPU time at 1K and 2.17 ms at 1900; on main the indexer op costs 190-198 us of host time but 24-34 us of GPU time per call (host time inflated by the profiler).

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_FP8BMM=0 VLLM_ROCM_USE_AITER_FP4BMM=0 \
vllm serve zai-org/GLM-5.1-FP8 --tensor-parallel-size 4 --kv-cache-dtype fp8 \
  --no-enable-prefix-caching --max-model-len 16384 --gpu-memory-utilization 0.9
# GLM-5.2: amd/GLM-5.2-MXFP4, same flags plus
#   --speculative-config '{"method":"mtp","num_speculative_tokens":3}'

# 1K c1 shown. 1900 c1: --random-input-len 1900 --seed 45.
# 1K c32: --num-prompts 512 --max-concurrency 32 --seed 46.
# 8K c1: --random-input-len 8192 --num-prompts 16 --seed 47.
vllm bench serve --backend vllm --model zai-org/GLM-5.1-FP8 --dataset-name random \
  --random-input-len 1024 --random-output-len 1 --temperature 0 \
  --num-prompts 40 --max-concurrency 1 --seed 44

lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --batch_size 1 \
  --model_args model=zai-org/GLM-5.1-FP8,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=64,tokenized_requests=False

VLLM_ROCM_USE_AITER=1 pytest -v tests/model_executor/layers/test_mla_short_prefill_indexer.py \
  "tests/v1/attention/test_sparse_mla_backends.py::test_rocm_fused_v32_dense_prefill_matches_mqa"

Details

• Placement: DeepseekV32MLAAttention._fused_attention, which already holds the MLA metadata. CUDA passes the layer name into its op only because the op cannot see that metadata.
• Conditions: the CUDA op's, plus the layer taking the dense-MHA branch. The flag in DeepseekV32Attention moves from is_cuda() to is_cuda_alike().
• K cache: still written by fused_norm_rope; only scoring is dropped.
• The gain is host time: an eager 1K prefill is host-bound and each indexer call costs about 0.2 ms on the host; at 1900 the step is GPU-bound and the indexer kernels take about 2 ms per step.
• Not measured: DeepSeek-V3.2 (same code path), GLM-5.1 with MTP, TP8.

Design notes

• Moving the check into rocm_aiter_sparse_attn_indexer would mirror CUDA but add an argument to an op that DeepSeek-V4 shares; the model already has the MLA metadata and the dense-MHA decision.
• The ROCm indexer op is built with skip_k_cache_insert=True. fused_norm_rope writes the indexer keys and clears the top-k buffer before the decision.
• PCP leaves the flag off, as on CUDA; ROCm sparse MLA rejects DCP. MTP layers keep scoring, as on CUDA.
• This model is not torch.compiled. Prefills run eagerly under the default FULL_DECODE_ONLY; under breakable cudagraphs the eager segment re-evaluates the decision on replay. Profile and dummy runs (no metadata) still call the op, so its workspace is reserved.

Not a duplicate: no open or closed PR skips indexer scoring on ROCm.

Duplicate search

gh pr list --repo vllm-project/vllm --state all --search with "short prefill indexer ROCm", "skip indexer ROCm", "indexer scoring ROCm", "dense mha indexer ROCm", "short prefill topk ROCm", "rocm_aiter_sparse_attn_indexer", "dense_mha_metadata_layer_name", "short prefill skip indexer", "GLM ROCm indexer prefill", "ROCm DSv3.2 indexer". Nearest: #56944 and #60601 (shard indexer prefill rows across TP), #61504 (sparse MLA decode indptr), #60294 (aiter mqa_logits for indexer prefill); none skips scoring.

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.

[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.