From 1dc03737776e6a07f58b608c6fcd4ec7b0a01c5f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 29 Sep 2026 23:14:36 +0000 Subject: [PATCH 1/3] feat(ds41): load lossless X4T experts through B12X Register exact_mxfp4 loading for compressed expert scales while retaining native dense and draft precision and file-backed Engram tables. Reuse serialized scale scratch and reject unsupported parallel execution modes. Preserve EXL3 prepared-weight validation and validate the SwiGLU clamp. Validation: 122 X4T, EXL3, and loader-registry tests pass with B12X #446. New code passes mypy; the five existing model.py findings reproduce on the KK base. Full-model serving evidence is linked in the companion B12X PR. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/quantization/test_ds41_exact_mxfp4.py | 217 +++++++++++++++++ vllm/config/load.py | 2 + vllm/config/model.py | 1 + vllm/model_executor/layers/fused_moe/b12x.py | 26 +- .../layers/quantization/__init__.py | 3 + vllm/model_executor/model_loader/__init__.py | 3 + .../model_loader/exact_mxfp4_loader.py | 59 +++++ vllm/models/deepseek_v4_1/exact_mxfp4.py | 229 ++++++++++++++++++ vllm/models/deepseek_v4_1/nvidia/model.py | 5 +- 9 files changed, 535 insertions(+), 10 deletions(-) create mode 100644 tests/quantization/test_ds41_exact_mxfp4.py create mode 100644 vllm/model_executor/model_loader/exact_mxfp4_loader.py create mode 100644 vllm/models/deepseek_v4_1/exact_mxfp4.py diff --git a/tests/quantization/test_ds41_exact_mxfp4.py b/tests/quantization/test_ds41_exact_mxfp4.py new file mode 100644 index 000000000000..2de6bcd4daa3 --- /dev/null +++ b/tests/quantization/test_ds41_exact_mxfp4.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from safetensors.torch import save_file + +from vllm.config.load import LoadConfig +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.model_loader.exact_mxfp4_loader import ExactMXFP4ModelLoader +from vllm.models.deepseek_v4_1.exact_mxfp4 import ( + DeepseekV41ExactMXFP4Config, + ExactMXFP4MoEMethod, +) + + +@pytest.fixture +def moe(): + return FusedMoEConfig( + num_experts=384, + num_local_experts=384, + num_logical_experts=384, + experts_per_token=6, + hidden_dim=5120, + intermediate_size=576, + in_dtype=torch.bfloat16, + device="cpu", + activation=MoEActivation.SILU, + swiglu_limit=10.0, + routing_method=RoutingMethodType.TopK, + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + ) + + +@pytest.fixture +def owner(): + return DeepseekV41ExactMXFP4Config.from_config( + {"format_version": 1, "checkpoint_root": "/x4t"} + ) + + +def test_exact_config_keeps_native_dense_precision_and_empty_expert_handles(moe, owner): + assert owner.is_scale_e8m0 and owner.weight_block_size == [32, 32] + assert owner.is_checkpoint_fp8_serialized and owner.scale_scratch is None + method = ExactMXFP4MoEMethod(moe, owner) + layer = torch.nn.Module() + layer.layer_name = "model.layers.17.ffn.experts" + method.create_weights(layer, 384, 5120, 576, torch.bfloat16) + assert not list(layer.named_parameters()) and not layer.state_dict() + assert method.layer_index == 17 and method.local_intermediate == 576 + assert method.get_fused_moe_quant_config(layer).weight_quant_dtype == "mxfp4" + with pytest.raises(ValueError, match="format_version"): + DeepseekV41ExactMXFP4Config.from_config({}) + + +def test_exact_config_and_loader_are_registered(): + from vllm.model_executor.layers.quantization import get_quantization_config + from vllm.model_executor.model_loader import get_model_loader + + assert get_quantization_config("exact_mxfp4") is DeepseekV41ExactMXFP4Config + assert isinstance( + get_model_loader(LoadConfig(load_format="exact_mxfp4")), ExactMXFP4ModelLoader + ) + + +@pytest.mark.parametrize("mode", ["normal", "pipeline", "ubatching", "wrong_loader"]) +def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( + monkeypatch, moe, owner, mode +): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.models.deepseek_v4_1 import exact_mxfp4 + + config = SimpleNamespace( + model_config=SimpleNamespace(hf_config=SimpleNamespace(num_hidden_layers=40)), + parallel_config=SimpleNamespace( + pipeline_parallel_size=2 if mode == "pipeline" else 1, + use_ubatching=mode == "ubatching", + ), + load_config=SimpleNamespace( + load_format="safetensors" if mode == "wrong_loader" else "exact_mxfp4" + ), + ) + monkeypatch.setattr(exact_mxfp4, "get_current_vllm_config", lambda: config) + native = object() + monkeypatch.setattr( + exact_mxfp4.DeepseekV41FP8Config, + "get_quant_method", + lambda *args, **kwargs: native, + ) + layer = Mock(spec=RoutedExperts) + layer.moe_config = moe + if mode == "normal": + assert isinstance( + owner.get_quant_method(layer, "model.layers.39.ffn.experts"), + ExactMXFP4MoEMethod, + ) + elif mode == "wrong_loader": + with pytest.raises(ValueError, match="load-format"): + owner.get_quant_method(layer, "model.layers.39.ffn.experts") + else: + with pytest.raises(NotImplementedError, match="PP1 without ubatching"): + owner.get_quant_method(layer, "model.layers.39.ffn.experts") + assert owner.get_quant_method(layer, "model.layers.40.ffn.experts") is native + assert owner.get_quant_method(torch.nn.Module(), "model.layers.0.attn") is native + + +@pytest.mark.parametrize( + "field,value", [("dp_size", 2), ("use_ep", True), ("enable_eplb", True)] +) +def test_exact_experts_reject_non_tp_execution(moe, owner, field, value): + parallel = replace(moe.moe_parallel_config, **{field: value}) + with pytest.raises(NotImplementedError, match="TP without EP/DP"): + ExactMXFP4MoEMethod(replace(moe, moe_parallel_config=parallel), owner) + + +def test_prepared_native_experts_validate_source_shape_and_swiglu_limit( + monkeypatch, moe, owner +): + from b12x.moe import fused_moe + + from vllm.model_executor.layers.fused_moe import b12x + + layer = torch.nn.Module() + layer.activation = moe.activation + layer.apply_router_weight_on_input = False + quant = ExactMXFP4MoEMethod(moe, owner).get_fused_moe_quant_config(layer) + backend = b12x.B12xExperts(moe, quant) + prepared = Mock(spec=fused_moe.PreparedExperts) + activation = fused_moe.ActivationSpec( + mode="a16", nonlinearity="silu", io_dtype=torch.bfloat16, swiglu_limit=10.0 + ) + prepared.plan = SimpleNamespace( + source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), + activation=activation, + ) + prepared.num_experts = moe.num_experts + prepared.hidden_size = moe.hidden_dim + prepared.intermediate_size = moe.intermediate_size_per_partition + monkeypatch.setattr( + b12x, "_register_b12x_moe_output_collective", lambda *a, **k: None + ) + backend.install_prepared_experts(layer, prepared) + assert layer._b12x_prepared_experts is prepared + assert backend._prepared_experts is prepared + prepared.plan.activation = replace(activation, swiglu_limit=None) + with pytest.raises(ValueError, match="activation"): + backend.install_prepared_experts(layer, prepared) + prepared.plan.activation = activation + prepared.hidden_size += 128 + with pytest.raises(ValueError, match="geometry"): + backend.install_prepared_experts(layer, prepared) + prepared.hidden_size = moe.hidden_dim + prepared.plan.source = fused_moe.PackedSource(format="modelopt_nvfp4") + with pytest.raises(TypeError, match="encoding"): + backend.install_prepared_experts(layer, prepared) + + +def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): + from b12x.moe.checkpoints.exact_mxfp4 import CODEC, SCHEMA + + tensor_dir = tmp_path / "tensors" + tensor_dir.mkdir() + name = "model-00001.safetensors" + tensors = { + "layers.0.ffn.experts.0.w1.weight": torch.ones(8, dtype=torch.uint8), + "layers.0.ffn.experts.0.w1.scale.exact_mxfp4_fixed": torch.ones( + 8, dtype=torch.uint8 + ), + "layers.40.ffn.experts.0.w1.weight": torch.full((8,), 42, dtype=torch.uint8), + "layers.1.engram.embedding.weight": torch.ones((8, 8), dtype=torch.uint8), + "norm.weight": torch.ones(8, dtype=torch.bfloat16), + } + save_file(tensors, tensor_dir / name) + common = {"schema": SCHEMA, "codec": CODEC, "family": "deepseek_v41"} + (tmp_path / "manifest.json").write_text( + json.dumps({**common, "shards": [{"file": name}]}) + ) + (tmp_path / "build-contract.json").write_text( + json.dumps({**common, "source_names": {k: name for k in tensors}}) + ) + config = SimpleNamespace( + hf_config=SimpleNamespace( + num_hidden_layers=40, + quantization_config={ + "quant_method": "exact_mxfp4", + "checkpoint_root": str(tmp_path), + }, + ) + ) + model = SimpleNamespace( + checkpoint_file_weight_filter=lambda n: ".engram.embedding." in n + ) + loader = ExactMXFP4ModelLoader(LoadConfig(load_format="exact_mxfp4")) + result = dict(loader.get_all_weights(config, model)) + assert len(result) == 3 + assert result["layers.1.engram.embedding.weight"].device.type == "meta" + assert result[ + "layers.1.engram.embedding.weight" + ]._vllm_file_tensor_source.shape == (8, 8) + assert torch.equal( + result["layers.40.ffn.experts.0.w1.weight"], + tensors["layers.40.ffn.experts.0.w1.weight"], + ) + model.checkpoint_weight_name_prefixes = ("layers.40.",) + assert list(dict(loader.get_all_weights(config, model))) == [ + "layers.40.ffn.experts.0.w1.weight" + ] diff --git a/vllm/config/load.py b/vllm/config/load.py index 5e24a617ba85..6729f6820c45 100644 --- a/vllm/config/load.py +++ b/vllm/config/load.py @@ -35,6 +35,8 @@ class LoadConfig: back to the pytorch bin format if safetensors format is not available. - "pt" will load the weights in the pytorch bin format. - "safetensors" will load the weights in the safetensors format. + - "exact_mxfp4" reads losslessly compressed DS4.1 MXFP4 expert scales and + retains file-backed loading for Engram tables. - "instanttensor" will load the Safetensors weights on CUDA devices using InstantTensor, which enables distributed loading with pipelined prefetching and fast direct I/O. diff --git a/vllm/config/model.py b/vllm/config/model.py index 3cb01da228c3..30df9ca07aea 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1284,6 +1284,7 @@ def _verify_quantization(self) -> None: "mxfp8", "modelopt_mixed", "exl3", + "exact_mxfp4", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) "mxfp4", diff --git a/vllm/model_executor/layers/fused_moe/b12x.py b/vllm/model_executor/layers/fused_moe/b12x.py index af958dca03ba..7f9f129ebe4c 100644 --- a/vllm/model_executor/layers/fused_moe/b12x.py +++ b/vllm/model_executor/layers/fused_moe/b12x.py @@ -499,28 +499,36 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: ) def install_prepared_experts(self, layer: torch.nn.Module, prepared: Any) -> None: - """Install EXL3 weights prepared through the common trellis API.""" + """Install canonical prepared experts without retaining source parameters.""" fused_moe = _require_b12x_fused_moe() - if ( - self._source_format != "exl3" - or not isinstance(prepared, fused_moe.PreparedExperts) - or not isinstance(prepared.plan.source, fused_moe.TrellisSource) - ): - raise TypeError("EXL3 installation requires B12X prepared trellis weights") + if not isinstance(prepared, fused_moe.PreparedExperts): + raise TypeError("Installation requires B12X PreparedExperts") + trellis = self._source_format == "exl3" and isinstance( + prepared.plan.source, fused_moe.TrellisSource + ) + packed = ( + self._source_format == "fp4_e8m0_k32" + and isinstance(prepared.plan.source, fused_moe.PackedSource) + and prepared.plan.source.format == "fp4_e8m0_k32" + ) + if not (trellis or packed): + raise TypeError("Prepared expert encoding does not match the backend") if ( prepared.num_experts != self.moe_config.num_experts or prepared.hidden_size != self.moe_config.hidden_dim or prepared.intermediate_size != self.moe_config.intermediate_size_per_partition or prepared.plan.activation.mode != "a16" - or prepared.plan.activation.rotation_dtype != torch.float16 + or (trellis and prepared.plan.activation.rotation_dtype != torch.float16) or prepared.plan.activation.io_dtype != self.moe_config.in_dtype or prepared.plan.activation.nonlinearity != _b12x_activation_name(layer.activation) + or prepared.plan.activation.swiglu_limit != self.moe_config.swiglu_limit or layer.apply_router_weight_on_input ): raise ValueError( - "Prepared EXL3 geometry, activation or routing does not match the layer" + "Prepared expert geometry, activation or routing " + "does not match the layer" ) self._apply_router_weight_on_input = False self._reuse_prepared_storage(layer, prepared) diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index c5e0f1446de1..95e79137d9d7 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -38,6 +38,7 @@ "gpt_oss_mxfp4", "deepseek_v4_fp8", "deepseek_v41_fp8", + "exact_mxfp4", "online", # Below are online quant shorthand names (see vllm.config.quantization). # Listed here as strings to avoid a circular import; kept in sync with @@ -132,6 +133,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: deepseek_config = DeepseekV4FP8Config + from vllm.models.deepseek_v4_1.exact_mxfp4 import DeepseekV41ExactMXFP4Config from vllm.models.deepseek_v4_1.quant_config import DeepseekV41FP8Config from .auto_awq import AutoAWQConfig @@ -182,6 +184,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "gpt_oss_mxfp4": GptOssMxfp4Config, "deepseek_v4_fp8": deepseek_config, "deepseek_v41_fp8": DeepseekV41FP8Config, + "exact_mxfp4": DeepseekV41ExactMXFP4Config, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/model_loader/__init__.py b/vllm/model_executor/model_loader/__init__.py index 6305c45f5d0b..d542fc862463 100644 --- a/vllm/model_executor/model_loader/__init__.py +++ b/vllm/model_executor/model_loader/__init__.py @@ -11,6 +11,7 @@ from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.exact_mxfp4_loader import ExactMXFP4ModelLoader from vllm.model_executor.model_loader.modelexpress_loader import ( ModelExpressModelLoader, ) @@ -36,6 +37,7 @@ "dummy", "fastsafetensors", "instanttensor", + "exact_mxfp4", "ipc_cache", "mistral", "modelexpress", @@ -53,6 +55,7 @@ "dummy": DummyModelLoader, "fastsafetensors": DefaultModelLoader, "instanttensor": DefaultModelLoader, + "exact_mxfp4": ExactMXFP4ModelLoader, "ipc_cache": IpcModelLoader, "mistral": DefaultModelLoader, "modelexpress": ModelExpressModelLoader, diff --git a/vllm/model_executor/model_loader/exact_mxfp4_loader.py b/vllm/model_executor/model_loader/exact_mxfp4_loader.py new file mode 100644 index 000000000000..1a81c4b703f1 --- /dev/null +++ b/vllm/model_executor/model_loader/exact_mxfp4_loader.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Read retained native tensors from a lossless MXFP4 container. + +Compressed routed experts are owned by the model's quantization method. +Engram tensors retain their immutable file-range descriptors, avoiding full +CPU/GPU staging allocations for the host-mapped embedding tables. +""" + +import time +from pathlib import Path + +import regex as re +from safetensors import safe_open + +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.model_executor.model_loader.weight_utils import ( + file_source_tensor, + safetensors_file_sources, +) + + +class ExactMXFP4ModelLoader(DefaultModelLoader): + def _root(self, model_config): + from b12x.moe.checkpoints.exact_mxfp4 import checkpoint_contract + + quant = model_config.hf_config.quantization_config + if quant.get("quant_method") != "exact_mxfp4": + raise ValueError("exact_mxfp4 loading requires the X4T model config") + root = Path(quant["checkpoint_root"]) + if not root.is_absolute(): + raise ValueError("X4T checkpoint_root must be an absolute local path") + return root, checkpoint_contract(str(root.resolve())) + + def download_model(self, model_config): + self._root(model_config) + + def get_all_weights(self, model_config, model): + root, contract = self._root(model_config) + if getattr(model, "secondary_weights", ()): + raise NotImplementedError("X4T does not support secondary weight sources") + file_filter = getattr(model, "checkpoint_file_weight_filter", None) + prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) + layers = model_config.hf_config.num_hidden_layers + self.counter_before_loading_weights = time.perf_counter() + for filename in sorted(set(contract["source_names"].values())): + path = str(root / "tensors" / filename) + descriptors = safetensors_file_sources(path) + with safe_open(path, framework="pt", device="cpu") as handle: + for name in sorted(descriptors): + if prefixes is not None and not name.startswith(prefixes): + continue + match = re.match(r"layers\.(\d+)\.ffn\.experts\.", name) + if match and int(match.group(1)) < layers: + continue + if callable(file_filter) and file_filter(name): + yield name, file_source_tensor(descriptors[name]) + else: + yield name, handle.get_tensor(name) diff --git a/vllm/models/deepseek_v4_1/exact_mxfp4.py b/vllm/models/deepseek_v4_1/exact_mxfp4.py new file mode 100644 index 000000000000..1e4536b4ae25 --- /dev/null +++ b/vllm/models/deepseek_v4_1/exact_mxfp4.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Losslessly compressed DS4.1 routed experts with native dense tensors.""" + +import regex as re +import torch + +from vllm.config import get_current_vllm_config +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import modular_kernel as mk +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.b12x import B12xExperts +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEQuantConfig, + FusedMoEQuantDesc, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.prepare_finalize.no_dp_ep import ( + MoEPrepareAndFinalizeNoDPEPModular, +) +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + +from .quant_config import DeepseekV41FP8Config + +logger = init_logger(__name__) + + +class DeepseekV41ExactMXFP4Config(DeepseekV41FP8Config): + """X4T main experts; unchanged native dense, shared and draft weights.""" + + checkpoint_root: str + scale_scratch: tuple[torch.Tensor, ...] | None + + @classmethod + def get_name(cls): + return "exact_mxfp4" + + @classmethod + def get_min_capability(cls): + return 120 + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + if ( + user_quant in (None, cls.get_name()) + and hf_quant_cfg is not None + and hf_quant_cfg.get("quant_method") == cls.get_name() + and getattr(hf_config, "model_type", None) + in ("deepseek_v41", "deepseek_v41_text") + ): + return cls.get_name() + return None + + @classmethod + def from_config(cls, config): + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("X4T requires format_version=1 and checkpoint_root") + result = cls( + is_checkpoint_fp8_serialized=True, + activation_scheme="dynamic", + weight_block_size=[32, 32], + ) + result.checkpoint_root = config["checkpoint_root"] + result.scale_scratch = None + return result + + def get_quant_method(self, layer, prefix): + if isinstance(layer, RoutedExperts): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", prefix) + if match is None: + raise ValueError("DS4.1 routed experts require a numbered layer") + config = get_current_vllm_config() + if int(match.group(1)) < config.model_config.hf_config.num_hidden_layers: + if ( + config.parallel_config.pipeline_parallel_size != 1 + or config.parallel_config.use_ubatching + ): + raise NotImplementedError( + "X4T shared scale scratch requires PP1 without ubatching" + ) + if config.load_config.load_format != "exact_mxfp4": + raise ValueError("X4T requires --load-format exact_mxfp4") + return ExactMXFP4MoEMethod(layer.moe_config, self) + return super().get_quant_method(layer, prefix) + + +class ExactMXFP4MoEMethod(FusedMoEMethodBase): + """Expand scales for routed experts into serialized, model-owned scratch.""" + + def __init__(self, moe, owner): + super().__init__(moe) + self.owner = owner + parallel = moe.moe_parallel_config + if ( + parallel.use_ep + or parallel.ep_size != 1 + or parallel.dp_size != 1 + or parallel.use_all2all_kernels + or parallel.enable_eplb + ): + raise NotImplementedError("X4T DS4.1 supports TP without EP/DP") + if ( + moe.activation != MoEActivation.SILU + or moe.in_dtype != torch.bfloat16 + or moe.has_bias + ): + raise ValueError("X4T DS4.1 requires bias-free BF16 SwiGLU experts") + + def create_weights( + self, + layer, + num_experts, + hidden_size, + intermediate_size_per_partition, + params_dtype, + **extra_weight_attrs, + ): + match = re.search(r"(?:^|\.)layers\.(\d+)\.", layer.layer_name) + if match is None or params_dtype != torch.bfloat16: + raise ValueError("X4T requires a numbered BF16 expert layer") + self.layer_index = int(match.group(1)) + self.num_experts, self.hidden_size = num_experts, hidden_size + self.local_intermediate = intermediate_size_per_partition + for name in ("w13_weight", "w2_weight"): + layer.register_buffer( + name, torch.empty(0, dtype=torch.uint8), persistent=False + ) + + def get_fused_moe_quant_config(self, layer): + return FusedMoEQuantConfig( + _a1=FusedMoEQuantDesc(), + _a2=FusedMoEQuantDesc(), + _w1=FusedMoEQuantDesc(dtype="mxfp4"), + _w2=FusedMoEQuantDesc(dtype="mxfp4"), + ) + + def process_weights_after_loading(self, layer): + from b12x.moe import fused_moe + from b12x.moe.checkpoints.exact_mxfp4 import read_exact_mxfp4_layer + + tp, rank = ( + get_tensor_model_parallel_world_size(), + get_tensor_model_parallel_rank(), + ) + device = layer.w13_weight.device + e, h, n = self.num_experts, self.hidden_size, self.local_intermediate + shapes = ((e, h // 32, 2 * n), (e, n // 32, h)) + if self.owner.scale_scratch is None: + # DS4.1 disallows ubatching: every layer consumes these scale grids + # before its successor may overwrite them on the same stream. + self.owner.scale_scratch = tuple( + torch.empty(s, dtype=torch.uint8, device=device) for s in shapes + ) + scratch = self.owner.scale_scratch + if any( + t.shape != shape or t.device != device for t, shape in zip(scratch, shapes) + ): + raise ValueError("X4T shared scale scratch geometry/device mismatch") + weights = read_exact_mxfp4_layer( + self.owner.checkpoint_root, + self.layer_index, + num_experts=e, + hidden_size=h, + intermediate_size=n * tp, + tp_rank=rank, + tp_size=tp, + device=device, + w13_scale_scratch=scratch[0], + w2_scale_scratch=scratch[1], + ) + plan = fused_moe.plan_weights( + source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), + activation=fused_moe.ActivationSpec( + mode="a16", + nonlinearity="silu", + io_dtype=torch.bfloat16, + swiglu_limit=self.moe.swiglu_limit, + ), + geometry=fused_moe.MoEGeometry( + num_experts=e, hidden_size=h, intermediate_size=n + ), + ) + prepared = fused_moe.prepare_weights(plan=plan, weights=weights) + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + backend = B12xExperts(self.moe, self.moe_quant_config) + backend.install_prepared_experts(layer, prepared) + self.moe_kernel = mk.FusedMoEKernel( + MoEPrepareAndFinalizeNoDPEPModular(), backend + ) + logger.info( + "X4T lossless MXFP4 layer %d rank %d/%d: BF16 activations, " + "compressed scales, shared scratch %d bytes", + self.layer_index, + rank, + tp, + sum(t.numel() for t in scratch), + ) + + def apply( + self, + layer, + x, + topk_weights, + topk_ids, + shared_experts, + shared_experts_input, + workspace=None, + ): + assert self.moe_kernel is not None + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + workspace=workspace, + ) diff --git a/vllm/models/deepseek_v4_1/nvidia/model.py b/vllm/models/deepseek_v4_1/nvidia/model.py index f2323ae3c915..f4b283eec113 100644 --- a/vllm/models/deepseek_v4_1/nvidia/model.py +++ b/vllm/models/deepseek_v4_1/nvidia/model.py @@ -369,7 +369,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) if vllm_config.lora_config is not None: raise ValueError("V4.1 native kernels do not support LoRA adapters") - if quant_config is None or quant_config.get_name() != "deepseek_v41_fp8": + if quant_config is None or quant_config.get_name() not in ( + "deepseek_v41_fp8", + "exact_mxfp4", + ): raise ValueError("V4.1 requires its native block32 quantization config") self.vocab_size = config.vocab_size self.hc_eps = config.hc_eps From c390adba0d35a4b91a83c9f8fd0d6b2475f7f5e4 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 14:23:14 +0000 Subject: [PATCH 2/3] Load Kimi-K3 X4T experts through the lossless MXFP4 reader Preserve BF16 dense weights, select SiTU expert activation, recognize nested checkpoint metadata and bound CPU shard-copy threads. Shared scale scratch remains serialized; PP and ubatching are rejected. Port the implementation qualified on TP16 without changing native MXFP4 arithmetic. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/quantization/test_ds41_exact_mxfp4.py | 97 +++++++++++++++++++ vllm/config/model.py | 1 + .../layers/quantization/__init__.py | 3 + .../layers/quantization/kimi_x4t.py | 58 +++++++++++ .../model_loader/exact_mxfp4_loader.py | 14 ++- vllm/models/deepseek_v4_1/exact_mxfp4.py | 45 ++++++--- vllm/models/kimi_k3/nvidia/model.py | 2 +- 7 files changed, 199 insertions(+), 21 deletions(-) create mode 100644 vllm/model_executor/layers/quantization/kimi_x4t.py diff --git a/tests/quantization/test_ds41_exact_mxfp4.py b/tests/quantization/test_ds41_exact_mxfp4.py index 2de6bcd4daa3..5b855fe478d9 100644 --- a/tests/quantization/test_ds41_exact_mxfp4.py +++ b/tests/quantization/test_ds41_exact_mxfp4.py @@ -73,6 +73,32 @@ def test_exact_config_and_loader_are_registered(): ) +def test_x4t_cpu_shard_reader_bounds_threads_and_restores_on_failure( + monkeypatch, moe, owner +): + from b12x.moe.checkpoints import exact_mxfp4 as checkpoint + + from vllm.models.deepseek_v4_1 import exact_mxfp4 + from vllm.utils.torch_utils import set_default_torch_num_threads + + monkeypatch.setattr(exact_mxfp4, "get_tensor_model_parallel_world_size", lambda: 4) + monkeypatch.setattr(exact_mxfp4, "get_tensor_model_parallel_rank", lambda: 0) + method = ExactMXFP4MoEMethod(moe, owner) + layer = torch.nn.Module() + layer.layer_name = "model.layers.17.ffn.experts" + method.create_weights(layer, 384, 5120, 576, torch.bfloat16) + + def fail_reader(*args, **kwargs): + assert torch.get_num_threads() == 1 + raise OSError("checkpoint read failed") + + monkeypatch.setattr(checkpoint, "read_exact_mxfp4_layer", fail_reader) + with set_default_torch_num_threads(4): + with pytest.raises(OSError, match="checkpoint read failed"): + method.process_weights_after_loading(layer) + assert torch.get_num_threads() == 4 + + @pytest.mark.parametrize("mode", ["normal", "pipeline", "ubatching", "wrong_loader"]) def test_target_uses_x4t_while_dense_and_draft_keep_native_methods( monkeypatch, moe, owner, mode @@ -215,3 +241,74 @@ def test_loader_preserves_file_backed_engram_and_native_draft(tmp_path): assert list(dict(loader.get_all_weights(config, model))) == [ "layers.40.ffn.experts.0.w1.weight" ] + + +def test_kimi_x4t_keeps_bf16_dense_and_situ_experts(monkeypatch, moe): + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod + from vllm.model_executor.layers.quantization import kimi_x4t + + owner = kimi_x4t.KimiX4TConfig.from_config( + {"format_version": 1, "checkpoint_root": "/x4t"} + ) + config = SimpleNamespace( + parallel_config=SimpleNamespace(pipeline_parallel_size=1, use_ubatching=False), + load_config=SimpleNamespace(load_format="exact_mxfp4"), + ) + monkeypatch.setattr(kimi_x4t, "get_current_vllm_config", lambda: config) + experts = Mock(spec=RoutedExperts) + experts.moe_config = replace( + moe, + activation=MoEActivation.SITU, + activation_situ_beta=4.0, + activation_situ_linear_beta=25.0, + swiglu_limit=None, + ) + assert isinstance( + owner.get_quant_method(experts, "model.layers.1.mlp.experts"), + ExactMXFP4MoEMethod, + ) + assert isinstance( + owner.get_quant_method(Mock(spec=LinearBase), "model.layers.0.q_proj"), + UnquantizedLinearMethod, + ) + config.parallel_config.use_ubatching = True + with pytest.raises(NotImplementedError, match="PP1 without ubatching"): + owner.get_quant_method(experts, "model.layers.1.mlp.experts") + + +def test_kimi_loader_skips_compressed_experts_with_nested_text_config(tmp_path): + from b12x.moe.checkpoints.exact_mxfp4 import CODEC, SCHEMA + + tensor_dir = tmp_path / "tensors" + tensor_dir.mkdir() + filename = "model-00001.safetensors" + expert = "language_model.model.layers.1.block_sparse_moe.experts.0.w1." + dense = "language_model.model.layers.0.mlp.gate_proj.weight" + tensors = { + expert + "weight_packed": torch.ones(8, dtype=torch.uint8), + expert + "weight_scale.exact_mxfp4_fixed": torch.ones(8, dtype=torch.uint8), + dense: torch.ones((8, 8), dtype=torch.bfloat16), + } + save_file(tensors, tensor_dir / filename) + common = {"schema": SCHEMA, "codec": CODEC, "family": "kimi_k3"} + (tmp_path / "manifest.json").write_text( + json.dumps({**common, "shards": [{"file": filename}]}) + ) + (tmp_path / "build-contract.json").write_text( + json.dumps({**common, "source_names": {k: filename for k in tensors}}) + ) + config = SimpleNamespace( + hf_config=SimpleNamespace(), + hf_text_config=SimpleNamespace( + num_hidden_layers=93, + quantization_config={ + "quant_method": "kimi_x4t", + "checkpoint_root": str(tmp_path), + }, + ), + ) + loader = ExactMXFP4ModelLoader(LoadConfig(load_format="exact_mxfp4")) + result = dict(loader.get_all_weights(config, SimpleNamespace())) + assert list(result) == [dense] + assert torch.equal(result[dense], tensors[dense]) diff --git a/vllm/config/model.py b/vllm/config/model.py index 30df9ca07aea..79b540db676d 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1285,6 +1285,7 @@ def _verify_quantization(self) -> None: "modelopt_mixed", "exl3", "exact_mxfp4", + "kimi_x4t", # Ensure heavy backends are probed last to avoid unnecessary # imports during override detection (e.g., MXFP4 imports Triton) "mxfp4", diff --git a/vllm/model_executor/layers/quantization/__init__.py b/vllm/model_executor/layers/quantization/__init__.py index 95e79137d9d7..e8ca734d6a74 100644 --- a/vllm/model_executor/layers/quantization/__init__.py +++ b/vllm/model_executor/layers/quantization/__init__.py @@ -39,6 +39,7 @@ "deepseek_v4_fp8", "deepseek_v41_fp8", "exact_mxfp4", + "kimi_x4t", "online", # Below are online quant shorthand names (see vllm.config.quantization). # Listed here as strings to avoid a circular import; kept in sync with @@ -148,6 +149,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: from .fp_quant import FPQuantConfig from .humming import HummingConfig from .inc import INCConfig + from .kimi_x4t import KimiX4TConfig from .modelopt import ( ModelOptFp8Config, ModelOptMixedPrecisionConfig, @@ -185,6 +187,7 @@ def get_quantization_config(quantization: str) -> type[QuantizationConfig]: "deepseek_v4_fp8": deepseek_config, "deepseek_v41_fp8": DeepseekV41FP8Config, "exact_mxfp4": DeepseekV41ExactMXFP4Config, + "kimi_x4t": KimiX4TConfig, "humming": HummingConfig, "online": OnlineQuantizationConfig, # MiniMax-style checkpoints tag `quant_method: "mxfp8"`; load with the diff --git a/vllm/model_executor/layers/quantization/kimi_x4t.py b/vllm/model_executor/layers/quantization/kimi_x4t.py new file mode 100644 index 000000000000..bfb5ab45b01b --- /dev/null +++ b/vllm/model_executor/layers/quantization/kimi_x4t.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Kimi-K3 lossless MXFP4 expert scales with unchanged BF16 dense weights.""" + +import torch + +from vllm.config import get_current_vllm_config +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4Config +from vllm.models.deepseek_v4_1.exact_mxfp4 import ExactMXFP4MoEMethod + + +class KimiX4TConfig(Mxfp4Config): + checkpoint_root: str + scale_scratch: tuple[torch.Tensor, ...] | None + + @classmethod + def get_name(cls): + return "kimi_x4t" + + @classmethod + def get_min_capability(cls): + return 120 + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant, hf_config=None): + if ( + user_quant in (None, cls.get_name()) + and hf_quant_cfg is not None + and hf_quant_cfg.get("quant_method") == cls.get_name() + and getattr(hf_config, "model_type", None) in ("kimi_k3", "kimi_linear") + ): + return cls.get_name() + return None + + @classmethod + def from_config(cls, config): + if config.get("format_version") != 1 or not config.get("checkpoint_root"): + raise ValueError("Kimi X4T requires format_version=1 and checkpoint_root") + result = cls() + result.checkpoint_root = config["checkpoint_root"] + result.scale_scratch = None + return result + + def get_quant_method(self, layer, prefix): + if isinstance(layer, RoutedExperts): + config = get_current_vllm_config() + if ( + config.parallel_config.pipeline_parallel_size != 1 + or config.parallel_config.use_ubatching + ): + raise NotImplementedError( + "X4T shared scale scratch requires PP1 without ubatching" + ) + if config.load_config.load_format != "exact_mxfp4": + raise ValueError("Kimi X4T requires --load-format exact_mxfp4") + return ExactMXFP4MoEMethod(layer.moe_config, self) + return super().get_quant_method(layer, prefix) diff --git a/vllm/model_executor/model_loader/exact_mxfp4_loader.py b/vllm/model_executor/model_loader/exact_mxfp4_loader.py index 1a81c4b703f1..9fa45d9dc206 100644 --- a/vllm/model_executor/model_loader/exact_mxfp4_loader.py +++ b/vllm/model_executor/model_loader/exact_mxfp4_loader.py @@ -24,8 +24,10 @@ class ExactMXFP4ModelLoader(DefaultModelLoader): def _root(self, model_config): from b12x.moe.checkpoints.exact_mxfp4 import checkpoint_contract - quant = model_config.hf_config.quantization_config - if quant.get("quant_method") != "exact_mxfp4": + text_config = getattr(model_config, "hf_text_config", model_config.hf_config) + quant = getattr(model_config.hf_config, "quantization_config", None) + quant = quant or text_config.quantization_config + if quant.get("quant_method") not in ("exact_mxfp4", "kimi_x4t"): raise ValueError("exact_mxfp4 loading requires the X4T model config") root = Path(quant["checkpoint_root"]) if not root.is_absolute(): @@ -41,7 +43,8 @@ def get_all_weights(self, model_config, model): raise NotImplementedError("X4T does not support secondary weight sources") file_filter = getattr(model, "checkpoint_file_weight_filter", None) prefixes = getattr(model, "checkpoint_weight_name_prefixes", None) - layers = model_config.hf_config.num_hidden_layers + text_config = getattr(model_config, "hf_text_config", model_config.hf_config) + layers = text_config.num_hidden_layers self.counter_before_loading_weights = time.perf_counter() for filename in sorted(set(contract["source_names"].values())): path = str(root / "tensors" / filename) @@ -50,7 +53,10 @@ def get_all_weights(self, model_config, model): for name in sorted(descriptors): if prefixes is not None and not name.startswith(prefixes): continue - match = re.match(r"layers\.(\d+)\.ffn\.experts\.", name) + match = re.search( + r"(?:^|\.)layers\.(\d+)\.(?:ffn|block_sparse_moe)\.experts\.", + name, + ) if match and int(match.group(1)) < layers: continue if callable(file_filter) and file_filter(name): diff --git a/vllm/models/deepseek_v4_1/exact_mxfp4.py b/vllm/models/deepseek_v4_1/exact_mxfp4.py index 1e4536b4ae25..3efced1878f1 100644 --- a/vllm/models/deepseek_v4_1/exact_mxfp4.py +++ b/vllm/models/deepseek_v4_1/exact_mxfp4.py @@ -25,6 +25,7 @@ MoEPrepareAndFinalizeNoDPEPModular, ) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.utils.torch_utils import set_default_torch_num_threads from .quant_config import DeepseekV41FP8Config @@ -106,11 +107,18 @@ def __init__(self, moe, owner): ): raise NotImplementedError("X4T DS4.1 supports TP without EP/DP") if ( - moe.activation != MoEActivation.SILU + moe.activation not in (MoEActivation.SILU, MoEActivation.SITU) or moe.in_dtype != torch.bfloat16 or moe.has_bias + or ( + moe.activation == MoEActivation.SITU + and ( + moe.activation_situ_beta != 4.0 + or moe.activation_situ_linear_beta != 25.0 + ) + ) ): - raise ValueError("X4T DS4.1 requires bias-free BF16 SwiGLU experts") + raise ValueError("X4T requires bias-free BF16 SwiGLU or SiTU(4,25) experts") def create_weights( self, @@ -152,7 +160,7 @@ def process_weights_after_loading(self, layer): e, h, n = self.num_experts, self.hidden_size, self.local_intermediate shapes = ((e, h // 32, 2 * n), (e, n // 32, h)) if self.owner.scale_scratch is None: - # DS4.1 disallows ubatching: every layer consumes these scale grids + # X4T disallows ubatching: every layer consumes these scale grids # before its successor may overwrite them on the same stream. self.owner.scale_scratch = tuple( torch.empty(s, dtype=torch.uint8, device=device) for s in shapes @@ -162,23 +170,28 @@ def process_weights_after_loading(self, layer): t.shape != shape or t.device != device for t, shape in zip(scratch, shapes) ): raise ValueError("X4T shared scale scratch geometry/device mismatch") - weights = read_exact_mxfp4_layer( - self.owner.checkpoint_root, - self.layer_index, - num_experts=e, - hidden_size=h, - intermediate_size=n * tp, - tp_rank=rank, - tp_size=tp, - device=device, - w13_scale_scratch=scratch[0], - w2_scale_scratch=scratch[1], - ) + # Per-expert CPU slices are too small to amortize intra-op barriers. + # Restore the serving thread policy before kernel preparation. + with set_default_torch_num_threads(1): + weights = read_exact_mxfp4_layer( + self.owner.checkpoint_root, + self.layer_index, + num_experts=e, + hidden_size=h, + intermediate_size=n * tp, + tp_rank=rank, + tp_size=tp, + device=device, + w13_scale_scratch=scratch[0], + w2_scale_scratch=scratch[1], + ) plan = fused_moe.plan_weights( source=fused_moe.PackedSource(format="fp4_e8m0_k32", w13_layout="w31"), activation=fused_moe.ActivationSpec( mode="a16", - nonlinearity="silu", + nonlinearity="situ" + if self.moe.activation == MoEActivation.SITU + else "silu", io_dtype=torch.bfloat16, swiglu_limit=self.moe.swiglu_limit, ), diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index f979c669e27f..8d313664a54b 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -766,7 +766,7 @@ def __init__( config, "min_moe_intermediate_per_partition", 256 ) native_b12x_mxfp4 = ( - vllm_config.model_config.quantization == "mxfp4" + vllm_config.model_config.quantization in ("mxfp4", "kimi_x4t") and vllm_config.kernel_config.moe_backend == "b12x" ) self._aligned_decode_projections = ( From 120a909da7ebd8139dc6d300d5b615981ae2d5d7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 30 Sep 2026 14:31:55 +0000 Subject: [PATCH 3/3] Declare Kimi X4T loader compatibility for runtime releases Co-authored-by: OpenAI Codex --- .lil/changes/vllm-kimi-x4t-experts.json | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 .lil/changes/vllm-kimi-x4t-experts.json diff --git a/.lil/changes/vllm-kimi-x4t-experts.json b/.lil/changes/vllm-kimi-x4t-experts.json new file mode 100644 index 000000000000..753195799499 --- /dev/null +++ b/.lil/changes/vllm-kimi-x4t-experts.json @@ -0,0 +1,16 @@ +{ + "schema": "local-inference-release-change/v1", + "id": "vllm-kimi-x4t-experts", + "category": "feature", + "summary": "Load lossless Kimi-K3 X4T experts with native dense precision.", + "models": [ + "Kimi-K3" + ], + "compatibility": "Requires B12X X4T checkpoint support, --quantization kimi_x4t and --load-format exact_mxfp4. Native MXFP4 and DS4.1 source precision are unchanged.", + "pull_requests": [ + 949 + ], + "authors": [ + "voipmonitor" + ] +}