CRITICAL FIX: Add missing q_b_norm (unweighted RMSNorm after q_b_proj)

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.
This commit is contained in:
2026-05-31 11:47:16 +00:00
parent dd50c355a6
commit daed594902

View File

@@ -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"],