[Bugfix] Fix Kimi K3 loading with interleaved weight streams (#53379)
Signed-off-by: Sherif Waly <sherif.waly@mistral.ai> Co-authored-by: Nicolò Lucchesi <nicolo.lucchesi@mistral.ai>
Sherif Waly committed
Sep 8, 2026 at 15:08 UTC
f998862d46e7af54467e10156f02cc51a67e268b
2 files changed
+80
-2
tests/models/kimi_k3/test_weight_loading.py
new
+72
@@ -0,0 +1,72 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.models.kimi_k3.nvidia.model import (
11
+ KimiK3ForConditionalGeneration,
12
+ KimiLinearForCausalLM,
13
+)
14
+
15
+pytestmark = pytest.mark.cpu_test
16
+
17
+
18
+class _FakeKimiLinearModel(nn.Module):
19
+ def __init__(self) -> None:
20
+ super().__init__()
21
+ self.tensor_a = nn.Parameter(torch.zeros(1))
22
+ self.tensor_c = nn.Parameter(torch.zeros(1))
23
+ self.finalized_values: list[tuple[float, float]] = []
24
+
25
+ def load_weights(self, weights):
26
+ params = dict(self.named_parameters())
27
+ loaded = set()
28
+ for name, value in weights:
29
+ params[name].data.copy_(value)
30
+ loaded.add(name)
31
+ return loaded
32
+
33
+ def finalize_mega_moe_weights(self) -> None:
34
+ self.finalized_values.append((self.tensor_a.item(), self.tensor_c.item()))
35
+ # MegaMoE finalization replaces its original weight parameters.
36
+ self.tensor_a = None
37
+ self.tensor_c = None
38
+
39
+
40
+def test_interleaved_composite_weights_finalize_kimi_once_after_loading() -> None:
41
+ language_model = object.__new__(KimiLinearForCausalLM)
42
+ nn.Module.__init__(language_model)
43
+ language_model.config = SimpleNamespace(tie_word_embeddings=False)
44
+ language_model.model = _FakeKimiLinearModel()
45
+
46
+ model = object.__new__(KimiK3ForConditionalGeneration)
47
+ nn.Module.__init__(model)
48
+ model.language_model = language_model
49
+ model.vision_tower = nn.Module()
50
+ model.vision_tower.tensor_b = nn.Parameter(torch.zeros(1))
51
+
52
+ loaded = model.load_weights(
53
+ iter(
54
+ [
55
+ ("language_model.model.tensor_a", torch.tensor([1.0])),
56
+ ("vision_tower.tensor_b", torch.tensor([2.0])),
57
+ ("language_model.model.tensor_c", torch.tensor([3.0])),
58
+ ]
59
+ )
60
+ )
61
+
62
+ assert loaded == {
63
+ "language_model.model.tensor_a",
64
+ "vision_tower.tensor_b",
65
+ "language_model.model.tensor_c",
66
+ }
67
+ assert language_model.model.finalized_values == []
68
+
69
+ model.process_weights_after_loading()
70
+
71
+ assert language_model.model.finalized_values == [(1.0, 3.0)]
72
+ assert model.vision_tower.tensor_b.item() == 2.0
vllm/models/kimi_k3/nvidia/model.py
+8
-2
@@ -1705,12 +1705,15 @@ class KimiLinearForCausalLM(
1705
1706
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1707
loader = AutoWeightsLoader(self)
1708
- loaded = loader.load_weights(weights)
1708
+ return loader.load_weights(weights)
1709
+
1710
+ def process_weights_after_loading(self) -> None:
1711
+ # A parent AutoWeightsLoader may invoke load_weights repeatedly for
1712
+ # non-contiguous streamed prefixes. Finalize only after the full stream.
1713
self.model.finalize_mega_moe_weights()
1714
# The fused MultiHeadLatentAttention's process_weights_after_loading
1715
# (W_UK_T / W_UV absorption) is driven by the loader's generic post-load
1716
# hook for any AttentionLayerBase, so no manual trigger is needed here.
1713
- return loaded
1717
1718
1719
def get_spec_layer_idx_from_weight_name(
@@ -2155,3 +2158,6 @@ class KimiK3ForConditionalGeneration(
2158
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
2159
loader = AutoWeightsLoader(self)
2160
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
2161
+
2162
+ def process_weights_after_loading(self) -> None:
2163
+ self.language_model.process_weights_after_loading()