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 #

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):

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).