[ROCm][DSv4] Functional fixes for DeepSeek V4 on MI300X/MI325X

[tuukkjs]

Purpose

This PR fixes DeepSeek V4 / DeepSeek V4 Flash functional issues on ROCm gfx942 (MI300X), while preserving the gfx950 path.

The main fixes are:

• Avoid unsupported FP8 arithmetic when converting UE8M0 block scales.
• Make the fused DeepSeek V4 qnorm/RoPE/KV-insert path emit the correct FP8 bytes on ROCm:
  • gfx942 uses FNUZ encoding.
  • gfx950 uses OCP encoding.
  • FNUZ uses max value 224.0.
• Propagate the FNUZ/OCP choice through DeepSeek V4 cache utility helpers and sparse MLA decode paths.
• Keep the sparse decode cache formats explicit: the main/SWA cache follows the platform FP8 encoding, while the compressed extra cache remains OCP.
• Provide the gfx942 Triton fallback for fp8_mqa_logits so the DSv4 ROCm path does not depend on an immediate aiter update for that kernel.
• Fix the fused-kernel test references:
  • allow a narrow 1 bf16 ULP tolerance for the RoPE region, where the GPU kernel and PyTorch reference can round opposite ways at ties;
  • use the correct ROCm FP8 store dtype in full-cache per-tensor FP8 reference checks, so gfx942 FNUZ bytes are decoded under the same encoding used by the kernel.

Continues #42893. AI assistance was used for this PR.

Co-authored-by: maeehart, markus.hartikainen@amd.com
Co-authored-by: ganyi, ygan@amd.com
Co-authored-by: Jin Tao, jintao12@amd.com

Test Plan

Build branch native extensions for gfx942 from:

vllm/vllm-openai-rocm:base-nightly-54bbf5166842932fa7abc34a14df850594daeb5e

Focused kernel test:

python3 -m pytest -q tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py

End-to-end validation:

• Start deepseek-ai/DeepSeek-V4-Flash with TP4, VLLM_ROCM_USE_AITER=1, --kv-cache-dtype fp8_e4m3, --tokenizer-mode deepseek_v4, --moe-backend triton_unfused, --max-model-len 32768, and --max-num-batched-tokens 8192.
• Disable the tilelang MHC path for this validation so the run uses the coherent non-tilelang MHC path; tilelang MHC correctness on gfx942 is a separate issue from this PR.
• Run a prompt sanity check against the OpenAI-compatible completions endpoint.
• Run full GSM8K, 5-shot, 1319 questions, local completions, concurrency 128.
• Run vllm bench serve with random 8192 input / 1024 output, --ignore-eos, and max concurrency 1,2,4,8,16,32.

Serve command:

export VLLM_ROCM_USE_AITER=1

vllm serve deepseek-ai/DeepSeek-V4-Flash \
  --tensor-parallel-size 4 \
  --kv-cache-dtype fp8_e4m3 \
  --max-model-len 32768 \
  --trust-remote-code \
  --tokenizer-mode deepseek_v4 \
  --moe-backend triton_unfused \
  --gpu-memory-utilization 0.85 \
  --distributed-executor-backend mp \
  --max-num-batched-tokens 8192 \
  --host 0.0.0.0 \
  --port 8000 \
  --disable-log-stats

GSM8K lm_eval command:

lm_eval --model local-completions \
  --model_args model=deepseek-ai/DeepSeek-V4-Flash,base_url=http://localhost:8000/v1/completions,num_concurrent=128,tokenized_requests=False,tokenizer_backend=none,max_length=32768 \
  --tasks gsm8k \
  --num_fewshot 5 \
  --batch_size 128

Bench command:

for C in 1 2 4 8 16 32; do
  vllm bench serve \
    --backend vllm \
    --base-url http://127.0.0.1:8000 \
    --model deepseek-ai/DeepSeek-V4-Flash \
    --tokenizer deepseek-ai/DeepSeek-V4-Flash \
    --trust-remote-code \
    --dataset-name random \
    --random-input-len 8192 \
    --random-output-len 1024 \
    --num-prompts $((3*C)) \
    --max-concurrency "${C}" \
    --ignore-eos \
    --seed 0
done

Test Result

Validation stack for the MI300X run:

vllm/vllm-openai-rocm:base-nightly-54bbf5166842932fa7abc34a14df850594daeb5e
  + this PR branch source
  + this PR branch native extensions rebuilt for gfx942
  + tilelang uninstalled for server/GSM8K validation to use the coherent non-tilelang MHC path

