[ROCm][Perf] Enable fused indexer-Q RoPE+quant kernel for DeepSeek/GLM sparse attention

[xaguilar-amd]

Summary

Enables the existing fused indexer-Q kernel (fused_indexer_q_rope_quant) on
ROCm for the DeepSeek Sparse Attention (DSA) indexer used by GLM-5.2 /
DeepSeek-V3.2-style models. On CUDA this kernel already collapses the indexer's
query-side RoPE + FP8 quantization + weight-scale fold into a single launch; on
ROCm the same work was falling back to a chain of small Triton/elementwise
kernels per sparse layer per decode token.

This is a low-risk, 2-line enablement — the kernel is already on vllm and is
architecture-aware. It removes redundant kernels and HBM round-trips on the
ROCm decode path (4 Q-side kernels → 1), yielding a small but consistent decode
improvement with no CUDA-side change and no measured regression.

Kernel-level evidence (torch.profiler, GLM-5.2-FP8, MI325X / gfx942, TP8, rank 0)

Before/after self-CUDA over the same steady-state decode capture. The indexer
decode kernels are the ones with ~21.5k calls (per decoded token × sparse
layers).

Fused away (present in baseline, gone after):

Kernel	Role	Self CUDA	Calls
per_token_group_quant_8bit_kernel	q → FP8 (ue8m0) quant	99.8 ms	21546
triton_poi_fused_3	RoPE/cat elementwise	95.8 ms	21526
triton_poi_fused_2	RoPE/cat elementwise	89.0 ms	20521
triton_poi_fused_mul_slice_unsqueeze_view_4	weights scale-fold	85.3 ms	20500

Introduced:

Kernel	Role	Self CUDA	Calls
_fused_indexer_q_rope_quant_kernel	fused q RoPE + FP8 quant + weight-fold (decode)	103.7 ms	21546
_fused_indexer_q_rope_quant_kernel_0	same, prefill variant	6.3 ms	42
triton_poi_fused_add_copy_index_select_mul_slice_…	residual K-side RoPE (now its own kernel)	90.6 ms	20500

The Q-side collapses from 4 dedicated kernels → 1 (the two triton_poi_fused_2/_3
RoPE kernels did q and k together; after the change q's RoPE is absorbed and
only a single K-RoPE kernel remains). Net ≈ 170 ms self-CUDA
reclaimed on rank 0 over the capture (≈0.7% of the ~24.9 s total). Untouched:
k_norm, indexer_k_quant_and_cache, _gluon_deepgemm_fp8_paged_mqa_logits,
topKPerRowDecode.

Accuracy

GLM-5.2-FP8, TP=8 on gfx942

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	5	exact_match	↑	0.9469	±	0.0062
		strict-match	5	exact_match	↑	0.9469	±	0.0062

End-to-end performance (GLM-5.2-FP8, TP8, MI325X / gfx942)

vllm bench serve, random dataset, --ignore-eos, --random-range-ratio 0.0.

Decode runs under CUDA graphs.

ISL	OSL	conc	TPOT_p50 base	TPOT_p50 fused	TPOT	ITL_p50 base	ITL_p50 fused	ITL	tok/s base	tok/s fused	tput
1024	1024	1	14.98	14.93	1.004x	14.98	14.94	1.003x	66.32	66.34	1.000x
1024	1024	8	18.39	18.32	1.004x	18.40	18.30	1.005x	422.81	423.68	1.002x
1024	1024	32	24.83	24.71	1.005x	24.31	24.19	1.005x	1220.94	1222.21	1.001x
1024	1024	64	32.39	32.30	1.003x	31.25	31.00	1.008x	1835.24	1842.09	1.004x
8192	1024	1	16.01	15.90	1.006x	15.98	15.88	1.006x	60.18	60.55	1.006x
8192	1024	8	22.26	22.06	1.009x	19.86	19.79	1.004x	323.27	324.17	1.003x
8192	1024	32	38.86	39.73	0.978x*	25.14	24.98	1.006x	715.20	718.98	1.005x
8192	1024	64	63.69	63.66	1.001x	29.72	29.80	0.997x	925.56	925.84	1.000x

Consistent small improvement across the sweep (~0.3–0.9%), with no credible
regression. *The single 8192/32 TPOT dip is run-to-run variance: ITL (1.006x)
and throughput (1.005x) both improve at that same point.

Related work / positioning

• [ROCm][Perf] DSv3.2: fuse indexer Q-RoPE+quant + K-norm/RoPE/quant/cache #43907 ([ROCm][Perf] DSv3.2: fuse indexer Q-RoPE+quant + K-norm/RoPE/quant/cache)
  fuses Q and K via an AITER kernel, but explicitly excludes GLM
  (model_type != "glm_moe_dsa") and depends on a companion aiter kernel behind
  an env flag. This PR covers glm_moe_dsa using the in-tree Triton kernel
  with no aiter dependency and no new flags — complementary, not overlapping.
