From f998862d46e7af54467e10156f02cc51a67e268b Mon Sep 17 00:00:00 2001 From: Sherif Waly Date: Tue, 8 Sep 2026 15:08:51 +0100 Subject: [PATCH] [Bugfix] Fix Kimi K3 loading with interleaved weight streams (#53379) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Sherif Waly Co-authored-by: Nicolò Lucchesi --- tests/models/kimi_k3/test_weight_loading.py | 72 +++++++++++++++++++++ vllm/models/kimi_k3/nvidia/model.py | 10 ++- 2 files changed, 80 insertions(+), 2 deletions(-) create mode 100644 tests/models/kimi_k3/test_weight_loading.py diff --git a/tests/models/kimi_k3/test_weight_loading.py b/tests/models/kimi_k3/test_weight_loading.py new file mode 100644 index 0000000000..b24ad17d9a --- /dev/null +++ b/tests/models/kimi_k3/test_weight_loading.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm.models.kimi_k3.nvidia.model import ( + KimiK3ForConditionalGeneration, + KimiLinearForCausalLM, +) + +pytestmark = pytest.mark.cpu_test + + +class _FakeKimiLinearModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.tensor_a = nn.Parameter(torch.zeros(1)) + self.tensor_c = nn.Parameter(torch.zeros(1)) + self.finalized_values: list[tuple[float, float]] = [] + + def load_weights(self, weights): + params = dict(self.named_parameters()) + loaded = set() + for name, value in weights: + params[name].data.copy_(value) + loaded.add(name) + return loaded + + def finalize_mega_moe_weights(self) -> None: + self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item())) + # MegaMoE finalization replaces its original weight parameters. + self.tensor_a = None + self.tensor_c = None + + +def test_interleaved_composite_weights_finalize_kimi_once_after_loading() -> None: + language_model = object.__new__(KimiLinearForCausalLM) + nn.Module.__init__(language_model) + language_model.config = SimpleNamespace(tie_word_embeddings=False) + language_model.model = _FakeKimiLinearModel() + + model = object.__new__(KimiK3ForConditionalGeneration) + nn.Module.__init__(model) + model.language_model = language_model + model.vision_tower = nn.Module() + model.vision_tower.tensor_b = nn.Parameter(torch.zeros(1)) + + loaded = model.load_weights( + iter( + [ + ("language_model.model.tensor_a", torch.tensor([1.0])), + ("vision_tower.tensor_b", torch.tensor([2.0])), + ("language_model.model.tensor_c", torch.tensor([3.0])), + ] + ) + ) + + assert loaded == { + "language_model.model.tensor_a", + "vision_tower.tensor_b", + "language_model.model.tensor_c", + } + assert language_model.model.finalized_values == [] + + model.process_weights_after_loading() + + assert language_model.model.finalized_values == [(1.0, 3.0)] + assert model.vision_tower.tensor_b.item() == 2.0 diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 8bd04a12e6..3e7645d8d2 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1705,12 +1705,15 @@ class KimiLinearForCausalLM( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - loaded = loader.load_weights(weights) + return loader.load_weights(weights) + + def process_weights_after_loading(self) -> None: + # A parent AutoWeightsLoader may invoke load_weights repeatedly for + # non-contiguous streamed prefixes. Finalize only after the full stream. self.model.finalize_mega_moe_weights() # The fused MultiHeadLatentAttention's process_weights_after_loading # (W_UK_T / W_UV absorption) is driven by the loader's generic post-load # hook for any AttentionLayerBase, so no manual trigger is needed here. - return loaded def get_spec_layer_idx_from_weight_name( @@ -2155,3 +2158,6 @@ class KimiK3ForConditionalGeneration( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): loader = AutoWeightsLoader(self) return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + def process_weights_after_loading(self) -> None: + self.language_model.process_weights_after_loading()