v4.0.7 — OLMo-2/3 Post-Norm + QK-Norm Fixes

🔧 Critical Architecture Fixes for OLMo-2/3

Three bugs that caused OLMo-2/3 family models to produce garbage output:

1. Post-norm layer ordering (CPU + GPU paths)

OLMo-2/3 uses post-norm: normalize the attention/FFN output, then add the residual.
ATLAS was incorrectly doing pre-norm (CPU) or normalizing the residual+output sum (GPU).
Fixed both paths to match the HuggingFace Olmo2DecoderLayer reference implementation.

2. QK-norm per-head weight slicing (GPU path)

q_norm.weight has shape [n_heads × head_dim] — each head has unique norm weights.
rmsnorm_inplace uses zip(w.iter()), so every head was getting weights[0..128] instead of weights[h*128..(h+1)*128]. Fixed to slice per-head.

3. QK-norm missing from CPU attention path

The CPU Attention::forward_token() had no QK-norm at all — RoPE was applied directly to raw Q/K projections. Added per-head QK-norm before RoPE to match the GPU path.

Results

Metric	Before	After
Top-1 token (capital of France)	yp (CPU) / décor (GPU)	Paris ✅
CPU/GPU logit agreement	max diff 20.0	max diff 0.000015
Output quality	décor décor décor...	Coherent reasoning traces
Throughput	—	15.4 tok/s (A100, BF16)

Stats

• 562 tests (+0 new), 0 failures
• 3 files changed, +158/-38 lines
• Zero external dependencies