[Bugfix][Model] Kimi-K3 NVIDIA: delegate regular FusedMoE padding to the selected quantization backend

tests/models/kimi_k3/test_moe_intermediate_padding.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Kimi-K3 MoE intermediate-size ownership tests."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
+    Mxfp4MoeBackend,
+    mxfp4_round_up_hidden_size_and_intermediate_size,
+)
+
+K3_MOE_INTERMEDIATE_SIZE = 3072
+K3_ROUTED_EXPERT_HIDDEN_SIZE = 3584
+
+
+def _make_kimi_config() -> SimpleNamespace:
+    return SimpleNamespace(
+        hidden_size=7168,
+        moe_intermediate_size=K3_MOE_INTERMEDIATE_SIZE,
+        routed_expert_hidden_size=K3_ROUTED_EXPERT_HIDDEN_SIZE,
+        num_experts=896,
+        num_experts_per_token=16,
+        num_shared_experts=None,
+        moe_renormalize=True,
+        use_grouped_topk=True,
+        num_expert_group=1,
+        topk_group=1,
+        moe_router_activation_func="sigmoid",
+        routed_scaling_factor=1.0,
+        hidden_act="situ",
+        activation_situ_beta=1.0,
+        activation_situ_linear_beta=1.0,
+        latent_moe_use_norm=False,
+        rms_norm_eps=1e-6,
+        min_moe_intermediate_per_partition=256,
+    )
+
+
+class _StubLinear(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.e_score_correction_bias = None
+
+
+class _RecordingExperts(nn.Module):
+    calls: list[dict] = []
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        type(self).calls.append(dict(kwargs))
+        self.moe_config = SimpleNamespace()
+        self.w13_weight = nn.Parameter(torch.ones(1), requires_grad=False)
+        self.w2_weight = nn.Parameter(torch.ones(1), requires_grad=False)
+
+
+def _stub_common_model_dependencies(monkeypatch, model_mod, tp_size: int) -> None:
+    monkeypatch.setattr(
+        model_mod, "get_tensor_model_parallel_world_size", lambda: tp_size
+    )
+    monkeypatch.setattr(model_mod, "GateLinear", _StubLinear)
+    monkeypatch.setattr(model_mod, "ReplicatedLinear", _StubLinear)
+    monkeypatch.setattr(model_mod, "RMSNorm", _StubLinear)
+    monkeypatch.setattr(model_mod, "KimiMLP", _StubLinear)
+    monkeypatch.setattr(model_mod, "KimiRoutedOutputTransform", _StubLinear)
+
+
+@pytest.mark.parametrize(
+    ("tp_size", "expected_model_size"),
+    [(8, 3072), (16, 4096), (32, 8192)],
+)
+def test_nvidia_regular_kimi_moe_delegates_logical_size(
+    monkeypatch: pytest.MonkeyPatch, tp_size: int, expected_model_size: int
+):
+    model_mod = pytest.importorskip("vllm.models.kimi_k3.nvidia.model")
+    _RecordingExperts.calls = []
+    _stub_common_model_dependencies(monkeypatch, model_mod, tp_size)
+    monkeypatch.setattr(model_mod, "FusedMoEFactory", _RecordingExperts)
+    monkeypatch.setattr(model_mod, "aux_stream", lambda: None)
+    monkeypatch.setattr(torch.cuda, "Event", lambda: SimpleNamespace())
+    monkeypatch.setattr(
+        model_mod,
+        "current_platform",
+        SimpleNamespace(
+            is_cuda=lambda: False,
+            is_device_capability_family=lambda *_: False,
+        ),
+    )
+    vllm_config = SimpleNamespace(
+        kernel_config=SimpleNamespace(moe_backend="marlin"),
+        parallel_config=SimpleNamespace(enable_expert_parallel=False),
+    )
+
+    module = model_mod.KimiMoE(_make_kimi_config(), vllm_config)
+
+    assert len(_RecordingExperts.calls) == 1
+    assert _RecordingExperts.calls[0]["intermediate_size"] == 3072
+    assert module.padded_moe_intermediate_size == expected_model_size
+    assert torch.count_nonzero(module.experts.w13_weight) == 1
+    assert torch.count_nonzero(module.experts.w2_weight) == 1
+    assert not hasattr(
+        module.experts.moe_config, "intermediate_size_per_partition_unpadded"
+    )
+
+
+@pytest.mark.parametrize(
+    ("backend", "expected_intermediate"),
+    [
+        (Mxfp4MoeBackend.MARLIN, 128),
+        (Mxfp4MoeBackend.BATCHED_MARLIN, 128),
+        (Mxfp4MoeBackend.DEEPGEMM_MXFP4, 128),
+        (Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8, 128),
+        (Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16, 128),
+        (Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8, 128),
+        (Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16, 128),
+        (Mxfp4MoeBackend.EMULATION, 96),
+        (Mxfp4MoeBackend.CPU, 96),
+    ],
+)
+def test_backend_rounding_of_tp32_partition(backend, expected_intermediate):
+    _, rounded = mxfp4_round_up_hidden_size_and_intermediate_size(
+        backend, K3_ROUTED_EXPERT_HIDDEN_SIZE, 96
+    )
+    assert rounded == expected_intermediate
+
+
+@pytest.mark.parametrize(
+    ("tp_size", "expected_local"), [(8, 384), (16, 256), (32, 128)]
+)
+def test_marlin_tp_sweep(tp_size: int, expected_local: int):
+    raw_local = K3_MOE_INTERMEDIATE_SIZE // tp_size
+    _, rounded = mxfp4_round_up_hidden_size_and_intermediate_size(
+        Mxfp4MoeBackend.MARLIN, K3_ROUTED_EXPERT_HIDDEN_SIZE, raw_local
+    )
+
+    assert rounded == expected_local

tests/models/kimi_k3/test_moe_padded_weight_loading.py

@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Generic MoE weight loading with an MXFP4-padded Kimi-K3 partition."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
+
+LOGICAL_INTERMEDIATE = 3072
+TP_SIZE = 32
+RAW_LOCAL = LOGICAL_INTERMEDIATE // TP_SIZE
+PADDED_LOCAL = 128
+MXFP4_BLOCK = 32
+SYNTHETIC_HIDDEN = 64
+
+
+class _Loader:
+    def __init__(self, tp_size: int = TP_SIZE):
+        self.moe_config = SimpleNamespace(
+            is_act_and_mul=True,
+            moe_parallel_config=SimpleNamespace(tp_size=tp_size),
+        )
+
+    _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim)
+    _narrow_expert_data_for_padding = staticmethod(
+        RoutedExperts._narrow_expert_data_for_padding
+    )
+    _load_w13 = RoutedExperts._load_w13
+    _load_w2 = RoutedExperts._load_w2
+
+
+def _rows(width: int, columns: int = 1) -> torch.Tensor:
+    return torch.arange(1, width + 1).unsqueeze(1).expand(width, columns).clone()
+
+
+def _columns(rows: int, width: int) -> torch.Tensor:
+    return torch.arange(1, width + 1).unsqueeze(0).expand(rows, width).clone()
+
+
+def test_mxfp4_allocates_expected_kimi_k3_shapes():
+    method = object.__new__(Mxfp4MoEMethod)
+    method.moe = SimpleNamespace(has_bias=False, w13_num_shards=2)
+    layer = nn.Module()
+
+    method.create_weights(
+        layer=layer,
+        num_experts=1,
+        hidden_size=3584,
+        intermediate_size_per_partition=PADDED_LOCAL,
+        params_dtype=torch.bfloat16,
+    )
+
+    assert layer.w13_weight.shape == (1, 256, 1792)
+    assert layer.w2_weight.shape == (1, 3584, 64)
+    assert layer.w13_weight_scale.shape == (1, 256, 112)
+    assert layer.w2_weight_scale.shape == (1, 3584, 4)
+    assert torch.count_nonzero(layer.w13_weight) == 0
+    assert torch.count_nonzero(layer.w2_weight) == 0
+    assert torch.count_nonzero(layer.w13_weight_scale) == 0
+    assert torch.count_nonzero(layer.w2_weight_scale) == 0
+
+
+@pytest.mark.parametrize("shard_id", ["w1", "w3"])
+@pytest.mark.parametrize("tp_rank", [0, 31])
+def test_w13_rank_boundaries_and_zero_tail(shard_id: str, tp_rank: int):
+    checkpoint = _rows(LOGICAL_INTERMEDIATE, SYNTHETIC_HIDDEN // 2)
+    parameter = torch.zeros(2 * PADDED_LOCAL, SYNTHETIC_HIDDEN // 2, dtype=torch.int64)
+
+    _Loader()._load_w13(parameter, 0, shard_id, checkpoint, tp_rank)
+
+    half_start = 0 if shard_id == "w1" else PADDED_LOCAL
+    half = parameter[half_start : half_start + PADDED_LOCAL]
+    expected = checkpoint.narrow(0, tp_rank * RAW_LOCAL, RAW_LOCAL)
+    torch.testing.assert_close(half[:RAW_LOCAL], expected)
+    assert torch.count_nonzero(half[RAW_LOCAL:]) == 0
+    other_start = PADDED_LOCAL if shard_id == "w1" else 0
+    assert torch.count_nonzero(parameter[other_start : other_start + PADDED_LOCAL]) == 0
+
+
+def test_w13_all_tp32_ranks_cover_checkpoint_once():
+    checkpoint = _rows(LOGICAL_INTERMEDIATE)
+    seen = torch.zeros(LOGICAL_INTERMEDIATE, dtype=torch.int32)
+
+    for tp_rank in range(TP_SIZE):
+        parameter = torch.zeros(2 * PADDED_LOCAL, 1, dtype=torch.int64)
+        _Loader()._load_w13(parameter, 0, "w1", checkpoint, tp_rank)
+        indices = parameter[:RAW_LOCAL, 0] - 1
+        seen[indices] += 1
+        assert torch.count_nonzero(parameter[RAW_LOCAL:PADDED_LOCAL]) == 0
+
+    assert torch.equal(seen, torch.ones_like(seen))
+
+
+@pytest.mark.parametrize("tp_rank", [0, 31])
+def test_w2_rank_boundaries_and_zero_tail(tp_rank: int):
+    packed_width = LOGICAL_INTERMEDIATE // 2
+    raw_packed_local = RAW_LOCAL // 2
+    checkpoint = _columns(SYNTHETIC_HIDDEN, packed_width)
+    parameter = torch.zeros(SYNTHETIC_HIDDEN, PADDED_LOCAL // 2, dtype=torch.int64)
+
+    _Loader()._load_w2(parameter, 1, checkpoint, tp_rank)
+
+    expected = checkpoint.narrow(1, tp_rank * raw_packed_local, raw_packed_local)
+    torch.testing.assert_close(parameter[:, :raw_packed_local], expected)
+    assert torch.count_nonzero(parameter[:, raw_packed_local:]) == 0
+
+
+def test_w2_all_tp32_ranks_cover_checkpoint_once():
+    packed_width = LOGICAL_INTERMEDIATE // 2
+    raw_packed_local = RAW_LOCAL // 2
+    checkpoint = _columns(1, packed_width)
+    seen = torch.zeros(packed_width, dtype=torch.int32)
+
+    for tp_rank in range(TP_SIZE):
+        parameter = torch.zeros(1, PADDED_LOCAL // 2, dtype=torch.int64)
+        _Loader()._load_w2(parameter, 1, checkpoint, tp_rank)
+        indices = parameter[0, :raw_packed_local] - 1
+        seen[indices] += 1
+        assert torch.count_nonzero(parameter[:, raw_packed_local:]) == 0
+
+    assert torch.equal(seen, torch.ones_like(seen))
+
+
+@pytest.mark.parametrize("shard_id", ["w1", "w3"])
+@pytest.mark.parametrize("tp_rank", [0, 31])
+def test_w13_scale_loading(shard_id: str, tp_rank: int):
+    checkpoint = _rows(LOGICAL_INTERMEDIATE, SYNTHETIC_HIDDEN // MXFP4_BLOCK)
+    parameter = torch.zeros(
+        2 * PADDED_LOCAL, SYNTHETIC_HIDDEN // MXFP4_BLOCK, dtype=torch.int64
+    )
+
+    _Loader()._load_w13(parameter, 0, shard_id, checkpoint, tp_rank)
+
+    half_start = 0 if shard_id == "w1" else PADDED_LOCAL
+    half = parameter[half_start : half_start + PADDED_LOCAL]
+    expected = checkpoint.narrow(0, tp_rank * RAW_LOCAL, RAW_LOCAL)
+    torch.testing.assert_close(half[:RAW_LOCAL], expected)
+    assert torch.count_nonzero(half[RAW_LOCAL:]) == 0
+
+
+@pytest.mark.parametrize("tp_rank", [0, 31])
+def test_w2_scale_loading(tp_rank: int):
+    logical_scale_width = LOGICAL_INTERMEDIATE // MXFP4_BLOCK
+    raw_scale_local = RAW_LOCAL // MXFP4_BLOCK
+    checkpoint = _columns(SYNTHETIC_HIDDEN, logical_scale_width)
+    parameter = torch.zeros(
+        SYNTHETIC_HIDDEN, PADDED_LOCAL // MXFP4_BLOCK, dtype=torch.int64
+    )
+
+    _Loader()._load_w2(parameter, 1, checkpoint, tp_rank)
+
+    expected = checkpoint.narrow(1, tp_rank * raw_scale_local, raw_scale_local)
+    torch.testing.assert_close(parameter[:, :raw_scale_local], expected)
+    assert torch.count_nonzero(parameter[:, raw_scale_local:]) == 0
+
+
+def test_tp8_unpadded_loading_is_unchanged():
+    tp_size = 8
+    local = LOGICAL_INTERMEDIATE // tp_size
+    checkpoint = _rows(LOGICAL_INTERMEDIATE)
+    parameter = torch.zeros(2 * local, 1, dtype=torch.int64)
+
+    _Loader(tp_size)._load_w13(parameter, 0, "w1", checkpoint, tp_rank=7)
+
+    torch.testing.assert_close(parameter[:local], checkpoint[-local:])
+    assert torch.count_nonzero(parameter[:local]) == parameter[:local].numel()

vllm/models/kimi_k3/nvidia/model.py

@@ -623,7 +623,7 @@ def __init__(
                 num_experts=num_experts,
                 top_k=num_experts_per_token,
                 hidden_size=self.moe_hidden_size,
-                intermediate_size=self.padded_moe_intermediate_size,
+                intermediate_size=moe_intermediate_size,
                 activation=config.hidden_act,
                 activation_situ_beta=activation_situ_beta,
                 activation_situ_linear_beta=activation_situ_linear_beta,
@@ -645,20 +645,6 @@ def __init__(
                 is_sequence_parallel=use_sequence_parallel,
                 runner_cls=LatentMoERunner if self.use_latent_moe else None,
             )
-        if self.padded_moe_intermediate_size != moe_intermediate_size:
-            w13_weight = getattr(self.experts, "w13_weight", None)
-            if w13_weight is None:
-                w13_weight = getattr(self.experts, "w13_weight_packed", None)
-            w2_weight = getattr(self.experts, "w2_weight", None)
-            if w2_weight is None:
-                w2_weight = getattr(self.experts, "w2_weight_packed", None)
-            if w13_weight is not None:
-                w13_weight.data.zero_()
-            if w2_weight is not None:
-                w2_weight.data.zero_()
-            self.experts.moe_config.intermediate_size_per_partition_unpadded = (
-                moe_intermediate_size // self.tp_size
-            )
 
     def _maybe_overlap_router_and_down_proj(
         self, hidden_states: torch.Tensor