feat(e7-18): thread pre_norm/post_norm through _build_attention + _build_mlp

Completes Spec-v2 §3.5 requirement that E7-6 left half-done. E7-6
(c1bc6f7) shipped norm_validation rules and the BrickContextPanel UI
dropdown, but the builders themselves hardcoded RMSNorm — a user
picking 'layernorm' or 'none' in the UI had zero effect on the model.

cppmega_v4/models/unified_superblock_v4.py:
- New _make_norm(kind, dim, eps) helper returns nn.RMSNorm /
  nn.LayerNorm / None; raises on unknown kind.
- _build_attention now reads params.pre_norm (default 'none' for
  legacy compat) and params.post_norm (default 'rmsnorm' for legacy
  compat); applies pre_norm to input before Q/K/V proj, post_norm to
  output projection result. self.norm alias preserved so existing
  checkpoint state-dict keys still load.
- _build_mlp gains the same pre_norm/post_norm wiring (defaults
  'none' / 'none' since the legacy MLP had no internal norm wrapper;
  callers typically wrap at the block level).

Tests (+15 pytest):
- _make_norm returns the right class for each kind, raises on unknown
- attention with 5 norm combinations (legacy default + 4 mixes)
  produces finite forward output
- mlp with 3 norm combos (none+none / rmsnorm+layernorm /
  layernorm+rmsnorm) produces finite output
- attention default == legacy (pre=None, post=RMSNorm)
- attention post_norm='layernorm' instantiates nn.LayerNorm
- mlp default has no pre/post norms (legacy)

Regression: pytest 2321 (was 2306; +15) / 2 skip / 15 xfail / 0 fail
(excl. pre-existing path_d).

Closes cppmega-mlx-bb0.18.

Author: apstenku123

cppmega_v4/models/unified_superblock_v4.py

@@ -74,25 +74,36 @@ def _build_mlp(hidden_size: int, params: dict) -> nn.Module:
 
     intermediate = params.get("intermediate_size", 4 * hidden_size)
     activation = params.get("activation", "glu")
+    norm_eps = params.get("norm_eps", 1e-6)
+    pre_norm_kind = params.get("pre_norm", "none")
+    post_norm_kind = params.get("post_norm", "none")
 
     class _MLP(nn.Module):
         def __init__(self):
             super().__init__()
             self.gate = nn.Linear(hidden_size, intermediate, bias=False)
             self.up = nn.Linear(hidden_size, intermediate, bias=False)
             self.down = nn.Linear(intermediate, hidden_size, bias=False)
+            self.pre_norm = _make_norm(pre_norm_kind, hidden_size, norm_eps)
+            self.post_norm = _make_norm(post_norm_kind, hidden_size, norm_eps)
 
         def __call__(self, x):
+            if self.pre_norm is not None:
+                x = self.pre_norm(x)
             up = self.up(x)
             if activation == "glu":
-                return self.down(mx.sigmoid(self.gate(x)) * up)
-            if activation in IS_GATED and IS_GATED[activation]:
-                return self.down(apply_activation(activation, up,
-                                                  gate=self.gate(x)))
-            if activation in IS_GATED:
-                return self.down(apply_activation(activation, up))
-            # Unknown name — fall back to GLU to keep training alive.
-            return self.down(mx.sigmoid(self.gate(x)) * up)
+                y = self.down(mx.sigmoid(self.gate(x)) * up)
+            elif activation in IS_GATED and IS_GATED[activation]:
+                y = self.down(apply_activation(activation, up,
+                                                gate=self.gate(x)))
+            elif activation in IS_GATED:
+                y = self.down(apply_activation(activation, up))
+            else:
+                # Unknown name — fall back to GLU to keep training alive.
+                y = self.down(mx.sigmoid(self.gate(x)) * up)
+            if self.post_norm is not None:
+                y = self.post_norm(y)
+            return y
     return _MLP()
 
 
@@ -363,11 +374,30 @@ def _build_mamba3(hidden_size: int, params: dict) -> nn.Module:
     return Mamba3ReferenceBlock(Mamba3Config(**cfg_kwargs))
 
 
+def _make_norm(kind: str, dim: int, eps: float):
+    """E7-6/E7-18: return RMSNorm / LayerNorm / None per kind."""
+    if kind == "rmsnorm":
+        return nn.RMSNorm(dim, eps=eps)
+    if kind == "layernorm":
+        return nn.LayerNorm(dim, eps=eps)
+    if kind == "none":
+        return None
+    raise ValueError(f"unknown norm kind {kind!r}; "
+                     "use 'rmsnorm' / 'layernorm' / 'none'")
+
+
 def _build_attention(hidden_size: int, params: dict) -> nn.Module:
