From 1d02758416df473e3d46a09de0ea1d843fc9bba7 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sat, 30 May 2026 22:45:14 +0000
Subject: [PATCH] Fix: kv_proj outputs hd=512 (1 KV head MQA), Z from
 compressor.gate_proj

---
 single_shot_inference.py | 16 +++++-----------
 1 file changed, 5 insertions(+), 11 deletions(-)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index fdcd5877..9e6e2ffd 100644
--- a/single_shot_inference.py
+++ b/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 ----