diag: print wo_a g_flat magnitude to find where zeros come from

single_shot_inference.py

@@ -366,6 +366,8 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
         a_grp = a_flat.reshape(T, o_groups, gid); oa_3d = oa_bf.reshape(o_groups, o_rank, gid)
         g_out = torch.bmm(a_grp.permute(1, 0, 2), oa_3d.transpose(1, 2))
         g_flat = g_out.permute(1, 0, 2).reshape(T, o_groups * o_rank)
+        if li < 3:
+            print(f"  L{li} wo_a: |g_flat|={g_flat.abs().max().item():.6f} shape={g_flat.shape}", flush=True)
         F_attn = prod_lin['o_b'](g_flat)
     else:
         # o_a_proj as full-rank BF16 linear (no low-rank)
@@ -375,6 +377,8 @@ def forward_attention(x_normed, w, li, cfg, rope_cos, rope_sin,
         else:
             log.warning(f"L{li}: No o_a_proj weight, zero attention output")
             F_attn = torch.zeros(T, cfg["hidden_size"], dtype=torch.bfloat16, device=dev)
+    if li < 3:
+        print(f"  L{li} F_attn: |F_attn|={F_attn.abs().max().item():.6f}", flush=True)
     return F_attn, q_a
 
 # =====================================================================