From fba1c06cad7fa217f08a47dbf793bca83b7c4436 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Mon, 1 Jun 2026 03:02:44 +0000
Subject: [PATCH] diag: check SE weight integrity

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

diff --git a/single_shot_inference.py b/single_shot_inference.py
index d57b4f1e..e0854e32 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -428,6 +428,10 @@ def moe_forward(x, li, moe_runner, se_runner, router, token_id):
             cuda_ok = False
             print(f"  L{li} SE CUDA ERROR: {e}", flush=True)
         print(f"  L{li} MoE shared: |out|={shared_out.abs().max().item():.4f} has_nan={torch.isnan(shared_out).any().item()} cuda_ok={cuda_ok}", flush=True)
+        # Check SE weight integrity
+        if hasattr(se_runner, 'l1_mat_b') and se_runner.l1_mat_b is not None:
+            w = se_runner.l1_mat_b
+            print(f"  L{li} SE l1_mat_b: |max|={w.abs().max().item():.4f} has_nan={torch.isnan(w.view(torch.float16)).any().item()} shape={w.shape}", flush=True)
         print(f"  L{li} SE gsa: l1={se_runner._l1_activation_global_scale:.6f} l2={se_runner._l2_activation_global_scale:.6f} gsb: l1={se_runner._l1_gsb[0].item():.6f} l2={se_runner._l2_gsb[0].item():.6f}", flush=True)
     return routed_out + shared_out