[shen-shanshan]

Purpose

#58671 added a paged MXFP4 sparse indexer for DeepSeek-V4.1 on ROCm gfx950. The indexer K cache is stored as MXFP4 (68 B per token instead of 132 B for FP8) in the preshuffled page order that aiter's pa_mqa_logits_mxfp4 reads. That kernel scores the cache in place, so prefill no longer gathers K first. The PR has two layers:

• the dense path, which scores the whole context. It is model-agnostic.
• the candidate-pool path, which only exists because of V4.1's two-level indexer.

This PR brings the dense path to DeepSeek-V4 (e.g. DeepSeek-V4-Pro: 30 indexer layers at compress ratio 4, 64 heads × 128, top-k 1024). Today V4 on ROCm can only use the FP8 indexer cache, because:

• dsa_indexer_uses_fp4() rejects every model except deepseek_v41 on ROCm;
• SparseAttnIndexer.forward_hip asserts that the cache is not FP4;
• V4's own MXFP4 Q/K Triton kernels pack FP4 with CUDA PTX (cvt.rn.satfinite.e2m1x2.f32), so they cannot run on ROCm.

The feature is opt-in, as on V4.1:

VLLM_ROCM_USE_AITER=1 vllm serve deepseek-ai/DeepSeek-V4-Pro -tp 8 \
  --attention-config '{"indexer_kv_dtype": "mxfp4"}'

The default path (FP8 indexer cache) is unchanged.

Changes

K write: ROCm pooling kernel + aiter's cache op

In V4 the indexer K is produced by the indexer's own DeepseekCompressor. One fused Triton kernel does pool → RMSNorm → RoPE → MXFP4 → cache store, and nothing before the norm ever reaches memory. On ROCm we split it into two launches:

1. A new Triton kernel, _compress_pool_kernel, does the softmax-gated overlap pooling. It is copied from V4's fused kernel, up to the norm. It writes one bf16 row per boundary token into a scratch buffer that all indexer layers share ([max_num_batched_tokens, 128]).
2. aiter's public k_norm_rope_mxfp4_cache(..., shuffle=layout) then applies norm → RoPE → MXFP4 and stores into the preshuffled pages. This is the same op V4.1 uses, through the same rocm_mxfp4_indexer_k_store wrapper.

This keeps the page layout and the quantization defined only in aiter, which is the principle #58671 follows (rocm_paged_mxfp4_cache_layout is the only place vLLM reads the layout).

The cost is one extra small launch per indexer layer, plus a bf16 round trip of the pooled row before the norm. That rounding error (~2⁻⁹) is two orders of magnitude below the MXFP4 quantization step. A fully fused ROCm kernel can follow if profiles show the extra launch matters.

Q quant

The Q path uses aiter's q_rope_mxfp4_quant, through the existing rocm_mxfp4_indexer_q_quant. It produces fp32 weights with the softmax and head scales folded in, as the aiter MQA-logits kernel expects.

Scoring and metadata

Both are reused unchanged:

• RocmSparseAttnIndexer, which calls rocm_mxfp4_sparse_attn_indexer with no candidate blocks;
• the model-agnostic RocmMxfp4IndexerMetadataBuilder.

A new DeepseekV4RocmMxfp4IndexerBackend subclasses DeepseekV4IndexerBackend, so V4's 256-token blocks and ROCm adaptive-verification support carry over. At compress ratio 4 that gives 64-entry pages, which is aiter's preferred page size.

Keeping CUDA untouched

Shared V4 code only gains extract-method / ClassVar hooks. The code inside them is moved verbatim, and the ROCm subclasses override the hooks.

Other changes

