MXFP4-CSF: load DeepSeek-V4.1-Flash checkpoints at TP3

[Defilan]

Purpose

The MXFP4-CSF reader admits DeepSeek-V4.1-Flash at TP1, TP2, TP4 and TP8 only, through the explicit TP allowlist in read_mxfp4_csf_layer. This adds TP3 for the deepseek_v41 family so the lossless CSF checkpoint can serve on three GPUs, for example three DGX Spark (GB10) nodes. The deepseek_v4_flash and kimi_k3 lists are unchanged.

Every other constraint on the path already holds at TP3. The routed-expert intermediate size is 2304, so a TP3 rank owns 768 channels: tp_extent needs 32-channel alignment (2304 is divisible by 96), the compressed scale-plane slicer needs whole 16-row slabs (rank boundaries 0, 768 and 1536 are multiples of 16), and the b12x compact W4A8 CSF geometry needs N64/K64 (w13 1536 x 5120 and w2 5120 x 768 per rank). Attention already handles TP3 through the fork's head padding (64 heads and 8 o_groups padded to 72 and 9 at runtime), so the checkpoint's config.json stays as published. DeepSeek-V4-Flash is left out on purpose: its 2048 channels do not split in three.

Not a duplicate: no open or closed PR in this repo covers MXFP4-CSF at TP3. The open TP3 PRs (#547, #601) are for GLM-5.3 models.

Release-note fragment: .lil/changes/vllm-mxfp4-csf-ds41-tp3.json.

Test Plan

• pre-commit run ruff-check and ruff-format on the changed files (ruff 0.14.0).
• .venv/bin/python -m pytest tests/quantization/test_mxfp4_csf.py -v. The existing byte-slicing test, test_tensor_sources_preserve_projection_order_and_tp_bytes, now also runs at TP3 (ranks 0, 1 and 2). It checks that each rank's nibbles and decoded compressed scale planes equal exactly its third of the source. It calls _load_mxfp4_csf_weights directly, below the TP guard. The new test_reader_admits_tp3_for_deepseek_v41_only checks that the reader gets past the TP guard at TP3 for deepseek_v41, and that TP6 for deepseek_v41 and TP3 for deepseek_v4_flash and kimi_k3 are rejected.
• Serving on hardware: local-inference-lab/DeepSeek-V4.1-Flash-MXFP4-CSF at revision 872da235, TP3 across three GB10 nodes on a switchless RoCE ring. The image is built from source at fork 1286ae9c (this branch's TP guard change applied as a one-line patch) and b12x 78ee52c3 plus [II] Align hybrid recurrent checkpoints with cache-group boundaries #457. The recipe is public: feat(cuda-gb10-vllm-dsv41-kk): karmic-kraken-beta vLLM + b12x for DSv4.1 TP3, fenced NCCL, MXFP4-CSF defilantech/llmkube-runtimes#56, variant cuda-gb10-vllm-dsv41-kk. Serving settings: --tensor-parallel-size 3, --quantization mxfp4_csf --load-format mxfp4_csf, 262K context, max-num-batched-tokens 4096, RoCEnante all-reduce on, DSpark off, KV cache 2.2 GiB per rank.

Test Result

• ruff check and format: clean.
• tests/quantization/test_mxfp4_csf.py on this branch, run on a GB10 (DGX Spark) in an image built from fork 1286ae9c and b12x 78ee52c3, with this branch's Python files under vllm/ overlaid (all pure Python): 72 passed. That includes the byte-slicing test at TP3 for ranks 0, 1 and 2, and all four cases of the new guard test. Negative control with the loader file from the base commit: 1 failed, 71 passed. The failure is the new test's TP3 deepseek_v41 case (ValueError: MXFP4-CSF deepseek_v41 supports TP (1, 2, 4, 8)). The byte-slicing cases pass on the base loader too, since _load_mxfp4_csf_weights sits below the TP guard; they show the slicer is correct at TP3, and the guard test shows the reader now admits it.
• Serving, DeepSeek-V4.1-Flash MXFP4-CSF at TP3 on three GB10s:
  • Loads at 90.0 GiB of weights per rank. KV pool is 1,394,380 tokens at 2.2 GiB per rank (5.32x at 262K). Ready in 509 s on the first boot (b12x kernel cache build) and 220 s on a cached restart.
  • Perplexity 2.9605 on our corpus of 200 sequences of 2048 tokens. On the same corpus the official MXFP4 checkpoint measured 2.9615 at TP3, and an EXL3 3.5 bpw build measured 2.9787.
  • Fixed 12-prompt battery (thinking off): 12/12.
  • LiveCodeBench subset (100 problems, seed 42, thinking off, max 8192 tokens): 84/100. Two runs of the EXL3 3.5 bpw build on the same subset scored 79 and 83.
  • Cold prefill: 2,678 tok/s at 10K tokens, 2,763 at 40K, and 2,668 at 100K. The EXL3 ring on the same nodes (DSpark on) measured 1,808, 1,777 and about 1,140.
  • Agent-style decode with DSpark off: 32.4 to 34.0 tok/s at concurrency 1, and 25.3 to 26.6 tok/s per stream at concurrency 2.

This PR was prepared with AI assistance (Claude Code), including the test and serving runs above on my lab hardware. I reviewed the change and the results, and I'm responsible for it.

Assisted-by: Claude Code

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[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: 44dc4cd0-b85b-4e9f-9684-1ecbc592f5ff

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

---

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.}

[voipmonitor]

Verified on beta e1ccbaf0 (overlaying this loader): DeepSeek-V4.1-Flash-lossless-CSF loads at TP3 on 3x RTX PRO 6000 (weights loaded in 38 s), then OOMs because the model itself needs ~93 GiB per GPU at TP3 on 96 GB cards (71.4 GiB at TP4); not related to this change. Fine for larger-memory nodes such as GB10. Unit tests and fragment (validated with dk release_changelog) look good.