[Fix][ROCm] MXFP4 MoE round-up inflates TP-sharded expert weights on CDNA3, starving KV cache

[vllmellm]

Purpose

Fixes #53961

While testing #53961, I found that the per-GPU VRAM used for weight loading is 135.35 GiB (without EP) vs 103.54 GiB (with EP). Upon closer investigation, it is caused by the rocm's branch of mxfp4_round_up_hidden_size_and_intermediate_size which rounds the intermediate moe size to 256 which inflated the expert weights usage. Before this PR, an existing special case for Kimi-K3 (AITER+SiLU/SiTU) already aligns to 128 instead of 256.

Concretely, for DeepSeek-V4-Pro (moe_intermediate_size=3072, TP=8, no EP):

• per-partition intermediate = 3072/8 = 384 fp4 elems (192 B)
• round_up(384, 256) = 512 → w13_weight created as (384, 1024, 3584) instead of (384, 768, 3584) → +33% expert weights = +30 GiB/GPU
• with the default max_num_batched_tokens (16384 since [Config] Update default _max_num_batched_tokens from 8192 to 16384 #51726), caused the observation from [Bug]: PR #51726's >=160GiB tier captures MI300X (191.98 GiB), doubling max_num_batched_tokens and breaking KV cache init for DeepSeek-V4-Pro at 1M #53961

EP wasn't affected because each GPU holds whole experts, so the intermediate dimension stays at its native 3072 (already aligned); without EP, TP slices it to 384 per GPU, which is not aligned and gets padded.

Change

I extended the existing 128-alignment exception to also cover TRITON_UNFUSED with SiLU/SiTU activation. The pre-existing AITER arm is unchanged: K3's AITER kernel (gfx950-only) handles its native intermediate size (moe_intermediate 3072; e.g. 384/partition at TP8), so K3 keeps 128 on CDNA4.

For the new TRITON_UNFUSED arm, the 128-alignment is gated to CDNA != 4, because opt_flags_amd sets matmul_ogs's block_k=128 whenever weight scales are present on CDNA != 4 (always the case for mxfp4 W4A16 checkpoints) — the unpadded 384-element shard already satisfies the kernel's even-K requirement. gpt-oss (SWIGLUOAI, tuned in #38043) keeps 256.

In essence, the checks keeps the pre-existing Kimi-K3, gpt-oss handling while letting DeepSeek-V4 Pro (TRITON_UNFUSED + SiLU/SiTU on CDNA3) keeps the 128 size without rounding up to 256.

Test Plan

Tested the before/after using this serve command:

# max batched tokens: 16384 will failed before this PR
VLLM_ROCM_USE_AITER=1 vllm serve deepseek-ai/DeepSeek-V4-Pro-0813 \
  --trust-remote-code \
  --kv-cache-dtype fp8 \
  --block-size 256 \
  --tensor-parallel-size 8 \
  --distributed-executor-backend mp \
  --gpu-memory-utilization 0.9 \
  --max-num-seqs 512 \
  --max-num-batched-tokens 16384 \
  --compilation-config '{"mode": 3, "cudagraph_mode": "FULL_DECODE_ONLY"}' \
  --tokenizer-mode deepseek_v4 \
  --tool-call-parser deepseek_v4 \
  --enable-auto-tool-choice \
  --reasoning-parser deepseek_v4 \
  --no-enable-prefix-caching

Test Result

• There is a separate PR [ROCm][DSv4] Fix sparse-indexer logits collapse on gfx950/gfx942 #50455 that tackles the accuracy issue directly. The logprobs and the subset of gsm8k tests here is to show there's no negative impact to the inference after the changes.

check	before	after
Model loading took (TP8, no EP)	135.35 GiB	103.54 GiB
1M-ctx serving @ util 0.9 + 16384	failed to start	successfully started
greedy outputs, 5 prompts, top-1 logprobs	baseline	240/240 tokens identical, 5/5 exact text
round-up unit check (5 backend/activation/cdna combos)	—	all match expected values
GSM8K (20-shot, limit 200, conc 32)	0.960 flex / 0.685 strict	0.955 flex / 0.700 strict
serving bench 8k-in/1k-out, conc 32, 128 prompts	369 out-tok/s	421 out-tok/s (+14%) — likely due to the unpadded K removes dead weight traffic

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

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

[Rohan138]

LGTM, nice catch!

[github-actions]

✅ @vllmellm, 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.

[Rohan138]

/ci run

[github-actions]

❌ @Rohan138, Only reviewers with write access can use CI commands before CI is delegated to the PR author.

[AndreasKaratzas]

/ci run

[github-actions]

❌ This PR is 250 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[Rohan138]

@vllmellm can you rebase and run ci?

[vllmellm]

/ci run

[github-actions]

❌ This PR is 1 commit behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[vllmellm]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89518 for commit ec07102e4d31.

[mergify]

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

[vllmellm]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89526 for commit 57bd71ef21bc.

[vllmellm]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89622 for commit b0157f762115.

[tjtanaa]

/ci run

[tjtanaa]

/amd-ci run

[github-actions]

✅ Triggered Buildkite CI #89626 for commit 87dafb76e680.

[github-actions]

❌ This PR is 1 commit behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /amd-ci run. To test this branch at your own risk, use /amd-ci run --allow-stale.