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