• [ROCm][DSv3.2] Eliminate per-decode FillFunctor launches in sparse-MLA hot loop #44527 ([ROCm][DSv3.2] Eliminate per-decode FillFunctor launches) removes
  the Fill kernels in the same hot loop; stacks cleanly on top of this change.

Risk / compatibility

Low. CUDA behavior is bit-for-bit unchanged; the ROCm change is opt-in via the
same static guard and only alters the previously-fragmented eager path. The
end-to-end gain is intentionally modest (kernel/HBM-traffic reduction under CUDA
graphs), but the change is trivial, carries no measured regression, and compounds
with related sparse-indexer cleanups (#44527). Requires the fused branch to fire
(is_inplace_rope == False), which is the default under torch.compile/Inductor.

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

[tjtanaa]

@xaguilar-amd please provide GSM8K with num-shot 30 as well. It validates the trigger of the sparse mla logic. num-shot 5 is short.

Please also provide the vllm serve commands.

[mergify]

Hi @xaguilar-amd, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

[xaguilar-amd]

@tjtanaa thanks for the review! Please find down below new accuracy numbers with num-shot 30:

local-completions ({'model': 'zai-org/GLM-5.2-FP8', 'base_url': 'http://localhost:8004/v1/completions', 'num_concurrent': 32, 'max_retries': 10, 'max_gen_toks': 2048, 'tokenizer_backend': 'None', 'tokenized_requests': 'False'}), gen_kwargs: ({}), limit: None, num_fewshot: 30, batch_size: auto

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	30	exact_match	↑	0.9454	±	0.0063
		strict-match	30	exact_match	↑	0.9462	±	0.0062

The server command was:

export VLLM_ROCM_USE_AITER=1
export VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1
export VLLM_ROCM_QUICK_REDUCE_QUANTIZATION=INT4

vllm serve zai-org/GLM-5.2-FP8
--enable_auto_tool_choice
--gpu_memory_utilization 0.83
--kv_cache_dtype fp8
--max-model-len 16384
--max_num_batched_tokens 16384
--no_enable_prefix_caching
--port 8004
--reasoning_parser glm45
--tensor_parallel_size 8
--tool_call_parser glm47
--trust_remote_code

[zhao18235158400]

Hi @xaguilar-amd — I read your ROCm enablement of the fused indexer-Q RoPE, FP8 quantization, and weight-scale fold for DeepSeek/GLM sparse attention. It’s a small code change that depends on careful kernel and backend validation. I’m Daniel, a user researcher with a product research team studying local AI systems for inference engineers. My main question is: what evidence or tooling do you need before reusing an existing fused CUDA-oriented path safely on ROCm? Would you be open to a Zoom conversation of up to 30 minutes when convenient? We’re preparing a prototype and may later provide test units or invite relevant participants as early users or technical advisors. This is research, not sales.

[maeehart]

@xaguilar-amd: can you check if this PR is still needed and let's drive it actively forward if it is.

[xaguilar-amd]

@maeehart the PR is still relevant until the serving migrates to the new implementation in vllm/models/deepseek_v32/amd (which is still experimental and for development and debugging purposes).

@zyongye could we please try to move this forward? It's just a 2-line change PR. It was already reviewed by TJ and CI is green. Thanks in advance!

[dllehr-amd]

@xaguilar-amd can you rerun this and confirm that this is still needed? It's been awhile and a lot of things have moved.

[amd-sriram]

@dllehr-amd Some testing on how this PR interacts with #51315 (AITER fused QK rope+quant+cache), since both touch the DSA indexer prologue.

With #51315 in place, its branch shadows this one whenever AITER >=v0.1.15, which is where our pin is. But this PR is still the fallback that matters: without it, if #51315 isn't available the indexer drops all the way to the fully-unfused path, since the ROCm in-place-rope branch is off under Inductor.

No conflict — they're consecutive tiers, so this one just sits one tier below:

if use_fused_indexer_qk:            # #51315, needs AITER >= v0.1.15
elif is_rocm and is_inplace_rope:   # off under Inductor
elif use_fused_indexer_q:           # this PR
else:                               # fully unfused

Worth merging both.

[maeehart]

Do I understand correctly that with #51315 and with the recent AITER version bump, this PR is not necessary anymore?

[xaguilar-amd]

yeah I guess, although we could have it as well as a safety net in case that AITER is not in place or whatever, it's just 2 lines of code changed. Your call @maeehart

[maeehart]

Yes, I think that it makes sense as a safety net. @dllehr-amd or @tjtanaa, can you check?