Focused Kernel Test

On gfx942:

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
155 passed

The full-cache fp8/bf16 cases are included in this result.

MI300X End-To-End Validation

Model/server:
DeepSeek-V4-Flash, TP4, fp8_e4m3 KV cache, triton_unfused MoE backend

GSM8K full 1319-question, 5-shot, local completions, concurrency 128:

Task	Filter	n-shot	Metric	Value	Stderr
gsm8k	flexible-extract	5	exact_match	0.9568	0.0056
gsm8k	strict-match	5	exact_match	0.9568	0.0056

Clean vllm bench serve, random 8192 input / 1024 output, --ignore-eos.
All rows had zero failed requests:

Max concurrency	Successful	Failed	Output tok/s	Total tok/s	Median TTFT (ms)	Median TPOT (ms)
1	3	0	17.16	154.41	196.24	58.16
2	6	0	33.50	301.48	317.18	59.46
4	12	0	65.47	589.22	535.18	60.63
8	24	0	125.91	1133.18	912.01	62.71
16	48	0	230.80	2077.16	1676.96	67.79
32	96	0	404.87	3643.80	3122.19	76.03

Additional Validation With PR #43950 Applied

This run validates the same PR branch with vLLM PR #43950 applied locally, leaving TileLang installed. PR #43950 changes ROCm MHC dispatch so the MI300X/AITER path uses AITER MHC when available, rather than selecting the TileLang fused MHC path just because TileLang is importable.

Validation stack:

vllm/vllm-openai-rocm:base-nightly-b8336c3c7c298e0878f22a7bf70f4e295b2f4e01
  + this PR branch source
  + vLLM PR #43950 applied locally
  + native extensions rebuilt for gfx942
  + tilelang installed and importable: tilelang==0.1.10

Focused Kernel Test With PR #43950

On gfx942:

tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py
155 passed

MI300X End-To-End Validation With PR #43950

Model/server:
DeepSeek-V4-Flash, TP4, fp8_e4m3 KV cache, triton_unfused MoE backend, VLLM_ROCM_USE_AITER=1, tilelang installed

Prompt sanity check:

Q: What is 17 + 25? Answer briefly.
A: 42.

GSM8K full 1319-question, 5-shot, local completions, concurrency 128:

Task	Filter	n-shot	Metric	Value	Stderr
gsm8k	flexible-extract	5	exact_match	0.9560	0.0056
gsm8k	strict-match	5	exact_match	0.9553	0.0057

Clean vllm bench serve, random 8192 input / 1024 output, --ignore-eos.
All rows had zero failed requests:

Max concurrency	Successful	Failed	Output tok/s	Total tok/s	Median TTFT (ms)	Median TPOT (ms)
1	3	0	31.09	279.84	1473.94	30.35
2	6	0	60.26	542.34	1570.09	31.21
4	12	0	115.31	1037.76	2394.49	31.62
8	24	0	204.51	1840.55	3127.39	36.02
16	48	0	327.45	2947.09	3288.08	43.95
32	96	0	478.61	4307.45	3390.08	61.28

[tuukkjs]

• Disable the tilelang MHC path for this validation so the run uses the coherent non-tilelang MHC path; tilelang MHC correctness on gfx942 is a separate issue from this PR.

I noticed this line. If this is the case we should add the _on_gfx942 condition to the mhc so that it is using torch implementation instead.

So that we can use it on mi300x out of the box, without the need to uninstall tilelang to make it work on mi300x. We must ensure that the mi355x still uses the tilelang code path.

Yes true. I opened a separate PR for the fix: #45931.

It adds an MHC-specific TileLang capability check that disables TileLang MHC only on gfx942, so MI300X/MI325X fall back to the existing torch/triton implementation while gfx950/MI355X keeps using TileLang.

[tjtanaa]

LGTM. Thanks. Let's land this and I will quickly look at the tilelang PR. All coauthors are included.

[mergify]

Hi @tuukkjs, 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.

[mergify]

Hi @tuukkjs, 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.

[mergify]

Hi @tuukkjs, 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.

[mergify]

Hi @tuukkjs, 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.

[mergify]

Hi @tuukkjs, 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.

[mergify]

Hi @tuukkjs, 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.

[tjtanaa]

A remark:

import needs to be made explicit else mypy fail to resolve it and cause precommit error.

$ mypy --python-version 3.10 
vllm/models/deepseek_v4/nvidia/ops/o_proj.py:48: error: Module not callable  [operator]
Found 1 error in 1 file (checked 686 source files)