[SM70] FP8-resident MTP experts for AWQ and NVFP4 targets

[Leonccaa]

Problem and implementation

Qwen4Exp MTP experts can consume substantial memory even when the target is AWQ or NVFP4. Keep the draft experts resident in FP8 while preserving the target model and standard rejection verification. Reduced proposal acceptance is allowed; the unchanged target distribution must be preserved.

Support both input routes:

• Opt-in online conversion of unquantized MTP experts with "mtp_expert_quantization": "fp8".
• Direct loading of serialized E4M3 MTP experts and their original 128x128 block scales, independently of the target's AWQ/NVFP4 quantization. Use the existing speculative model field for a separate draft checkpoint; ModelOpt mixed checkpoints and regular block-FP8 metadata are recognized.

This backports vllm-project/vllm#55513, merged as 60ad959b6f1a5c8f602edbd608c8decbc0788c50: ModelOpt FP8_PB_WO / FP8_BLOCK_SCALES dispatch and AMD/NVIDIA MTP layer-metadata remapping. The new work is SM70 resident execution, TP padding that preserves original FP8 bytes and block-scale offsets, and the independent online expert option. It does not depend on 1CatAI#553, whose SM70 fallback retains FP16 expert weights. vllm-project/vllm#55498 overlaps the already-merged generic loading fix.

On the tested TP4 shape, logical width 160 needs physical width 256. Native checkpoint loading preserves offsets 0/32/64/96 inside original scale blocks, without dequantizing and requantizing. The SM70 kernel uses FP16 scales; overflowing/underflowing scales are rejected. Online conversion is restricted to unquantized draft experts. SM70/FP16, Qwen4ExpMTP, ordinary TP without PP/EP, and standard rejection sampling are required.

Validation

Both loading routes completed full-model tests with AWQ and NVFP4 targets. This remains a fork draft for human review.

• CPU: .venv/bin/python -m pytest tests/models/qwen4_exp/test_mtp_fp8_experts.py tests/models/qwen4_exp/test_mtp_fp8_checkpoint.py -q: 48 passed. Includes both ModelOpt aliases, exclusions, metadata remapping, TP1/2/4/8 byte/scale preservation, incomplete checkpoints, the real allocator/TP loader with AWQ/NVFP4/mixed configuration names, and online conversion of unquantized experts omitted from mixed metadata.
• V100: .venv/bin/python -m pytest tests/models/qwen4_exp/test_mtp_fp8_experts_gpu.py -q: 20 passed, covering online and checkpoint-native methods, four TP4 offsets, M=1/2/8/64, reference reconstruction and exact CUDA Graph replay.
• Original NVIDIA checkpoint experts with original BF16 scales: four V100 processes × M=1/2/8/64 passed. Maximum absolute error 0.0008544921875, relative L2 below 0.00082; graph replay exact. These scales retain their values exactly in the kernel's FP16 representation.
• Checkpoint-native AWQ: same-source FP16-control versus original FP8 experts, 30992 to 29982 MiB/rank (1010 MiB saved). Both arms completed 16 complete greedy and 16 stochastic requests. Full choices matched in 14/16 cases and usage in 15/16; probability and list-versus-tuple explanations differed in wording while retaining their conclusions. The unchanged FP16 control also varied on a serial wording retest; original pair results are retained.
• Checkpoint-native NVFP4: same-source FP16-control versus original FP8 experts, 31960 to 31068 MiB/rank (892 MiB saved). 16/16 complete greedy choices and token usage identical, plus 16 successful finite-logprob stochastic requests per arm. Every rank retained 975 MiB of packed MTP expert weights/scales; original expert parameters were absent.
• NVFP4 online conversion: the target checkpoint's own unquantized MTP loaded and retained 975 MiB expert payload/rank, with no original FP16 expert parameters. Whole-GPU usage was 31048 MiB/rank after both campaigns. Completed 16 greedy and 16 finite-logprob stochastic requests; all greedy choices/usage matched the same NVFP4 target control. This uses a different MTP source from the native FP16 control, so it is an integration/output check rather than a same-source quantization comparison.
• Earlier online AWQ paired validation: fixed 5-GiB KV/rank, TP4/MTP3/C2, 16/16 final greedy responses identical to FP16 baseline, 16 stochastic requests per arm successful with finite logprobs. Actual packed expert weights/scales were 975 versus 1200 MiB/rank; whole-GPU usage was 29962 versus 30992 MiB/rank. The extra saving beyond 225 MiB of payload was not individually attributed.
• Existing rejection-sampler tests passed in the native runtime, including adversarial draft distributions at K1/K3 with 200,000 trials each, greedy verification and calibrated nucleus sampling.
• Applicable pre-commit hooks passed. These are local checks, not a claim of hosted CI success.

All paired full-model arms use TP4/MTP3/C2, FP16 execution/KV, CUDA Graphs and fixed 5-GiB KV/rank. Memory values are idle snapshots after both greedy and stochastic campaigns, identical across all four ranks. Whole-GPU reductions include runtime overhead and must not be attributed entirely to the 225-MiB expert-payload reduction. Actual test prompts reach approximately 7700 tokens; configured 262144 context is not a full-context validation.

The native runtime uses extensions built from 752f86495f, with changed Python modules overlaid from this branch based on fe67339ddf; it is not a clean native rebuild. The first checkpoint-native AWQ run preceded two additional input guards; the 20 GPU tests and NVFP4 runs include them. Main target weights, shared-head handling and rejection sampling are not modified. Tests are bounded correctness/integration evidence, not a universal bitwise or benchmark-quality guarantee. Standard rejection sampling preserves the unchanged target distribution using actual proposal probabilities. No throughput or loading-peak improvement is claimed.

See design and validation details. This remains a fork draft for human review; no production deployment or upstream submission is included.

AI assistance: Codex assisted with implementation, tests and documentation under the submitter's direction.