[Qwen4-Exp] Load nvidia/Qwen3.8-Flash-Next-NVFP4 (ModelOpt MIXED_PRECISION) on qwen4-main-squashed

[Jiminator]

Motivation

#37752 makes nvidia/Qwen3.8-Flash-Next-NVFP4 load, but it is stacked on #37500 (qwen4-main-squashed-rebased). This PR is the same change rebased onto qwen4-main-squashed (#36497). The two #37752 commits are cherry-picked with their original authorship; the port notes are in the commit messages and repeated below.

#36845 and #36806 needed no rebase: both were merged straight into qwen4-main-squashed and are its two tip commits (99c9362e66, 78c5024e9d).

The NVIDIA export (revision fab0aecb) is a ModelOpt MIXED_PRECISION checkpoint. RadixArk/Qwen3.8-Flash-Next-NVFP4 is not.

part	quant	checkpoint metadata
routed experts (48 layers)	NVFP4 g16	quantized_layers["model.language_model.layers.N.mlp.experts"]
PLE n-gram table	FP8 per-tensor	quantized_layers["model.language_model.layers.1.ple.ple_embedding.ngram_embedding"], 128 fp8 shard_N.weight + weight_scale
MTP experts	FP8_BLOCK_SCALES g128	quantized_layers["mtp.layers.0.mlp.experts"], weight_scale_inv
everything else	bf16	292-entry glob ignore list

config.json carries quantization_config (quant_method: modelopt, quant_algo: MIXED_PRECISION) and no text_config.ple_embedding_dtype. On qwen4-main-squashed the server dies while loading weights:

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

With --no-ple-offload-embedding it dies at CUDA-graph capture instead:

NotImplementedError: Unsupported moe_runner_backend for NVFP4 MoE: MoeRunnerBackend.FLASHINFER_TRTLLM. Use --moe-runner-backend flashinfer_cutlass instead.

Modifications

Commits 1 and 2 are #37752:

• ModelOptMixedPrecisionConfig: resolve model.language_model.* keys against the model.* prefixes Qwen4-Exp modules use; default exclude_modules to []; map FP8_BLOCK_SCALES to the 128x128 block Fp8LinearMethod / Fp8MoEMethod; expose resolve_quant_algo().
• Qwen4ExpNGramEmbedding: use fp8 table storage when the mixed config marks the ngram table FP8, so the pinned-host offload path sees fp8 storage up front.
• Qwen3 MoE family override: default modelopt_mixed to flashinfer_trtllm on sm100 when the expert layers are NVFP4/FP8; route W4A16_NVFP4 experts to marlin, or reject an explicit non-marlin runner.
• _mtp_quant_config: keep the mixed config for the MTP head when it lists mtp.* layers.

Port differences, forced by what qwen4-main-squashed has:

• The Qwen3 MoE override is still inline in python/sglang/srt/arg_groups/overrides.py (this branch has no arg_groups/model_overrides/qwen3_moe.py), so the change and its test target that file.
• ModelOptMixedPrecisionConfig has no fp8_pb_wo_config here, so a 128x128 block Fp8Config is added as fp8_block_config and only FP8_BLOCK_SCALES maps to it.
• The override test mocks with patch.object(overrides_module, "is_sm100_supported", ...), the idiom this branch's test file already uses.

Commits 3 and 4 came out of validating on 1x RTX PRO 6000 (sm120). They are separate so either can be dropped or moved:

• Commit 3: Fp8MoEMethod.create_moe_runner resolves flashinfer_cutlass / flashinfer_cutedsl to triton. The MoE runner backend is global. sm120 has no trtllm-gen, so NVFP4 experts need --moe-runner-backend flashinfer_cutlass, and for that backend Fp8MoEMethod created no runner at all; the FP8_BLOCK_SCALES MTP head then died in the draft CUDA-graph warmup with AttributeError: 'Fp8MoEMethod' object has no attribute 'runner'. On sm100 the override picks flashinfer_trtllm, which Fp8MoEMethod supports, so [Qwen4-Exp] Load nvidia/Qwen3.8-Flash-Next-NVFP4 (ModelOpt MIXED_PRECISION: NVFP4 experts, fp8 PLE, fp8 MTP) #37752 never hit this. Unit test: test/registered/unit/layers/quantization/test_fp8_moe_runner_fallback.py.
• Commit 4: cherry-pick of Fix KV cache pool sized far too small when weight-loading memory is still referenced #36583 (gc.collect() before profiling the KV budget). qwen4-main-squashed predates it. Without it, loader temporaries still referenced at profiling time understate free memory by about 6 GiB on a 96 GB card, and the pinned mamba pool plus KV no longer fit (Loaded weights leave no GPU memory for the KV cache under --mem-fraction-static=0.93). This is not specific to the checkpoint. Drop it here if the base branch picks it up from main.