• FusedKVCompressNormRopeInsertIndexerTritonKernel.get_warmup_keys() returns no keys on ROCm with an MXFP4 indexer cache. The warmup is a process-wide singleton that every non-CUDA compressor registers, including V4's head-512 attention compressors, so without this guard the PTX MXFP4 variant was compiled at startup and failed (error: unknown asm constraint 'f'). ROCm never runs that kernel: it writes the cache with aiter's op.
• dsa_indexer_uses_fp4(): the ROCm branch now accepts deepseek_v4 as well as deepseek_v41. The CDNA4 and aiter checks are unchanged, and the CUDA branch is untouched.
• _aiter_indexer_cache_ops, rocm_mxfp4_indexer_k_store and rocm_mxfp4_indexer_q_quant move from models/deepseek_v41/amd/rocm.py to v1/attention/ops/rocm_paged_mxfp4_indexer.py, so V4 and V4.1 share them. They are still importable from the old module.

Test Plan

Unit Test

pytest tests/kernels/attention/test_rocm_dsv4_mxfp4_indexer.py
pytest tests/kernels/attention/test_rocm_paged_mxfp4_indexer.py \
    tests/v1/attention/test_rocm_paged_mxfp4_indexer_plan.py
pytest tests/models/test_deepseek_v4_rocm_compressor_gemm_fusion.py \
    tests/models/test_deepseek_v4_rocm_wo_a.py \
    tests/models/test_deepseek_v4_dspark_rocm.py \
    tests/v1/attention/test_deepseek_v4_rocm_adaptive.py \
    tests/kernels/test_compressor_kv_cache.py \
    tests/kernels/test_deepseek_v4_cpu_kernels.py

Benchmark

Benchmark with SA InferenceX 8k1k and agentic workload with server launch commands below.

Fixed sequence length:

vllm serve /shared/models/deepseek-ai/DeepSeek-V4-Pro \
    --port 8333 \
    --tensor-parallel-size 8 \
    --data-parallel-size 1 \
    --max-model-len 9472 \
    --max-num-seqs 512 \
    --max-num-batched-tokens 16384 \
    --async-scheduling \
    --no-enable-prefix-caching \
    --distributed-executor-backend mp \
    --gpu-memory-utilization 0.8 \
    --kv-cache-dtype fp8 \
    --trust-remote-code \
    --moe-backend aiter \
    --tokenizer-mode deepseek_v4 \
    --reasoning-parser deepseek_v4 \
    --speculative-config '{"method": "mtp", "num_speculative_tokens": 2}' \
    --compilation-config '{"mode":3,"cudagraph_mode":"FULL_AND_PIECEWISE"}' \
    --load-format fastsafetensors \
    --attention-config '{"indexer_kv_dtype": "mxfp4"}'

Agentic replay:

vllm serve /shared/models/deepseek-ai/DeepSeek-V4-Pro \
    --host 0.0.0.0 \
    --port 8333 \
    --trust-remote-code \
    --async-scheduling \
    --distributed-executor-backend mp \
    --kv-cache-dtype fp8 \
    --max-num-batched-tokens 8192 \
    --tensor-parallel-size 1 \
    --data-parallel-size 8 \
    --prefill-schedule-interval 8 \
    --long-prefill-token-threshold 16384 \
    --gpu-memory-utilization 0.86 \
    --moe-backend aiter \
    --compilation-config '{"mode":3,"cudagraph_mode":"FULL_AND_PIECEWISE"}' \
    --speculative-config '{"method": "mtp", "num_speculative_tokens": 3}' \
    --tokenizer-mode deepseek_v4 \
    --tool-call-parser deepseek_v4 \
    --reasoning-parser deepseek_v4 \
    --enable-auto-tool-choice \
    --enable-prefix-caching \
    --no-disable-hybrid-kv-cache-manager \
    --max-num-seqs 64 \
    --load-format fastsafetensors \
    --attention-config '{"indexer_kv_dtype": "mxfp4"}'

Accuracy Test

Accuracy test with gsm8k dataset using lm_eval:

