From dcbb74841a002a427178c510494be3411b30b2c1 Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 31 May 2026 00:27:48 +0000
Subject: [PATCH] =?UTF-8?q?Remove=20emergency=20RMSNorm=20from=20mHC=20pos?=
 =?UTF-8?q?t=5Fblock=20=E2=80=94=20MoE=20provides=20balance=20now?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 single_shot_inference.py | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index 09d48998..9dff6bcc 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -107,12 +107,7 @@ class mHCBlock:
         mixed_residual = torch.einsum('tij,tjh->tjh', comb_mix, residual.float())
         post_term = post_mix.unsqueeze(-1) * F_out.unsqueeze(1).float()
         residual_next = mixed_residual + post_term
-        # Emergency RMSNorm (remove once MoE provides balance)
-        _T = residual_next.shape[0]
-        rn_f = residual_next.reshape(_T, self.n_hc, -1)
-        rms = rn_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
-        return (rn_f * rms).bfloat16()
-
+        return residual_next.bfloat16()
 
 # =====================================================================
 # RoPE