From 79ba7e6636c53e2257061941cd2a9e8b217d59e1 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 31 May 2026 03:10:05 +0000
Subject: [PATCH] diag: add mHC diagnostics for first 3 layers

---
 single_shot_inference.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index b4949944..bb2c6498 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -505,6 +505,10 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
 
     # -- mHC post_block (attention) --
     X_mid = attn_mhc.post_block(X_l, F_attn, attn_ctx)  # (T, n_hc, H)
+    # Diagnostic: check mHC is stabilizing the residual
+    if li < 3:
+        B_l, C_l = attn_ctx
+        print(f"  L{li} attn: |X_l|={X_l.abs().max().item():.2f} |F_attn|={F_attn.abs().max().item():.2f} |B|={B_l.abs().max().item():.4f} |C|={C_l.abs().max().item():.4f} |X_mid|={X_mid.abs().max().item():.2f}", flush=True)
 
     # ==================================================================
     # FFN SUB-BLOCK