Fix A_log shape in KimiDeltaAttention (num_heads -> head_dim)
KimiDeltaAttention._init_weights builds A_log with num_heads, but the released checkpoint stores it with head_dim, so from_pretrained fails to load the model's own weights:
RuntimeError: Error(s) in loading state_dict for KimiDeltaAttention:
size mismatch for A_log: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([96]).
In config.json, linear_attn_config: num_heads=96, head_dim=128. The checkpoint's A_log is (128,), matching head_dim (and o_norm, the other per-head-dim tensor) -- KDA uses a per-head-dim decay. This one-line change makes the released weights load.
- self.num_heads, dtype=torch.float32).uniform_(1, 16)))
+ self.head_dim, dtype=torch.float32).uniform_(1, 16)))
Found while loading the checkpoint for NVFP4 quantization with NVIDIA TensorRT Model Optimizer. Verified the checkpoint loads with this fix (all shards); the forward path was not independently re-verified, but the checkpoint's stored A_log shape is authoritative.
Same failure here, same diagnosis—but head_dim sizes the parameter to the padding rather than to the parameter.
A_log is per-head, and flash-linear-attention documents it that way in the two kernels K3 calls: chunk_kda takes "A_log (shape [HV])", fused_recurrent_kda takes "Decay parameter of shape [HV]", and both index it per-head (tl.load(A_log + i_hv)). fla/ops/kda/gate.py carries a "modified and supported by the Moonshot AI Team" header, so that shape isn't an outside reading.
vLLM loads it per-head too. vllm/models/kimi_k3/nvidia/kda.py allocates torch.empty(self.local_num_heads) and installs a_log_weight_loader, which does loaded_weight.narrow(shard_axis, tp_rank * shard_size, shard_size). At tp=1 that's A_log[:96]; at any TP degree the union across ranks is [0, 96), so the last 32 lanes are never read by any rank. That loader also keeps a branch for an older (1, 1, H, 1) export—the layout has already moved once, which is an argument for absorbing it at load rather than reshaping the module to match whatever ships.
On the checkpoint side b_proj is (96, 7168), and that's the tensor that settles it. 96*128 == 128*96, so q_proj (12288, 7168) and dt_bias (12288,) are consistent with either reading; b_proj isn't, and it agrees with linear_attn_config.
To be fair to this patch: sizing to 128 loads and runs correctly. The gate kernels take H from g.shape[-2] and never check A_log's length, so the extra lanes sit unread. So this isn't a numerics argument—it's that the parameter would then describe a per-head quantity as per-channel, and everything downstream has to re-derive that.
#150 slices checkpoint-side with a zero-tail assert, same shape as vLLM's loader. That's the one I'd take. Not my PR to route—just didn't want the shape decided by the padding. No rush.
Hi @pjordanandrsn , thanks a lot for pointing that out and through looking our deployment logs by SGLang, we also confirmed A_log (1, 1, 6, 1) — 6 = 96/16, matching local_num_heads at TP=16, not head_dim.
Closing this PR and deferring the fix to #150.