Fix: kv_proj outputs hd=512 (1 KV head MQA), Z from compressor.gate_proj

single_shot_inference.py

@@ -228,17 +228,11 @@ def forward_layer(x, w, li, cfg, rope_cos, rope_sin):
     # ---- Reshape for attention ----
     q_heads = q.reshape(T, n_h, hd).permute(1, 0, 2)  # (n_h, T, hd)
     
-    kv_dim = kv.shape[-1]
-    if compress == 0:  # SWA
-        k = kv.reshape(T, 1, hd).permute(1, 0, 2)
-    elif compress == 128:  # HCA
-        c, z = kv.chunk(2, dim=-1)
-        k = c.reshape(T, 1, hd).permute(1, 0, 2)
-    elif compress == 4:  # CSA
-        # kv has 4 streams: Ca, Cb, Za, Zb
-        # For baseline decode with no cache, just use Ca
-        ca = kv[..., :hd]
-        k = ca.reshape(T, 1, hd).permute(1, 0, 2)
+    kv_dim = kv.shape[-1]  # Should be hd=512 for all layer types
+    
+    # kv_proj outputs (hd,) = 1 KV head for MQA
+    # The Z (compression weights) come from compressor.gate_proj separately
+    k = kv.reshape(T, 1, hd).permute(1, 0, 2)  # (1, T, hd) — MQA
     v = k.clone()
     
     # ---- Apply RoPE ----