[YIHONG-JIN]

What

Follow-up to #4025 (which landed the NVFP4 W4A4 NaN-clamp). Adds the
prefix=maybe_prefix(prefix, "lm_head") argument that the custom
Qwen3MoeLLMForCausalLM (the Thinker text backbone) was missing when building
lm_head, mirroring vLLM's stock Qwen3MoeForCausalLM.

Why

Without it, loading an NVFP4 W4A4 Qwen3-Omni checkpoint on vLLM 0.23.0 crashes at
load (in vocab_parallel_embedding.py, loading lm_head.weight):

RuntimeError: The size of tensor a (1024) must match the size of tensor b (2048) at non-singleton dimension 1

Root cause: on 0.23.0 the Thinker lm_head is kept BF16 only via ModelOpt's
prefix-based exclude-list (the W4A4 checkpoint ignores *.lm_head); with
prefix="" that match fails, the head is wrongly NVFP4-quantized (FP4-packed →
hidden 1024), and the BF16 [vocab, 2048] lm_head weight then fails to copy into
the 1024-wide param.

This only manifests on vLLM ≥ 0.23.0 — vLLM 0.22.x is unaffected. On 0.22,
ModelOpt.get_quant_method only matched LinearBase, so ParallelLMHead was
never a quantization candidate and the missing prefix had no effect — the
W4A4 checkpoint loaded fine (this is the stack the original W4A4 accuracy /
throughput validation ran on). 0.23.0 added ParallelLMHead to the
quantizable set, which is what makes the head's BF16 status depend on the
prefix. So this is a 0.23.0-migration fix, not a regression introduced here —
the omission has been latent since the original Qwen3-Omni port (#55).

BF16 Qwen3-Omni is also unaffected on any version (no quant_config, so the
head is never a quant candidate regardless of the prefix).

Reproduce

vLLM 0.23.0 + vllm-omni on Blackwell (B200 sm_100 / RTX 5090 sm_120), serving the
published W4A4 checkpoint:

vllm serve YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip --omni

The Thinker worker dies at weight load (StageEngineCoreProc died during READY,
worker traceback qwen3_omni_moe_thinker.py:load_weights -> ... -> vocab_parallel_embedding.py) with the 1024 vs 2048 RuntimeError above. A debug
print at Qwen3MoeLLMForCausalLM.__init__ shows lm_head built as
[152064, 1024] (FP4-packed) instead of [152064, 2048]. With this PR it is
[152064, 2048] and the model loads.

To then serve end-to-end (independent of this PR) also set
VLLM_USE_FLASHINFER_MOE_FP4=1 VLLM_FLASHINFER_MOE_BACKEND=throughput to select
the Cutlass NvFp4 MoE backend — the default selects TRT-LLM, which separately
fails on this checkpoint's expert shapes.

Test

Verified on B200 (sm_100, FlashInfer Cutlass FP4) and RTX 5090 (sm_120): with the
fix the W4A4 Thinker loads (embed_tokens and lm_head both [152064, 2048])
and serves coherent text (no !!!!). Before the fix it crashes at load as above.

🤖 Generated with Claude Code

[amy-why-3459]

Thanks, LGTM.

[hsliuustc0106]

Bugfix looks correct.

[YIHONG-JIN]

@hsliuustc0106 Can you help approve this if everything looks fine? I suspect the impact is not limited to NVFP4

[david6666666]

LGTM