From 354749b24884af0f906c5be82bd0919f6023152b Mon Sep 17 00:00:00 2001 From: ghorses10 Date: Fri, 31 Jul 2026 20:40:47 +0800 Subject: [PATCH] [Bugfix][Model] Kimi-K3 NVIDIA: delegate regular FusedMoE padding to the selected quantization backend Assisted-by: OpenAI Codex Signed-off-by: ghorses10 --- .../kimi_k3/test_moe_intermediate_padding.py | 139 ++++++++++++++ .../kimi_k3/test_moe_padded_weight_loading.py | 171 ++++++++++++++++++ vllm/models/kimi_k3/nvidia/model.py | 16 +- 3 files changed, 311 insertions(+), 15 deletions(-) create mode 100644 tests/models/kimi_k3/test_moe_intermediate_padding.py create mode 100644 tests/models/kimi_k3/test_moe_padded_weight_loading.py diff --git a/tests/models/kimi_k3/test_moe_intermediate_padding.py b/tests/models/kimi_k3/test_moe_intermediate_padding.py new file mode 100644 index 000000000000..9983eb1da236 --- /dev/null +++ b/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 diff --git a/tests/models/kimi_k3/test_moe_padded_weight_loading.py b/tests/models/kimi_k3/test_moe_padded_weight_loading.py new file mode 100644 index 000000000000..68bff2a9222b --- /dev/null +++ b/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() diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 78093de02a35..ac596849e633 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/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