lm_eval \
  --model local-completions \
  --model_args model=/shared/models/deepseek-ai/DeepSeek-V4-Pro,base_url=http://0.0.0.0:8333/v1/completions,num_concurrent=64,max_retries=10,max_gen_toks=2048,timeout=60000 \
  --batch_size auto \
  --tasks gsm8k \
  --num_fewshot 5 \
  --output_path . 2>&1 | tee -a eval.log

Test Result

Unit Test

• New tests: 4 passed. The ROCm compressor's pages match aiter's writer fed by torch pooling byte for byte (0 mismatched bytes). As a sanity check, feeding a shifted window mismatches 21.8% of bytes, so the test does detect wrong pooling.
• V4.1 MXFP4 indexer tests: 37 passed.
• V4 ROCm / compressor tests: 99 passed, 137 skipped (platform-gated).

Benchmark

Fixed sequence length:

Conc	Version	Output token throughput (tok/s)	Change	Mean TTFT (ms)	Change	Mean TPOT (ms)	Change
1	Main (FP8 Indexer)	145.99	—	597.24	—	6.24	—
1	This PR (MXFP4 Indexer)	156.34	7.09% ↑	413.91	30.70% ↓	5.96	4.49% ↓
4	Main (FP8 Indexer)	454.91	—	504.07	—	8.02	—
4	This PR (MXFP4 Indexer)	466.92	2.64% ↑	437.29	13.25% ↓	7.95	0.87% ↓
16	Main (FP8 Indexer)	1020.90	—	744.88	—	14.51	—
16	This PR (MXFP4 Indexer)	1030.30	0.92% ↑	681.89	8.46% ↓	14.46	0.34% ↓
64	Main (FP8 Indexer)	1629.53	—	1637.36	—	36.85	—
64	This PR (MXFP4 Indexer)	1656.28	1.64% ↑	1622.58	0.90% ↓	36.29	1.52% ↓

In summary, this PR delivers a clear overall improvement over Main: throughput increases across all concurrency levels, Mean TTFT decreases substantially—especially at low and medium concurrency—and Mean TPOT remains slightly better across the board. The results suggest that the optimization is beneficial for both latency-sensitive workloads and overall serving throughput, with the most pronounced gains appearing at lower concurrency.

Agentic replay:

Indexer	Output Token Throughput (tok/s)	Mean TTFT (ms)	Mean TPOT (ms)	KV Cache Capacity (tokens)
Main (FP8 Indexer)	697.25	2,442.24	46.01	4,589,684
This PR (MXFP4 Indexer)	675.83 (3.07% ↓)	2,321.05 (4.96% ↓)	43.42 (5.62% ↓)	5,010,749 (9.17% ↑)

• Latency improves: mean TPOT -5.62%, mean TTFT -4.96%. TPOT is the most consistent signal: both MXFP4 runs (44.60 / 42.24 ms) had lower TPOT than both FP8 runs (46.98 / 45.03 ms). TTFT varied by up to 13% between runs of the same configuration, so treat its gain as indicative.
• Output token throughput is within noise (-3.07%). On this workload, aggregate throughput is driven mainly by the trace replay and by output lengths.
• KV cache capacity: +9.17%. The MXFP4 indexer cache takes 68 B per compressed token instead of 132 B, and the freed memory goes to the KV pool. Peak KV usage stayed below 40% in all runs, so this workload does not yet turn the extra KV capacity into throughput. Higher-concurrency or longer-context workloads should benefit more.

Accuracy Test

GSM8K:

Task	Filter	n-shot	Metric	Main (FP8)	This PR (MXFP4)	Change
gsm8k	flexible-extract	5	exact_match ↑	0.9492	0.9477	-0.16% ↓
gsm8k	strict-match	5	exact_match ↑	0.9500	0.9484	-0.17% ↓

The MXFP4 indexer in this PR shows very small accuracy degradation compared with the FP8 indexer in main. On GSM8K, the flexible-extract exact-match score decreases from 94.92% to 94.77% (-0.16%), while the strict-match score decreases from 95.00% to 94.84% (-0.17%). Overall, the accuracy impact is negligible (<0.2% relative), indicating that the MXFP4 indexer largely preserves the model's original accuracy.

