model: add per-head QK-norm (Qwen3-compat) for xserv export

xserv's Qwen3 forward unconditionally applies per-head RMSNorm to Q and K
(q_norm/k_norm, shape [head_dim]) before RoPE — even gamma=1 is a real RMS
divide, not identity. xtrain never had this, so an exact xserv<->xtrain loop
was structurally impossible. Add it (reusing the 2D rms_norm op on the
[seq*nh, hd] head rows, inserted between reshape and rope to mirror
qwen3.rs's order) so the trained model is genuinely Qwen3-compatible.

params() inserts q_norm,k_norm after wv; num_params() counts them; the
PyTorch parity refs (parity.py / adamw_parity.py) + their name lists add the
same step so the dumps stay self-consistent.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

crates/xtrain-model/tests/parity.py

@@ -98,7 +98,7 @@ lm_head = load("lm_head")
 layers = []
 for l in range(NL):
     layers.append({p: load(f"l{l}_{p}") for p in
-                   ["attn_norm", "wq", "wk", "wv", "wo",
+                   ["attn_norm", "wq", "wk", "wv", "q_norm", "k_norm", "wo",
                     "ffn_norm", "w_gate", "w_up", "w_down"]})
 
 idx = torch.tensor(ids, dtype=torch.long)
@@ -111,6 +111,9 @@ for L in layers:
     q = (x @ L["wq"]).reshape(SEQ, NH, HD)
     k = (x @ L["wk"]).reshape(SEQ, NH, HD)
     v = (x @ L["wv"]).reshape(SEQ, NH, HD)
+    # Per-head QK-norm (Qwen3-style), before RoPE.
+    q = rms_norm(q, L["q_norm"])
+    k = rms_norm(k, L["k_norm"])
     q = rope(q).transpose(0, 1)  # [nh, seq, hd]
     k = rope(k).transpose(0, 1)
     v = v.transpose(0, 1)