Update moondream.py

moondream.py

@@ -77,60 +77,64 @@ class KVCache(nn.Module):
             "v_cache", torch.zeros(*cache_shape, device=device, dtype=dtype)
         )
 
-    # in class KVCache
+    # --- replace the whole method in KVCache ---
     def update(self, pos_ids, k, v):
         """
         Supports both:
-        - PREFILL:  pos_ids shape == (q_len,), k/v shape == (B, H, q_len, D)
-        - STEP-DECODE (batched): pos_ids shape == (B,), k/v shape == (B, H, 1, D)
-        - STEP-DECODE (single):  scalar pos_ids,   k/v shape == (1, H, 1, D)
+          • Prefill:  k,v  = (B, n_kv_heads, q_len, d),  pos_ids = (q_len,)
+          • 1-step:   k,v  = (B, n_kv_heads, 1, d),      pos_ids = (B,)
+        Writes into caches shaped (B, n_kv_heads, T_max, d).
         """
-        kout, vout = self.k_cache, self.v_cache  # (Bcache, H, T, D)
-        B, H, Q, D = k.shape
+        kout, vout = self.k_cache, self.v_cache
     
-        # Case A: PREFILL — a vector of all time indices
-        if torch.is_tensor(pos_ids) and pos_ids.ndim == 1 and pos_ids.numel() == Q and Q > 1:
-            # broadcast batch dimension into cache if needed
-            if kout.size(0) != B:
-                # grow/shrink the first dim to match B (this happens after you cloned
-                # image caches to B rows for batched prefill)
-                new_k = kout.new_zeros((B,) + tuple(kout.shape[1:]))
-                new_v = vout.new_zeros((B,) + tuple(vout.shape[1:]))
-                # copy row 0 as base (image prefix) into all rows
-                new_k[:] = kout[0]
-                new_v[:] = vout[0]
-                self.k_cache = kout = new_k
-                self.v_cache = vout = new_v
-    
-            # write the whole segment for all rows at once
+        if not torch.is_tensor(pos_ids):
+            # Scalar position for singleton batch (legacy)
             kout[:, :, pos_ids, :] = k
             vout[:, :, pos_ids, :] = v
             return kout, vout
     
-        # Case B: STEP-DECODE (batched) — one position per row, q_len == 1
-        if torch.is_tensor(pos_ids) and pos_ids.ndim == 1 and pos_ids.numel() == B and Q == 1:
+        # Normalize dtype
+        pos_ids = pos_ids.to(dtype=torch.long, device=k.device)
+    
+        # Shapes
+        if k.dim() != 4 or v.dim() != 4:
+            raise RuntimeError(f"KV update expects k,v 4D. Got k={tuple(k.shape)} v={tuple(v.shape)}")
+        B, Hkv, q_len, D = k.shape
+    
+        # Ensure cache batch matches B
+        if kout.size(0) != B:
+            raise RuntimeError(f"KV cache batch mismatch: cache.B={kout.size(0)} vs k.B={B}")
+    
+        # Case A: PREFILL — per-row write of a whole span of positions
+        if pos_ids.dim() == 1 and pos_ids.numel() == q_len:
+            for i in range(B):
+                kout[i, :, pos_ids, :] = k[i]  # (Hkv, q_len, D)
+                vout[i, :, pos_ids, :] = v[i]
+            return kout, vout
+    
+        # Case B: STEP DECODE — one new position per row (q_len must be 1)
+        if pos_ids.dim() == 1 and pos_ids.numel() == B and q_len == 1:
             for i in range(B):
-                pi = int(pos_ids[i])
+                pi = int(pos_ids[i].item())
                 kout[i, :, pi, :] = k[i, :, 0, :]
                 vout[i, :, pi, :] = v[i, :, 0, :]
             return kout, vout
     
-        # Case C: STEP-DECODE (single) — scalar pos, B==1, q_len==1
-        if (not torch.is_tensor(pos_ids)) or pos_ids.ndim == 0:
-            pi = int(pos_ids)
-            kout[:B, :, pi, :] = k[:, :, 0, :]
-            vout[:B, :, pi, :] = v[:, :, 0, :]
+        # Optional legacy: scalar pos for everyone
+        if pos_ids.dim() == 0 and q_len == 1:
+            pi = int(pos_ids.item())
+            kout[:, :, pi, :] = k[:, :, 0, :]
+            vout[:, :, pi, :] = v[:, :, 0, :]
             return kout, vout
     
-        # Fallback: shape combo we didn't expect
         raise RuntimeError(
-            f"KVCache.update: unsupported shapes pos_ids={tuple(pos_ids.shape) if torch.is_tensor(pos_ids) else '()'}, "
-            f"k={tuple(k.shape)}, v={tuple(v.shape)}"
+            f"Unsupported KV update combo: k={tuple(k.shape)}, pos_ids={tuple(pos_ids.shape)}"
         )
 
 
 
 
+
 class MoondreamModel(nn.Module):
 
     def __init__(