Add per-layer NaN tracking for mHC debug

single_shot_inference.py

@@ -437,6 +437,17 @@ def main():
             ffn_mhc = ffn_mhc_blocks.get(li)
             rc, rs = rope_caches[gpu]
             X = forward_layer(X, layer_weights[li], li, cfg, rc, rs, attn_mhc, ffn_mhc)
+            
+            # NaN check every 5 layers or on first NaN
+            if li % 10 == 0 or (X.float().abs().max().item() > 1e4 or torch.isnan(X.float()).any().item()):
+                has_nan = torch.isnan(X.float()).any().item()
+                xmax = X.float().abs().max().item()
+                print(f"    L{li}: nan={has_nan}, max_abs={xmax:.2f}")
+                if has_nan:
+                    # Debug: check mHC outputs
+                    x_out = X[:, 0, :]
+                    print(f"    L{li} stream0: nan={torch.isnan(x_out.float()).any().item()}, max={x_out.float().abs().max().item():.2f}")
+                    break
         
         # Back to gpu0
         X = X.to('cuda:0')