fix(aprender-serve): M32d Step 5 — apply per-head Q/K RMSNorm in forward_qwen3_moe (GH-279) — gibberish → coherent English

[noahgift]

TL;DR

Root cause of the M32d %%%%%%%% gibberish output found and fixed. Qwen3 per-head Q/K RMSNorm (GH-279) was applied in the dense path but missing from forward_qwen3_moe.rs. With this fix, apr run against the cached 17.3 GB Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf goes from %%%%%%%% to coherent English text.

Live dogfood evidence on lambda-vector RTX 4090

PRE-FIX:

$ apr run <17.3 GB Qwen3-Coder GGUF> --prompt "What is 2+2?" --max-tokens 8
Output: %%%%%%%%

POST-FIX (this PR):

$ apr run <GGUF> --prompt "What is 2+2?" --max-tokens 8
Output: Human: What is 2+

$ apr run <GGUF> --prompt "Hello" --max-tokens 16
Output: Human: What is the difference between a function and a method in Python?

Output is coherent English. Math completion and chat-template handling are separable downstream issues; the forward pass is now producing context-aware logits.

Five-whys

Why	Answer
1. %%%%%%%%?	Greedy argmax repeats one token
2. Why?	Logits dominated by one direction regardless of context
3. Why?	Hidden state is context-invariant through 48 layers
4. Why?	Qwen3 per-head Q/K RMSNorm (GH-279) missing from forward_qwen3_moe
5 (root). Why missing?	forward_qwen3_moe (M32c.2.2.2.1.1) was authored mirroring the OLD dense forward (pre-GH-279); the GH-279 wiring never propagated to the MoE path. No regression test asserted it for MoE.

How the bug was pinned

• Step 2 (PR feat(aprender-serve): M32d Step 2 — forward_qwen3_moe_traced for per-layer ActivationStats #1222): forward_qwen3_moe_traced per-layer ActivationStats

• Step 2.5 (PR feat(apr-cli): M32d Step 2.5 — wire apr trace --payload to forward_qwen3_moe_traced #1226): apr trace --payload dispatch for qwen3_moe

• Live dogfood:

  layer[0].output_stats.std_dev  = 0.07
  layer[47].output_stats.std_dev = 2.82
  

  → 40× std-dev growth across 48 layers — the exact signature of attention scores compounding without per-head Q/K norm.

  This is the rank-3 component prior (15%) in the M34 FAST PATH plan. The diagnostic surface from Step 2 / 2.5 confirmed it before the HF FP16 fixture was even produced.

The fix

In forward_qwen3_moe.rs, mirror adaptive_ffn.rs:174-179 (dense path's GH-279 reference impl):

if let Some(ref q_norm) = layer.attn_q_norm_weight {
    ops::apply_per_head_rms_norm(&mut q, q_norm, self.config.num_heads, self.config.eps);
}
if let Some(ref k_norm) = layer.attn_k_norm_weight {
    ops::apply_per_head_rms_norm(&mut k, k_norm, self.config.num_kv_heads, self.config.eps);
}

Applied AFTER bias, BEFORE RoPE (matches GH-279).

Regression test

crates/aprender-serve/tests/qwen3_moe_qk_norm_regression.rs::f_qw3_moe_step5_001_context_aware_argmax asserts:

1. Two distinct prompts produce distinct argmax tokens (pre-fix they were the same gibberish)
2. Top-1 vs top-2 logit gap < 50 (pre-fix the gap was much larger because logits collapsed to a single direction)

Live PASS on lambda-vector RTX 4090 in 6.67s.

Hot-path safety

• Production apr run → run_qwen3_moe_generate → forward_qwen3_moe now applies the norm.
• Sibling test f_qw3_moe_c22211_001_full_forward_one_token still passes (finite logits + correct shape unchanged).
• forward_qwen3_moe_traced (PR feat(aprender-serve): M32d Step 2 — forward_qwen3_moe_traced for per-layer ActivationStats #1222) does NOT yet have this fix — follow-up will sync the two paths once feat(aprender-serve): M32d Step 2 — forward_qwen3_moe_traced for per-layer ActivationStats #1222 lands.

Stack research

Per CLAUDE.md "Stack research reference repos" memory:

• HuggingFace Qwen3MoeForCausalLM applies per-head q_norm/k_norm in Qwen3MoeAttention.forward
• llama.cpp ggml_qwen3_moe_kv_norm does the same on GGUF tensors attn_q_norm.weight, attn_k_norm.weight

Both confirm: load-bearing per-arch fix, not a Qwen3-specific quirk.

Test plan

• cargo check -p aprender-serve --lib — clean
• cargo clippy -p aprender-serve --lib -- -D warnings — clean
• cargo fmt -p aprender-serve --check — clean
• cargo test -p aprender-serve --test qwen3_moe_qk_norm_regression --release — 1/1 PASS
• cargo test -p aprender-serve --test qwen3_moe_forward_one_token --release — 1/1 PASS (sibling, hot-path safety)
• Live apr run --prompt "What is 2+2?" → "Human: What is 2+" (was %%%%%%%%)
• Live apr run --prompt "Hello" → "Human: What is the difference between a function and a method in Python?" (was %%%%%%%%)

Refs

• M32d Step 5 of M34 FAST PATH (paiml/claude-code-parity-apr § "M32d FAST PATH")
• feat(aprender-serve): M32d Step 2 — forward_qwen3_moe_traced for per-layer ActivationStats #1222 (Step 2: forward_qwen3_moe_traced — dogfood surface that surfaced this)
• feat(apr-cli): M32d Step 2.5 — wire apr trace --payload to forward_qwen3_moe_traced #1226 (Step 2.5: apr trace dispatch — diagnostic that pinned it)
• Loss trait must be Send + Sync for multithreaded training #279 (Qwen3 per-head Q/K RMSNorm)
• FALSIFY-QW3-MOE-FORWARD-003

🤖 Generated with Claude Code