# opus_ple_moe_out — Round 4 Track A: fold `moe_out` into the next HC read + fuse the PLE norms

Patch: `/tmp/opus_ple_moe_out/opus_ple_moe_out.patch` (`git -C vendor/strata apply --check` passes against the current dirty tree).
Scratch tree: `/tmp/opus_ple_moe_out/strata` (git-tracked; base commit = rsync of vendor/strata at start). Test build: `/tmp/opus_ple_moe_out/build`.

## What changed

| file | change |
|---|---|
| `metal/kernels/prompt.metal` | `gr_norm_h_moe` (moe_out + gr_norm_h in one, **one threadgroup per token covering all 4 streams** so bo is formed once per token), `gr_apply_moe` (final layer: R' only, elementwise) |
| `metal/src/engine/msl/ple.metal` | `ple_moe_gate_norm` (prompt: layer 0's moe_out + gr_norm R-update + key norm + query norm + gate + gated/normalized, 10 floats/thread in registers), `ple_query_gate_norm` (decode: query norm + gate, keyn precomputed), shared `gnorm_regs` / `gate_norm_regs` helpers |
| `metal/src/engine/window.cpp` | `prompt_chunk`: no `moe_out` dispatch; `hc_attn` reads use `gr_norm_h_moe` (grid `(1, T)`); l==0 PLE = `ple_moe_gate_norm` + `ple_conv_add` (was gr_norm + 2x ple_gnorm + ple_gate_norm + conv); final `gr_norm` -> `gr_apply_moe`. `prompt_gemm_init` / `prompt_gemm_need_bytes`: drop `p.keyn`, `p.qn`; `p.bo` aliases `p.yr` (own rows only when `STRATA_METAL_DUMP_MOE` is set, see below). `Window::run`: ple_gnorm(R)+ple_gate_norm -> `ple_query_gate_norm`, `qn_` gone |
| `metal/tests/ple_moe_parity.cpp` | new standalone bitwise test |

`moe_out`, `ple_gnorm`, `ple_gate_norm` kernels are kept (prompt_chain.inc, decode keyn, prompt_parity still use them).

## Results (measured, M5 Air, standalone; no :8080, no full-model CLI)

`ple_moe_parity` (0.8 s): **ALL PASS, bitwise (memcmp)** for T = 1, 7, 64:
- `gr_norm_h_moe` == moe_out + gr_norm_h (R', xnh)
- `gr_apply_moe` == moe_out + gr_norm (R')
- `ple_moe_gate_norm` == moe_out + gr_norm + ple_gnorm(key) + ple_gnorm(query) + ple_gate_norm (R', gated, normW)
- `ple_query_gate_norm` == ple_gnorm(query) + ple_gate_norm (gated, normW)

Falsifier check: perturbing the gnorm eps and the gr_norm_h_moe residual by 1 ulp-ish made the matching lines FAIL (exit 1); reverted -> ALL PASS.
Regression: `fuse_parity` ALL PASS (5.5 s), `prompt_parity` ALL PASS (1.1 s). Scratch `window.cpp` passes `-fsyntax-only` with the real build's flags.

GPU time at T = 1024 (the laptop's streaming chunk; 3 runs, stable to ±0.5%; throwaway bench `/tmp/opus_ple_moe_out/bench_ple_moe.cpp`):

| chain | old | new | delta |
|---|---|---|---|
| moe_out + gr_norm_h (x48 layers/chunk) | 1338 µs | **1002 µs** | -25%, ≈ -16 ms/chunk |
| l==0: moe_out + gr_norm + 2x ple_gnorm + ple_gate_norm (once/chunk) | 4454 µs | **2561 µs** | -42%, ≈ -1.9 ms/chunk |

Note: a first `gr_norm_h_moe` with gr_norm_h's (stream, token) grid was *slower* (1470 µs vs 1347) since each of the 4 stream threadgroups re-read yr+ys; the per-token layout (all 4 streams in one threadgroup, same per-thread d order and same simd/part reduction order -> still bitwise) is what wins. The same layout would likely speed up the plain `gr_norm_h` for `hc_ffn` too (bo read once per token instead of 4x) — not done here, suggested follow-up.

## Memory

Prompt scratch (`prompt_gemm_need_bytes`), with ple on:
- `p.keyn` + `p.qn`: -2·MT·10240·4 B = **-83.9 MB at MT=1024 (laptop default stream chunk), -167.8 MB at MT=2048**
- `p.bo` aliased onto `p.yr`: -MT·2560·4 B = -10.5 MB / -21.0 MB
- decode: `qn_` gone (-328 KB, MAXT=8)

DRAM traffic per chunk (estimate, not measured with counters): moe_out's bo write + gr_norm_h's bo read = 2·T·N·4 per layer -> ≈1.0 GB/chunk at T=1024 (2.0 GB at 2048); PLE stage drops keyn/qn write+read, the unused fp32 xn write into YE and gated re-read (≈0.3 GB at 1024); final gr_norm's unused xn write (42 MB at 1024) also gone.

## Why bo can alias yr (and the exception)

After the fold, `p.bo` only carries the attention GEMM's output to the same layer's `hc_ffn` read. `p.yr` lives from `mm_combine` (layer l) to the next layer's `hc_attn` read (or the l==0 PLE kernel / final apply). Every writer here is pushed with `barrier = true` (no `nb_` before the attention `gemm(..., p.bo, ...)` nor before `mm_combine`), so the lifetimes are disjoint and serialized. Exception: `STRATA_METAL_DUMP_MOE`'s `at_cut` dumps layer s-1's `yr` at layer s's handshake, i.e. after layer s's attention has written bo — so `prompt_bo_alias()` returns false when that env var is set and bo gets its own rows (both in `make()` and `prompt_gemm_need_bytes()`, so the `need` check stays consistent). Free path unaffected: the alias is not in `all_bufs`.

## Not verified / residual risk

- **No end-to-end run of `Window::prompt_chunk` / `Window::run`** (rules: no :8080, no full-model CLI). Kernel math is bitwise-proven in isolation; the window.cpp wiring (arg order, grids, inj[ci] parity, the alias lifetime) is checked by reading and compiling only. Before merging, run one greedy prefill+decode token-equality check vs. the unpatched binary (e.g. strata_window_cli on a fixed prompt, compare tokens/logits) when the server is free.
- `gr_norm_h_moe` assumes hc <= 4 and `ple_*` fused kernels assume n_embd <= 2560 (10 floats/thread register arrays); both are model constants (`N = 2560, HC = 4` in window.cpp). No runtime guard added.
- `STRATA_METAL_SKIP=moe_out` timing hook no longer does anything for the prompt path (dispatch gone).
- Decode path still does `gr_norm` (possibly `gr_norm_comb`) before `ple_query_gate_norm`; folding gr_norm_comb's R update into the PLE kernel too was left out (comb deferral state makes it messier, and Tn <= 8 so the win is tiny).