Load native FP4 experts with optional lossless CSF scales

[voipmonitor]

VLLM_B12X_MOE_FP4_CSF=1 loads ordinary native MXFP4/NVFP4 expert checkpoints with lossless CSF scale storage. The standard safetensors loader stages scale parameters on the CPU, prepares one layer on the GPU, and releases the original scale tensors. Model-scoped shared scratch avoids retaining full scale planes for every layer. The flag defaults to 0; checkpoint files, dense-layer quantization, and expert activation precision retain their existing behavior.

Scale staging covers the native MXFP4, ModelOpt NVFP4, and compressed-tensors MXFP4 loaders. Kimi normalizes its native metadata to the mxfp4 handler and uses --quantization mxfp4. Default compressed-tensors scale placement remains on the GPU when online compression is disabled.

This stacks on #956 and adds native-checkpoint preparation; it does not duplicate #956's stored-CSF loader or #963's beta import. It requires paired B12X #458. PP1/DP1, no EP/EPLB, and serialized execution are required; unsupported overlap is rejected. TP workers compress independently, and GPU preparation requires no extra thread setting.

Implemented and qualified for the tested staging, ownership, and bounded startup contracts:

• python -m pytest --confcutdir=tests/kernels/moe tests/kernels/moe/test_b12x.py -k 'compressed_tensors_mxfp4_preserves or online_csf' -xq: 6 GPU cases passed, including source-scale release, scratch lifetime/isolation, and enabled/disabled compressed-tensors placement.
• Applicable repository pre-commit hooks pass, including mypy, formatting, and accelerator API checks.
• Paired B12X codecs/index/MoE: 166 GPU cases and 24 compute-sanitizer memcheck cases passed. Complete prepared buffers match the CPU oracle, including real GLM checkpoint data; real Kimi scales reconstruct exactly.
• Native GLM TP2/no-spec: 8 successful generation requests per launch, 85.71 GiB/rank. The paired GPU index reduces preparation of 42 MoE layers to 2.38–2.40 s/rank including first compilation.
• Native Kimi TP16/DCP16/no-spec: 8 requests with finite logprobs and a 10,401-token prefill/64-token continuation passed; 86.23 GiB/rank, 92 MoE layers prepared in 4.35–5.34 s/rank. Expert activations remain BF16 and dense-MLA partials FP32. The tested runtime includes merged B12X [II] Preserve Kimi parallel-draft KV group identity #441 to provide that attention API.
• Native DeepSeek-V4-Flash-0731 TP2/no-spec (distinct from DS4.1): 43 MoE layers per rank compressed online, 71.89 GiB/rank, native MXFP8 expert activations preserved. Eight short probes, a 1024-token decode, uncached 8192/32,768-token prefills with continuation, and eight concurrent 128-token generations passed; all 19 requests completed with zero errors or preemptions. Preparation totals 8.010–8.036 s/rank including staging, weight preparation, and first compilation. DS4F conditions, identities, receipts, and limits. This is bounded execution validation; KLD and a native-storage throughput comparison are unmeasured.

Real DS4.1 TP4 (384 experts) and Qwen TP1 (512 experts), both projections of layer 3, pass exact source-byte reconstruction and baseline encoded-plane comparisons. Qwen also matches the complete CPU-index prepared storage. Eight balanced warm cycles measure Qwen encoding plus index at 105.394 → 2.923 ms (36.06x) and DS4.1 encoding at 2.833 → 2.834 ms (1.00x). DS4.1 uses 160/18 scale columns per row, so the at-most-eight-column MXFP4 optimization does not apply. These component timings are research-only under NVML 0x400, RTX PRO 6000 Blackwell, 600 W, memory 16,365 MHz. Complete DS4.1/Qwen startup receipts cover the separately pinned loading implementation; whole-model startup after the GPU-index change was not repeated.

Full-model KLD and full-vocabulary logit parity are unmeasured. Full-model GLM responses differ even between unchanged control launches; byte/component oracles establish exactness. Encoding timings are research-only under NVML event mask 0x400, and do not claim improved decode throughput or equal I/O/compiler-cache history.

Preparation measurements and exact runtime identities separate component speedups from complete startup. Native loading/serving evidence also records Qwen TP1/MTP3 and DS4.1 TP4/no-spec loading, serving comparisons, and configuration limits.

AI assistance: implemented, reviewed, and tested with OpenAI Codex at the repository owner's request. Human review has not been represented as complete.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

🗂️ Base branches to auto review (1)

• dev/*

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Advanced

Run ID: 30a3a6ca-3888-462e-a259-04afde2202dd

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

• Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Autopilot is currently an internal CodeRabbit preview.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}