From daed59490293cfb9ad9c229e0db3b8cdea08b6cf Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 31 May 2026 11:47:16 +0000
Subject: [PATCH] CRITICAL FIX: Add missing q_b_norm (unweighted RMSNorm after
 q_b_proj)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The HuggingFace reference (DeepseekV4ForCausalLM) applies an unweighted
RMSNorm after q_b_proj, normalizing Q before attention. Without it, Q
magnitudes are too large, causing attention scores to collapse to uniform
(entropy ~3.2 with 24 positions) and the model to produce garbage.

q_b_norm has no learnable parameters — just q / RMS(q).
This explains the nearly-uniform attention weights we've been seeing.
---
 single_shot_inference.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/single_shot_inference.py b/single_shot_inference.py
index 742495aa..11082a42 100644
--- a/single_shot_inference.py
+++ b/single_shot_inference.py
@@ -429,6 +429,13 @@ def forward_layer(X_l, w, li, cfg, rope_cos, rope_sin,
                      w[f"{pre}.q_b_proj.weight_scale"],
                      w[f"{pre}.q_b_proj.weight_scale_2"])  # (T, n_h * hd)
 
+    # q_b_norm — unweighted RMSNorm after q_b_proj (paper §2.3.1)
+    # This is critical: normalizes Q before attention, preventing score collapse.
+    # No learnable parameters — just q / RMS(q).
+    q_f = q.float()
+    q_rms = q_f.pow(2).mean(-1, keepdim=True).add(1e-6).rsqrt()
+    q = (q_f * q_rms).bfloat16()
+
     # -- KV projection (MQA: 1 KV head) + KV norm --
     kv = nvfp4_linear(x_normed,
                       w[f"{pre}.kv_proj.weight"],