-    """Standard multi-head self-attention (causal). Used for `attention`."""
+    """Standard multi-head self-attention (causal). Used for `attention`.
+
+    Honors params.pre_norm (default 'none', matches legacy behavior)
+    and params.post_norm (default 'rmsnorm', also matches legacy).
+    See cppmega_v4/buildspec/norm_validation.py for diagnostics.
+    """
     num_heads = params.get("num_heads", max(1, hidden_size // 64))
     head_dim = params.get("head_dim", hidden_size // num_heads)
     norm_eps = params.get("norm_eps", 1e-6)
+    pre_norm_kind = params.get("pre_norm", "none")
+    post_norm_kind = params.get("post_norm", "rmsnorm")
 
     class _SelfAttn(nn.Module):
         def __init__(self):
@@ -376,11 +406,16 @@ def __init__(self):
             self.k_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
             self.v_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
             self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
-            self.norm = nn.RMSNorm(hidden_size, eps=norm_eps)
+            self.pre_norm = _make_norm(pre_norm_kind, hidden_size, norm_eps)
+            # 'norm' alias preserves state-dict keys from earlier
+            # checkpoints where only one (post) norm existed.
+            self.norm = _make_norm(post_norm_kind, hidden_size, norm_eps)
             # Zero-init out so the block is identity at init.
             self.o_proj.weight = mx.zeros_like(self.o_proj.weight)
 
         def __call__(self, x):
+            if self.pre_norm is not None:
+                x = self.pre_norm(x)
             B, S, _ = x.shape
             q = self.q_proj(x).reshape(B, S, num_heads, head_dim)
             k = self.k_proj(x).reshape(B, S, num_heads, head_dim)
@@ -396,7 +431,10 @@ def __call__(self, x):
             w = mx.softmax(scores.astype(mx.float32), axis=-1).astype(scores.dtype)
             o = mx.matmul(w, v)
             o = mx.transpose(o, (0, 2, 1, 3)).reshape(B, S, num_heads * head_dim)
-            return self.norm(self.o_proj(o))
+            o = self.o_proj(o)
+            if self.norm is not None:
+                o = self.norm(o)
+            return o
 
     return _SelfAttn()
 

tests/v4/test_norm_builder_params.py

@@ -0,0 +1,92 @@
+"""E7-6-builder tests: pre_norm/post_norm params actually thread into
+_build_attention and _build_mlp (Spec-v2 §3.5)."""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import mlx.nn as nn
+import pytest
+
+from cppmega_v4.models.unified_superblock_v4 import (
+    BLOCK_BUILDERS, _make_norm,
+)
+
+
+HIDDEN = 32
+
+
+def test_make_norm_returns_rmsnorm():
+    n = _make_norm("rmsnorm", HIDDEN, 1e-6)
+    assert isinstance(n, nn.RMSNorm)
+
+
+def test_make_norm_returns_layernorm():
+    n = _make_norm("layernorm", HIDDEN, 1e-5)
+    assert isinstance(n, nn.LayerNorm)
+
+
+def test_make_norm_returns_none_for_none():
+    assert _make_norm("none", HIDDEN, 1e-6) is None
+
+
+def test_make_norm_rejects_unknown_kind():
+    with pytest.raises(ValueError, match="unknown norm kind"):
+        _make_norm("batchnorm", HIDDEN, 1e-6)
+
+
+@pytest.mark.parametrize("pre,post", [
+    ("none", "rmsnorm"),  # legacy default
+    ("rmsnorm", "none"),
+    ("layernorm", "rmsnorm"),
+    ("rmsnorm", "rmsnorm"),
+    ("layernorm", "layernorm"),
+])
+def test_attention_with_norm_combos_forward(pre, post):
+    """Each valid norm combination produces finite forward output."""
+    mlp = BLOCK_BUILDERS["attention"](HIDDEN, {
+        "pre_norm": pre, "post_norm": post,
+    })
+    x = mx.random.normal((1, 4, HIDDEN), key=mx.random.key(0))
+    y = mlp(x)
+    assert y.shape == (1, 4, HIDDEN)
+    assert mx.isfinite(y).all().item()
+
+
+@pytest.mark.parametrize("pre,post", [
+    ("none", "none"),
+    ("rmsnorm", "layernorm"),
+    ("layernorm", "rmsnorm"),
+])
+def test_mlp_with_norm_combos_forward(pre, post):
+    """MLP same — pre/post wrap the activation core."""
+    mlp = BLOCK_BUILDERS["mlp"](HIDDEN, {
+        "intermediate_size": 64,
+        "pre_norm": pre, "post_norm": post,
+        "activation": "swiglu",
+    })
+    x = mx.random.normal((1, 4, HIDDEN), key=mx.random.key(0))
+    y = mlp(x)
+    assert y.shape == (1, 4, HIDDEN)
+    assert mx.isfinite(y).all().item()
+
+
+def test_attention_default_matches_legacy_only_post_norm():
+    """Without pre_norm/post_norm params, behavior matches legacy
+    (single RMSNorm at output, no pre-norm)."""
+    mlp = BLOCK_BUILDERS["attention"](HIDDEN, {})
+    assert mlp.pre_norm is None
+    assert isinstance(mlp.norm, nn.RMSNorm)
+
+
+def test_attention_layernorm_post_switch_instantiated():
+    mlp = BLOCK_BUILDERS["attention"](HIDDEN, {
+        "post_norm": "layernorm",
+    })
+    assert isinstance(mlp.norm, nn.LayerNorm)
+
+
+def test_mlp_default_no_norms_legacy_behavior():
+    """MLP legacy default: no pre/post norm wrappers (caller wraps)."""
+    mlp = BLOCK_BUILDERS["mlp"](HIDDEN, {})
+    assert mlp.pre_norm is None
+    assert mlp.post_norm is None