Fix A_log loading via checkpoint-side slice instead of resizing the model parameter (alternative to #144)

modeling_kimi_linear.py

@@ -540,6 +540,29 @@ class KimiDeltaAttention(nn.Module):
             self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
         self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)
 
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        # The released checkpoint stores A_log as [head_dim] rather than
+        # [num_heads], which raises a size mismatch on a plain load. Inspecting
+        # the checkpoint shows elements [num_heads:head_dim] are exactly zero in
+        # every KDA layer -- i.e. num_heads trained per-head decays, zero-padded
+        # to head_dim by the export pipeline. The KDA kernel indexes A_log per
+        # head (the head count used for indexing comes from the shape of v at
+        # call time, not from A_log's own shape), so the model's [num_heads]
+        # parameter is correct; only the checkpoint tensor needs slicing down to
+        # match it.
+        key = prefix + "A_log"
+        if key in state_dict:
+            ckpt_val = state_dict[key]
+            want = self.A_log.shape[0]
+            if ckpt_val.shape[0] > want:
+                tail = ckpt_val[want:]
+                if not torch.all(tail == 0):
+                    raise ValueError(
+                        f"{key}: checkpoint tail beyond index {want} is non-zero; "
+                        "refusing to truncate real values")
+                state_dict[key] = ckpt_val[:want]
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
     def forward(
         self,
         hidden_states: torch.Tensor,