Fix A_log loading via checkpoint-side slice instead of resizing the model parameter (alternative to #144)
Problem
self.A_log in KimiDeltaAttention.__init__ is built with shape [num_heads] (96), but the released checkpoint stores it as [head_dim] (128), so from_pretrained raises a size mismatch and the model cannot load its own released weights.
#144 proposes fixing this by resizing the model's A_log parameter to [head_dim] (128) to match the checkpoint. I don't think that's the right direction, and this PR proposes the opposite fix instead.
Why not resize to 128
Inspecting the raw checkpoint bytes for A_log across all 69 KDA layers (not sampled — checked every one): elements [96:128] are exactly 0.0 in every layer, while [0:96] hold real trained values (nonzero mean/std, varying per layer). This is 96 trained per-head decay values, zero-padded to 128 by whatever exported the checkpoint.
This is also what the fla KDA kernel expects: the head count used to index A_log (HV in fused_recurrent_kda/chunk_kda) comes from the shape of v at call time — rearrange(v, '... (h d) -> ... h d', d=head_dim) applied to a head_dim * num_heads-wide projection — not from A_log's own shape. That resolves to 96, matching the model's current [96] parameter. Resizing the parameter to 128 doesn't change what the kernel reads; it just makes A_log.view(H, 1) operate on a differently-shaped tensor.
I tried the #144 fix locally to see what would actually happen: it doesn't load silently and produce wrong output — A_log.view(96, 1) on a 128-element tensor raises RuntimeError: shape '[96, 1]' is invalid for input of size 128 immediately, on the first forward call. So the resize approach fails fast rather than silently — better than I first assumed — but the checkpoint still doesn't load after applying it.
This PR's fix
Add a _load_from_state_dict override on KimiDeltaAttention that, when the checkpoint's A_log is wider than the model's parameter, slices it down to size — after asserting the truncated tail is all zero, so it fails loudly (not silently) if a future checkpoint revision ever has real values there.
def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
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)
No other change to the model. num_heads/head_dim and every other tensor's shape are untouched.
Verification
- Checked all 2,460 dense (non-expert) tensors in the checkpoint against the model's parameter shapes:
A_logis the only one with a size mismatch, and it's the only one with a non-zero-vs-zero-tail discrepancy — no other silent-padding cases in the KDA family (dt_bias,conv1d,o_norm, beta/gate projections all match exactly). - This exact patched file loads the real checkpoint's
A_logtensor (pulled directly from the released safetensors shards) viaload_state_dict, and the loaded values matchcheckpoint[:96]exactly. - A synthetic negative test (non-zero tail) confirms the guard raises instead of silently truncating real data.
- With this fix (plus the usual
trust_remote_code=True/eagerattention setup), I've run the full model end-to-end on a single machine with disk-offloaded MoE experts and gotten coherent generations across code/math/chat/prose prompts, so the loaded values produce sane output, not just a shape that happens to match.
Confirmed independently, from a hand-rolled loader rather than from_pretrained. You've got wider coverage than me on the padding—I only sampled layers 0, 1, 45, 90, you did all 69—so just the two additions:
vLLM already does this narrow at load. vllm/models/kimi_k3/nvidia/kda.py sizes the parameter self.local_num_heads and installs a_log_weight_loader → loaded_weight.narrow(shard_axis, tp_rank * shard_size, shard_size). At tp=1 that's A_log[:96], and at any TP degree the union across ranks is [0, 96), so no rank reads the padding. Your zero-tail assert is the part that loader skips. It also carries a branch for an older (1, 1, H, 1) export, so the layout has moved once before.
And b_proj (96, 7168) is the checkpoint-side tensor that pins num_heads=96: 96*128 == 128*96, so q_proj and dt_bias fit either reading and b_proj doesn't.
Nothing needed from me on this one.