Accuracy Tests

1x RTX PRO 6000 Blackwell (96 GB, sm120), TP1, lmsysorg/sglang:qwen38flashnext with this branch overlaid (sglang-kernel 0.4.6.post1 matches this branch's pin; the image ships flashinfer 0.6.18, the branch pins 0.6.17). GSM8K is the chat endpoint, 200 questions, thinking off, greedy, max_tokens=1024, at 64 concurrent requests (no MTP) or 16 (MTP).

Common flags: --fp4-gemm-backend flashinfer_cutlass --moe-runner-backend flashinfer_cutlass --ple-offload-embedding --page-size 64 --chunked-prefill-size 4096 --context-length 262144 --mamba-radix-cache-strategy extra_buffer_lazy --mamba-ssm-dtype bfloat16 --reasoning-parser qwen3, env SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK=1, no --quantization (resolved from the checkpoint). No MTP adds --max-running-requests 64 --max-mamba-cache-size 192 --mem-fraction-static 0.93. MTP adds --speculative-algorithm NEXTN --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4 --max-running-requests 16 --max-mamba-cache-size 48 --mem-fraction-static 0.96.

checkpoint	PLE offload	MTP	startup (launch to ready)	GSM8K (200, no think)	accept len
nvidia/Qwen3.8-Flash-Next-NVFP4	on	off	134 s (weights 61 s)	0.965 (197 stop / 3 length)	-
nvidia/Qwen3.8-Flash-Next-NVFP4	on	on	112 s (weights 59 s + draft 11 s)	0.975 (197 stop / 3 length)	3.35

Both boots log quant=modelopt_mixed, quant_algo=MIXED_PRECISION, keep 81.35 GB of weights on the GPU with the fp8 ngram table pinned in host RAM (scheduler RSS about 71 GB), and resolve the linear-attention kernels to triton. The MTP boot logs Fp8MoEMethod has no FLASHINFER_CUTLASS path; using triton for its fp8 experts. Pools: no MTP has 192 mamba slots (10.58 GB) and a 97,600-token KV pool, with 5.25 GB free after graph capture and 4.0 GB left at peak; MTP has 48 slots, 3.59 GB of intermediate SSM state and a 169,664-token KV pool, with 4.12 GB free after capture and 2.9 GB left at peak. Neither server log has an error.

Unit tests in the same container, plus pre-commit run --files <changed files>, all pass:

test/registered/unit/model_loader/test_modelopt_loader.py -k MixedPrecision        10 passed
test/registered/unit/test_model_overrides.py -k qwen3_moe                           2 passed
test/registered/unit/models/test_shared_experts_fusion_gates.py                    25 passed
test/registered/unit/layers/quantization/test_fp8_moe_runner_fallback.py            3 passed

Speed Tests and Profiling

bench_serving --backend sglang-oai --dataset-name random --random-input-len 1024 --random-output-len 256 --random-range-ratio 1 --request-rate inf --flush-cache, same servers as above:

config	concurrency (prompts)	output tok/s	mean TTFT	mean TPOT
no MTP, 64	16 (32)	536.8	1,159 ms	25.3 ms
no MTP, 64	64 (192)	881.3	3,318 ms	54.8 ms
MTP, 16	1 (8)	154.4	162 ms	5.86 ms
MTP, 16	16 (32)	669.5	727 ms	19.1 ms

RadixArk/Qwen3.8-Flash-Next-NVFP4 on the qwen4-main-squashed-rebased image with the same flags (#37995) measured GSM8K 97.0-97.5 and TPOT 5.8-5.9 ms at concurrency 1 with MTP, and GSM8K 97.0-98.0 and 856-861 tok/s at 64 without. The NVIDIA export lands in the same place.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

https://claude.ai/code/session_01V5B2WK8ciABmJ8pMBtGcgN

---

CI States

Latest PR Test (Base): ❌ Run #33963603943
Latest PR Test (Extra): ❌ Run #33963603948
Latest PR Test (AMD ROCm 7.2): ❌ Run #33963604062