diff --git a/single_shot_inference.py b/single_shot_inference.py index aa411cf1..608d81a1 100644 --- a/single_shot_inference.py +++ b/single_shot_inference.py @@ -359,6 +359,11 @@ class Compressor: bi = torch.arange(n_complete, device=dev) pos_idx = ((bi + 1) * r - 1).clamp(max=positions.numel() - 1) comp_pos = positions[pos_idx] + # INDEXER PROBE: Compressor output shape + ident = f"hd={self.hd} kv_dim={self.kv_dim} ratio={self.ratio} is_csa={self.is_csa}" + print(f" COMPRESSOR OUT [{ident}]: compressed.shape={tuple(compressed.shape)} " + f"dtype={compressed.dtype} stride={compressed.stride()} " + f"contig={compressed.is_contiguous()}", flush=True) return compressed, comp_pos, torch.zeros(1, T, n_complete, dtype=torch.float32, device=dev) # ===================================================================== @@ -394,12 +399,41 @@ class Indexer: self.compressor = Compressor(4, self.ihd, 7168, dev) self.compressor.load(w, f"{pfx}.compressor", dev) - def forward(self, q_lora, hidden_states, comp_indexer_kv, positions): - if self.q_b_lin is None or comp_indexer_kv is None or comp_indexer_kv.shape[0] == 0: return None + def forward(self, q_lora, hidden_states, comp_indexer_kv, positions, layer_idx=None): + if self.q_b_lin is None or comp_indexer_kv is None or comp_indexer_kv.shape[0] == 0: + print(f" INDEXER SKIP L{layer_idx}: q_b_lin={self.q_b_lin is not None} " + f"comp_idx_kv={tuple(comp_indexer_kv.shape) if comp_indexer_kv is not None else None}", flush=True) + return None dev = q_lora.device; T = q_lora.shape[0]; n_comp = comp_indexer_kv.shape[0] + # INDEXER PROBE: print shapes at layer_idx==0 only + li = layer_idx + if li == 0: + print(f"\n=== INDEXER PROBE L0 ===", flush=True) + print(f" q_lora: shape={tuple(q_lora.shape)} dtype={q_lora.dtype}", flush=True) + print(f" comp_idx_kv: shape={tuple(comp_indexer_kv.shape)} " + f"dtype={comp_indexer_kv.dtype} stride={comp_indexer_kv.stride()} " + f"contig={comp_indexer_kv.is_contiguous()}", flush=True) + print(f" self.n_ih={self.n_ih} self.ihd={self.ihd} n_ih*ihd={self.n_ih * self.ihd}", flush=True) + print(f" self.q_b_lin.in_features={self.q_b_lin.in_features} out_features={self.q_b_lin.out_features}", flush=True) + print(f" self.wp_lin.in_features={self.wp_lin.in_features} out_features={self.wp_lin.out_features}", flush=True) + if self.compressor is not None: + print(f" self.compressor.kv_dim={self.compressor.kv_dim} ratio={self.compressor.ratio} hd={self.compressor.hd}", flush=True) q_idx = self.q_b_lin(q_lora).reshape(T, self.n_ih, self.ihd) w_h = self.wp_lin(hidden_states) # (T, n_ih) - k_idx = comp_indexer_kv.reshape(n_comp, self.n_ih, self.ihd) + # INDEXER PROBE: try reshape, catch failure + try: + k_idx = comp_indexer_kv.reshape(n_comp, self.n_ih, self.ihd) + except RuntimeError as e: + print(f" !!! RESHAPE FAILURE L{li} !!!", flush=True) + print(f" comp_indexer_kv.shape = {tuple(comp_indexer_kv.shape)}", flush=True) + print(f" tried to reshape to ({n_comp}, {self.n_ih}, {self.ihd})", flush=True) + print(f" total elements: have {comp_indexer_kv.numel()}, need {n_comp * self.n_ih * self.ihd}", flush=True) + raise + if li == 0: + print(f"--- INDEXER L0 SCORING TENSORS ---", flush=True) + print(f" q_idx: shape={tuple(q_idx.shape)} dtype={q_idx.dtype}", flush=True) + print(f" k_idx: shape={tuple(k_idx.shape)} dtype={k_idx.dtype}", flush=True) + print(f" w_h: shape={tuple(w_h.shape)} dtype={w_h.dtype}", flush=True) scores = torch.einsum('tnd,cnd->tnc', q_idx.float(), k_idx.float()) scores = F.relu(scores); total = (scores * w_h.unsqueeze(-1).float()).sum(1) tk = min(self.top_k, n_comp); _, idx = total.topk(tk, -1); return idx @@ -558,7 +592,7 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin, # 4. Indexer top-k (CSA) topk_idx = None if indexer is not None and ratio == 4: - topk_idx = indexer.forward(q_a, x_normed, kv_cache.comp_idx_kv, positions) + topk_idx = indexer.forward(q_a, x_normed, kv_cache.comp_idx_kv, positions, layer_idx=li) # 5. Gather KV _pt('gather_start')