GPQA Diamond (newly added):

Indexer	GPQA Diamond	Empty Responses (reasoning over budget)
Main (FP8 Indexer)	0.7121 ± 0.0323	13 / 198
This PR (MXFP4 Indexer)	0.7121 ± 0.0323	11 / 198

• Accuracy is unchanged: both indexers score 0.7121. Per question, 124 are correct in both runs, and the disagreements are symmetric (17 vs 17, McNemar p = 1.0).
• This exercises the indexer's long-context path: with thinking, reasoning traces run to tens of thousands of tokens. Most decode steps are therefore far past 4096 tokens (top-k 1024 × compress ratio 4), where the indexer starts its sparse top-k selection.
• KV cache capacity grows by 9.19%, since the MXFP4 indexer cache stores 68 B per token instead of 132 B.
• The 11–13 empty responses are questions whose greedy reasoning did not finish within the 131072-token budget. They affect both runs alike.

---

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.

[simondanielsson]

This looks good, just some small comments. Thanks!

Suggestion: Should we add gpqa and niah to ensure the implementation works at longer context? Don't think gsm8k would catch a potential issue in the indexer

[simondanielsson]

Suggestion: Would it be better to do this with the workspace manager?

[shen-shanshan]

Switched to the workspace manager. The indexer layers run one after another, so they all share one persistent buffer. forward claims it on the profiling run, before the workspace locks.

[simondanielsson]

Reminder to self: I think we need to run :amd: (MI355 DPX) Attention Kernels Shard %N manually in CI to make it run on mi350 as it's optional

[simondanielsson]

Suggestion: can we check if we need to register this for warmup?

[shen-shanshan]

With a fresh Triton cache, it was compiling during the CUDA graph warmup run, after JIT warmup. Registering it alone wasn't enough on ROCm: Triton's AMD backend gives pointers into storage of at most 2 GB a pointer_range=32 specialization, but TritonWarmupTensor always reported a small range, while the state cache views the KV pool, which is over 2 GB. So I added ptr_range_bytes to TritonWarmupTensor (default 0, only read by the ROCm backend) and warm up both pool sizes. With that, the variant used at runtime compiles during JIT warmup and nothing compiles after it.

[simondanielsson]

Good, I actually added the same fix in #60919 but I think we might merge this one first :)

[shen-shanshan]

This looks good, just some small comments. Thanks!

Suggestion: Should we add gpqa and niah to ensure the implementation works at longer context? Don't think gsm8k would catch a potential issue in the indexer

@simondanielsson Sorry for late reply, I'm busy with another work recently... Here is GPQA-Diamond A/B test results (also added into the PR body).

Indexer	GPQA Diamond	Empty Responses (reasoning over budget)
Main (FP8 Indexer)	0.7121 ± 0.0323	13 / 198
This PR (MXFP4 Indexer)	0.7121 ± 0.0323	11 / 198

• Accuracy is unchanged: both indexers score 0.7121. Per question, 124 are correct in both runs, and the disagreements are symmetric (17 vs 17, McNemar p = 1.0).
• This exercises the indexer's long-context path: with thinking, reasoning traces run to tens of thousands of tokens. Most decode steps are therefore far past 4096 tokens (top-k 1024 × compress ratio 4), where the indexer starts its sparse top-k selection.
• KV cache capacity grows by 9.19%, since the MXFP4 indexer cache stores 68 B per token instead of 132 B.
• The 11–13 empty responses are questions whose greedy reasoning did not finish within the 131072-token budget. They affect both runs alike.

[shen-shanshan]

/ci run

[github-actions]

✅ Triggered Buildkite CI #96632 for commit 8f51c20f74c3.

[simondanielsson]

LGTM