[Model] Qwen4-Exp: honour per-layer FP8 for the PLE table and the MTP experts in ModelOpt mixed checkpoints

[zhendonghua]

Motivation

nvidia/Qwen3.8-Flash-Next-NVFP4 is a three-precision ModelOpt MIXED_PRECISION checkpoint: NVFP4 routed experts, an FP8 PLE n-gram table (128 F8_E4M3 shards + a scalar weight_scale), and block-FP8 MTP experts. Two decisions in the Qwen4-Exp model code were made before consulting the checkpoint's per-layer quantized_layers map, and both break on it. RadixArk/Qwen3.8-Flash-Next-NVFP4 (plain NVFP4, PLE declared via text_config.ple_embedding_dtype, bf16 MTP) and Qwen/Qwen3.8-Flash-Next-FP8 (fp8) are on other code paths and are not affected.

1. Startup failure. The PLE table's storage dtype is fixed at construction and only became fp8 for a whole-checkpoint fp8 quant name or an explicit text_config.ple_embedding_dtype. This checkpoint sets neither, so the table was built bf16, moved into pinned host memory by the default --ple-offload-embedding, and load_weights then hit

ValueError: fp8 PLE auto-switch is unsupported with ple_offload_embedding; set text_config.ple_embedding_dtype="float8_e4m3fn" instead

2. Silent draft corruption. _mtp_quant_config dropped the quant config for every modelopt_mixed checkpoint, assuming embedded MTP weights are bf16. Here mtp.layers.0.mlp.experts is block-FP8, so the draft's FusedMoE was built unquantized: fp8 expert values were cast to bf16 without their scales and weight_scale_inv was skipped by the loader's ignore list. The server ran and target accuracy was fine, but MTP accept length fell to ~1.6.

Mirrors vllm-project/vllm#54882 (PLE) and #55513 (MTP).

Modifications

• models/qwen4_exp.py: _ple_table_is_fp8() additionally asks the mixed config's _resolve_quant_algo() for the embedding's own prefix (...ple.ple_embedding.ngram_embedding), so the table is fp8 from construction; the prefix is threaded through Qwen4ExpPLELayer → Qwen4ExpNGramEmbedding.
• models/qwen3_5_mtp.py: for modelopt_mixed, keep the quant config when quantized_layers lists any mtp. layer. Every other draft module still resolves to bf16 through the same per-layer lookup (verified: qkv_proj, shared_expert, gate, lm_head all resolve to None).

Depends on #38726 (block-FP8 FusedMoE dispatch in ModelOptMixedPrecisionConfig). This branch is stacked on that one, so its diff also shows #38726's commit until it merges; please review only the second commit here.

Test

• New test/registered/unit/models/test_qwen4_exp_mixed_precision.py (cpu): PLE fp8 detection from the per-layer map incl. the language_model.model. prefix drift; draft keeps / drops the quant config depending on mtp. entries. Verified red on the pre-change code.
• E2E on 4x B300, user launch command (TP4, NEXTN steps 3 / topk 1 / draft 4, no --quantization flag): loads; sgl-eval GSM8K 200 examples, thinking, max_tokens 16384: 0.97 (CI threshold 0.94), stop_rate 1.0; accept length 3.2–3.6 (was 1.6); draft weights 1.53 GB (was 2.47 GB bf16).

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel if you have any questions.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34413177538
Latest PR Test (Extra): ❌ Run #34413177514
Latest PR Test (AMD ROCm 10): ❌ Run #34413177487

[janbernloehr]

We independently reproduced and isolated the PLE startup failure addressed by this PR.

Configuration:

• Model: nvidia/Qwen3.8-Flash-Next-NVFP4
• Hardware: 4× GB300
• Architecture: arm64
• Parallelism: TP4
• Base: lmsysorg/sglang:dev, reporting SGLang 0.0.0.dev1+g00143e9c2
• Quantization automatically detected as modelopt_mixed
• ple_offload_embedding=True
• No explicit --quantization argument
• No explicit --moe-runner-backend argument
• Speculative decoding disabled

On the unpatched image, the FP8 PLE n-gram table was constructed with the wrong storage dtype and startup failed with:

ValueError: fp8 PLE auto-switch is unsupported with ple_offload_embedding

We then applied only #38727’s runtime changes to qwen4_exp.py and qwen3_5_mtp.py. The isolation controls verified all of the following:

1. _ple_table_is_fp8 was absent before patching.
2. It was present after patching.
3. [Quant] ModelOpt mixed precision: dispatch block-FP8 MoE experts and derive the block size #38726’s _BLOCK_FP8_ALGOS changes remained absent.
4. [Config] Qwen3 MoE family: default modelopt_mixed to the flashinfer_trtllm MoE runner on SM100 #38728’s modelopt_mixed MoE-runner override remained absent.
5. The patched files compiled successfully.

With #38727 alone, the PLE failure disappeared. SGLang detected the NVFP4 checkpoint, proceeded past PLE construction, and loaded all checkpoint shards. This directly confirms that resolving the PLE table dtype from the per-layer quantized_layers entry fixes the reported startup defect.

The run subsequently failed later, in NVFP4 expert post-processing, because the unresolved default remained:

• quantization=modelopt_mixed
• moe_runner_backend=auto

The later failure was:

AssertionError: The intermediate size required padding, but padding is also implemented for gated activations

That is the separate failure addressed by #38728. Therefore, our result is:

• [Model] Qwen4-Exp: honour per-layer FP8 for the PLE table and the MTP experts in ModelOpt mixed checkpoints #38727 independently fixes its PLE dtype-selection defect.
• [Model] Qwen4-Exp: honour per-layer FP8 for the PLE table and the MTP experts in ModelOpt mixed checkpoints #38727 alone is not sufficient to serve this checkpoint under the default TP4 configuration.
• The _mtp_quant_config portion of [Model] Qwen4-Exp: honour per-layer FP8 for the PLE table and the MTP experts in ModelOpt mixed checkpoints #38727 was not exercised because speculative decoding was disabled and the mtp.* layers were never loaded.

We also tested the complete #38726/#38727/#38728 stack with the same model and TP4 shape on arm64 GB300 and GB200. In that configuration, the backend resolved to flashinfer_trtllm, checkpoint loading completed, the server became ready, and both serving and validation workloads completed.