[ROCm][DSv4.1][Perf] Dequantize the MXFP8 weight once when dot_scaled cannot be used

[JohnQinAMD]

Purpose

RocmDotScaledMxfp8LinearKernel needs K % 128 == 0 for tl.dot_scaled. Where that does not hold, apply_weights dequantizes the weight instead:

w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)

layer.weight and layer.weight_scale are constant after loading, so that rebuilds the same tensor on every forward. This converts it once in process_weights_after_loading and keeps BF16 — what EmulationMxfp8LinearKernel already does under VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD — and has apply_weights dispatch on the weight's element size.

Why it is worth a patch

dequant_mxfp8_to_bf16 is six pointwise kernels, and its FP32 intermediate is 4x the FP8 source and round-trips through HBM twice, so a 2.95 MB weight costs 56 MB of traffic.

On DeepSeek-V4.1-Flash it fires on one linear per layer, and which one depends on TP: moe_intermediate_size is 2304, and the shared expert's down_proj is row-parallel, so its K is sharded while gate_up's stays at 5120.

TP	down_proj K	K % 128	path
1	2304	0	dot_scaled
2	1152	0	dot_scaled
4	576	64	dequantize
8	288	32	dequantize

TP1 and TP2 never see this; TP4 and TP8 pay it on all 40 layers — 2.24 GB of HBM traffic per decode token to produce 236 MB of weights that never change.

Test Plan

4x MI355X (gfx950), ROCm 7.2.3, TP4, on vllm/vllm-openai-rocm:nightly-eed1f3d0c6043bd494424a22443ee198dd56f657. Server as the recipe composes it — --tokenizer-mode deepseek_v41 --tensor-parallel-size 4 --gpu-memory-utilization 0.9 --moe-backend aiter_triton_mxfp4_bf16, VLLM_ROCM_USE_AITER=1, VLLM_ROCM_USE_AITER_MOE=1 — with cudagraph_mode=FULL_DECODE_ONLY pinned, and a fresh server per arm.

vllm bench serve --backend vllm --model /model \
  --served-model-name DeepSeek-V4.1-Flash --dataset-name random --ignore-eos \
  --random-input-len 8192 --random-output-len 1024 \
  --max-concurrency {1,8} --num-prompts {8,24}

The board drifts between processes, so the unpatched arm was run both first and last.

Test Result

arm	median TPOT, c=1	median TPOT, c=8
unpatched, first	10.60 ms	12.34 ms
patched	10.10 ms	12.05 ms
unpatched, last	10.57 ms	12.33 ms
	-4.6%	-2.3%

The two unpatched arms are 0.28% apart at c=1 and 0.08% at c=8, so the effect is 16-29x the drift over that span.

A decode profile confirms the mechanism rather than inferring it from the timing: exp2_kernel_vectorized4 and CUDAFunctorOnSelf_add<float> go from 2600 calls each to exactly zero, and bfloat16_copy_kernel_cuda from 2665 to 65 (the remainder is an unrelated once-per-step collapse). Weight memory rises 0.11 GiB per rank, which is the 236 MB of BF16 predicted from the weight shapes less the 118 MB of FP8 it replaces.

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

[ChuanLi1101]

Please fix before merge:

1. apply_weights drops the unaligned fallback. If the weight is still FP8 you always call dot_scaled, which is invalid when K % 128 != 0. Emulation still dequants on the fly in that case. Load keys off K % 128, apply keys off element_size() >= 2. If those disagree (e.g. process_weights_after_loading not run), this hits the illegal path. Please keep the per-step dequant fallback.

2. No test. Emulation already has test_mxfp8_linear_emulation_bf16_at_load (and those Ks are 128-aligned). Add the same for this kernel with K % 128 != 0: load-time BF16 matches per-step dequant, and a still-FP8 unaligned weight does not call dot_scaled.

3. The file header still says unaligned K falls through to emulation. It does not: can_implement is always True and the selector never sees K. Fix the comment.

[JohnQinAMD]

Please fix before merge:

1. apply_weights drops the unaligned fallback. If the weight is still FP8 you always call dot_scaled, which is invalid when K % 128 != 0. Emulation still dequants on the fly in that case. Load keys off K % 128, apply keys off element_size() >= 2. If those disagree (e.g. process_weights_after_loading not run), this hits the illegal path. Please keep the per-step dequant fallback.
2. No test. Emulation already has test_mxfp8_linear_emulation_bf16_at_load (and those Ks are 128-aligned). Add the same for this kernel with K % 128 != 0: load-time BF16 matches per-step dequant, and a still-FP8 unaligned weight does not call dot_scaled.
3. The file header still says unaligned K falls through to emulation. It does not: can_implement is always True and the selector never sees K. Fix the comment.

Thanks @ChuanLi1101 for the comments. All three comments addressed in 603edd28b.

1. The unaligned fallback. You are right that the two paths keyed off different things, and that the apply path stopped being safe on its own. The per-step dequant is back as a third branch, so an FP8 weight with an unaligned K never reaches dot_scaled regardless of whether process_weights_after_loading ran. The alignment the load and apply paths have to agree on is named now rather than repeated as a literal.

2. The test. test_mxfp8_rocm_native_unaligned_k_dequantizes_at_load in tests/kernels/test_minimax_m3_amd_ops.py, parametrised on K = 2080 and 2048, asserts what you asked for: the load-time BF16 weight matches a per-step dequant of the same bits, and an unaligned weight that never went through process_weights_after_loading does not reach _mxfp8_dot_scaled_linear. Without the restored branch the unaligned case fails.

3. The header. Corrected. The header now describes what the code does.

[ChuanLi1101]

Looks good overall. One small test gap:

• The regression test uses synthetic K=2080/2048.
• The actual DSv4.1 problematic shapes are K=576 / 288.
• Please add at least one real production shape to make sure the original issue is directly covered.

[JohnQinAMD]

Looks good overall. One small test gap:

• The regression test uses synthetic K=2080/2048.
• The actual DSv4.1 problematic shapes are K=576 / 288.
• Please add at least one real production shape to make sure the original issue is directly covered.

Thanks @ChuanLi1101 for the comments, the production shapes have been added in ceab01d: (5120, 576) and (5120, 288) — the DSv4.1 shared expert's down_proj is Linear(moe_intermediate_size=2304 -> hidden_size=5120) and row-parallel, so K is 576 at TP4 and 288 at TP8.

Dropped K=2080 at the same time: 2080 % 128 == 32, the same residue as 288, so it covered nothing 288 does not. K=2048 stays as the aligned control that has to keep reaching dot_scaled.

58 passed on gfx950.

[mergify]

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

[github-actions]

✅ @JohnQinAMD, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[mergify]

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

[JohnQinAMD]

/ci run

[github-actions]

✅ Triggered Buildkite CI #88951 for commit 7487778ac5fe.