[ROCm][DSv4][WIP] Sparse-MLA bring-up on MI300X (FP8 encoder/decoder symmetry + cudagraph fixes)

[maeehart]

Summary

Bring up DeepSeek-V4 sparse-MLA on MI300X (gfx942) using the
ROCm/AITER vLLM stack. End-to-end the model now runs with TP=4,
--kv-cache-dtype fp8, VLLM_ROCM_USE_AITER=1, and cudagraph
FULL_AND_PIECEWISE enabled, and produces coherent topical output
("capital of France" -> Paris, "largest planet" -> Jupiter,
"27 times 43 = 1161" correct on some samples). Accuracy is still
degraded vs CUDA baseline (GSM8K@200 ~1-2%), which is being tracked
as follow-up FP8 work outside the DSv4-specific Triton kernels (FMOE,
linear). Throughput at ISL=1000/OSL=100, max-concurrency=16 is
~301 tok/s output / ~3300 tok/s total on a single MI300X node.

What's in this branch

Two cherry-picks from @ganyi1996pku in
vllm-project#41451
provide the C++ kernel + initial MI300X plumbing, and a set of
follow-on commits make the DSv4-specific Triton kernels and Python
glue symmetric on FNUZ-only hardware.

#	Commit	Origin
1	mi300 support	cherry-pick from @ganyi1996pku (#41451)
2	accuracy right (gfx950 gate on the C++ SWA-K encoder)	cherry-pick from @ganyi1996pku (#41451)
3	[Build] hipify.py: copy already-HIP-native .cu sources to .hip	mine
4	[ROCm][DSv4] Remove stale ffn_norm call from _forward_rocm	mine
5	[ROCm][DSv4] Use tl.float8e4b8 for FNUZ on MI300X sparse MLA kernels	mine
6	[ROCm][DSv4] Fix compressed K cache dequant to match Triton OCP encoder (Plan A; superseded by vllm-project#7 + vllm-project#8)	mine
7	[ROCm][DSv4] Make compressed-K Triton encoder FNUZ-aware on gfx942	mine
8	[ROCm][DSv4] Make sparse indexer Q quant FNUZ-aware on gfx942	mine
9	[ROCm][DSv4] Make sparse-attn output FP8 quant FNUZ-aware on gfx942	mine

Commits vllm-project#6 and vllm-project#7 should be squashed before any upstream submission;
the standalone Plan-A change reverts cleanly inside Plan-B and is kept
here only so the bring-up narrative is reproducible.

Root cause story (what the FP8 commits fix)

DSv4 sparse-MLA touches four distinct FP8 encoders, all writing into
buffers that downstream Triton/AITER readers interpret as
tl.float8e4nv (OCP). On gfx942 MI300X the MFMA instructions only
support FNUZ, so Triton silently lowers .to(tl.float8e4nv) casts
to FNUZ. That means the bytes in those buffers were FNUZ-encoded but
the readers interpreted them as OCP. The four mismatched paths and
their fixes:

1. Compressed-K cache (this PR commits Support beam search & parallel generation vllm-project/vllm#7 + retired Automatically configure KV cache size vllm-project/vllm#6) -
   _fused_kv_compress_norm_rope_insert_sparse_attn /
   _fused_kv_compress_norm_rope_insert_indexer_attn (Triton) wrote
   bytes with FP8_MAX=448.0 and tl.float8e4nv. Decoder
   (dequantize_and_gather_k_cache) reads them back. On MI300X
   both now use USE_FNUZ=True with FP8_MAX=240.0 and the cast
   target tl.float8e4b8; the decoder keeps
   use_fnuz=current_platform.is_fp8_fnuz().

2. SWA-K cache (already correct after cherry-pick Support tensor parallel vllm-project/vllm#2) - C++ kernel
   fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert now writes
   FNUZ-encoded bytes on gfx942 and OCP elsewhere; decoder uses
   use_fnuz=current_platform.is_fp8_fnuz().

3. Sparse indexer Q (this PR commit Add miscellaneous updates vllm-project/vllm#8) -
   _fused_indexer_q_rope_quant_kernel was producing bytes consumed
   by AITER's fp8_mqa_logits. Plumbed FP8_MAX and USE_FNUZ
   constexprs through the kernel; on FNUZ hardware it casts to
   tl.float8e4b8 with FP8_MAX=240.0 and the tensor dtype
   becomes torch.float8_e4m3fnuz so the dot-product reader uses
   the matching bias.

4. Sparse-attention output (this PR commit Implement LLaMA vllm-project/vllm#9) -
   _fused_inv_rope_fp8_quant_per_head feeds the wo_b einsum at the
   end of every block. Same shape of fix: USE_FNUZ constexpr,
   current_platform.fp8_dtype() for the buffer, fp8_max from
   the platform-aware finfo.

Verification

Container: vllm-dsv4-mi300 on chi-mi300x-004 (single MI300X
node, TP=4). Launch (in /host_logs/run_smoke.sh):

vllm serve deepseek-ai/DeepSeek-V4-Flash \
  --tensor-parallel-size 4 --max-model-len 32768 \
  --max-num-batched-tokens 8192 --gpu-memory-utilization 0.85 \
  --kv-cache-dtype fp8 --trust-remote-code --port 8000

Environment: VLLM_ROCM_USE_AITER=1,
HIP_VISIBLE_DEVICES=0,1,2,3, AITER fp8_mqa_logits patched to
BLOCK_KV=64 in the container's installed aiter package (gfx942
has 64 KB LDS / CU; the upstream 128 wants 96 KB).

Smoke (/host_logs/dsv4_smoke.sh, three back-to-back runs)

Sample outputs (max_tokens=32, temperature=0):

=== capital of France ===
run 1: "afamily of the French Republic. ... capital of the region is Toulouse..."
run 2: "a city in the country of France."
run 3: "a country in Europe. ... Its capital is Paris. Its official language is French..."

=== /v1/chat/completions math (27*43) ===
run 1: "Number of the beast"
run 2: "to the question 'What is 27*43?' is 1161."      <- correct
run 3: "original research 1. Introduction 2. Literature Review..."

The non-determinism between identical deterministic requests
(temperature=0) suggests residual workspace / state-cache uninitialised
memory; documented as follow-up.

lm_eval gsm8k

Both with --num_fewshot=5 against the running server:

Mode	Limit	flexible-extract	strict-match
num_concurrent=32	200	0.01 (1%)	0.00
num_concurrent=1	50	0.02 (2%)	0.00

Coherent end-to-end output but precision is still off; the next
FP8 paths to audit are AITER FMOE and the FP8 block-quant linear
weights (fp8.py / process_weights_after_loading).

vllm bench serve ISL=1000/OSL=100, max-concurrency=16

Successful requests: 64 / 64
Benchmark duration: 21.25 s
Output token throughput: 301.16 tok/s   (peak 448.00)
Total token throughput:  3312.80 tok/s
Mean TTFT:  1103 ms     Median TTFT: 949 ms     P99 TTFT: 1992 ms
Mean TPOT:  42.4 ms     Median TPOT: 40.2 ms    P99 TPOT: 52.4 ms
Mean ITL:   42.4 ms     Median ITL:  36.4 ms    P99 ITL:  271 ms

What's still broken (next steps)

1. GSM8K accuracy ~1-2% despite coherent output. Next paths to
   audit on FNUZ hardware:
   • AITER FMOE FP8 kernels (block-quant FP8 weights + activations).
   • Generic FP8 linear weight load in
     vllm/model_executor/layers/quantization/fp8.py -
     process_weights_after_loading already handles UE8M0 and FNUZ
     for scales but the activation quant in
     Fp8LinearOp should be double-checked end-to-end.
   • fp8_einsum consumer of the inv-rope output (DeepGEMM /
     AITER) - now reads FNUZ buffers, verify GEMM-side bias.
2. AITER BLOCK_KV=64 patch lives in the container's installed
   aiter package, not under git. Either upstream a small AITER PR
   gating BLOCK_KV on shared-memory budget, or expose a kwarg in
   rocm_aiter_mla_sparse.py so callers can pick the safe block.
3. Non-determinism across identical temperature=0 requests even
   in --enforce-eager mode. Suspect: a workspace tensor allocated
   with torch.empty and only partially overwritten when seq-length
   padding kicks in. Worth bisecting once accuracy is restored.
4. Full DeepSeek-V4 at TP=8 deferred until the accuracy regression
   above is fixed - infrastructure (build, container, run scripts)
   is in place at conversations/vllm-fp8-mla-dense-prefill-pr42509/.

Attribution

The C++ kernel + initial wiring for DSv4 sparse-MLA on MI300X were
written by @ganyi1996pku in
vllm-project#41451;
cherry-picks #1 and vllm-project#2 in the table above are theirs. The Triton
encoder / decoder symmetry fixes and the build/cudagraph adjustments
are mine.

Test plan

• Smoke (/host_logs/dsv4_smoke.sh) - model produces coherent
  topical output across multiple back-to-back runs.
• lm_eval gsm8k --limit 200 --num_fewshot 5 - completes; result
  recorded above (1% flexible-extract).
• vllm bench serve ISL=1000/OSL=100 max-concurrency=16 -
  completes; result recorded above.
• Restore GSM8K to the CUDA baseline (separate follow-up).
• Re-run on full DeepSeek-V4 at TP=8 (separate follow-up).

Made with Cursor

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