From 690c0a1121370c925ec9552ae8b7afc900a227c2 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 31 May 2026 03:16:07 +0000
Subject: [PATCH] CRITICAL FIX: mHC base/scale ordering was wrong

Checkpoint order is [pre, post, res] not [pre, res, post]:
- base[0:4] = S_pre, base[4:8] = S_post, base[8:24] = S_res
- scale[0] = alpha_pre, scale[1] = alpha_post, scale[2] = alpha_res
- W_stacked rows: [W_pre(4), W_post(4), W_res(16)]
- Projection split: A_raw=proj[:,0:4], C_raw=proj[:,4:8], B_raw=proj[:,8:24]

This was causing B_l to be near-identity and C_l to be near-2.0,
leading to exponential residual stream growth.
---
 single_shot_inference.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index aa828864..017e4f3e 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -138,17 +138,20 @@ class mHCBlock:
         """Load from checkpoint tensors.
 
         fn:   (24, 28672) FP32
-        base: (24,)
-        scale: (3,) — [alpha_pre, alpha_res, alpha_post]
+        base: (24,) — [pre(4), post(4), res(16)]
+        scale: (3,) — [alpha_pre, alpha_post, alpha_res]
         """
         n = self.n_hc
+        # CRITICAL: checkpoint base order is [pre, post, res], not [pre, res, post]
+        # This matches the old working code: base[:4]=pre, base[4:8]=post, base[8:24]=res
         self.W_stacked = fn.to(device=self.device, dtype=torch.float32).contiguous()
         self.S_pre = base[0:n].reshape(1, n).to(device=self.device, dtype=torch.float32)
-        self.S_res = base[n:n + n*n].reshape(n, n).to(device=self.device, dtype=torch.float32)
-        self.S_post = base[n + n*n:].reshape(n, 1).to(device=self.device, dtype=torch.float32)
+        self.S_post = base[n:2*n].reshape(n, 1).to(device=self.device, dtype=torch.float32)
+        self.S_res = base[2*n:2*n + n*n].reshape(n, n).to(device=self.device, dtype=torch.float32)
+        # CRITICAL: checkpoint scale order is [alpha_pre, alpha_post, alpha_res]
         self.alpha_pre = scale[0].item() if scale.numel() > 0 else 0.01
-        self.alpha_res = scale[1].item() if scale.numel() > 1 else 0.01
-        self.alpha_post = scale[2].item() if scale.numel() > 2 else 0.01
+        self.alpha_post = scale[1].item() if scale.numel() > 1 else 0.01
+        self.alpha_res = scale[2].item() if scale.numel() > 2 else 0.01
 
     def _project_and_rms(self, X_flat):
         """Compute RMSNorm(X_flat) @ W_stacked.T → (T, N_proj) BF16.
@@ -173,11 +176,11 @@ class mHCBlock:
         X_flat = X_l.reshape(T, self.K_proj)  # (T, K_proj)
         proj = self._project_and_rms(X_flat).float()  # (T, N_proj) FP32
 
-        # Split projection
-        i0, i1, i2, i3 = 0, n, n + n*n, self.N_proj
-        A_raw = proj[:, i0:i1]
-        B_raw = proj[:, i1:i2]
-        C_raw = proj[:, i2:i3]
+        # Split projection — order matches W_stacked rows: [pre(4), post(4), res(16)]
+        i0, i1, i2, i3 = 0, n, 2*n, self.N_proj
+        A_raw = proj[:, i0:i1]           # (T, n_hc) — pre
+        C_raw = proj[:, i1:i2]           # (T, n_hc) — post
+        B_raw = proj[:, i2:i3]           # (T, n_hc²) — res
 
         # Add biases and scale by gating alphas (paper eqs. 3-5)
         A_tilde = self.alpha_pre  * A_raw + self.S_pre