opus_r6_ple — Round 6 Track A: PLE gated elimination, ple_moe_gate_norm_c, PromptGemm PLE/bo/yr scratch aliasing #

patch: /tmp/opus_r6_ple/opus_r6_ple.patch (git -C vendor/strata apply --check passes against the current working tree) scratch tree: /tmp/opus_r6_ple/strata (baseline commit = vendor/strata working tree as of task start), build: /tmp/opus_r6_ple/build files: metal/src/engine/msl/ple.metal, metal/tests/ple_moe_parity.cpp, metal/src/engine/window.cpp

what changed #

ple.metal #

• gate_norm_regs no longer writes gated; it writes R = rb + gated (rb = the residual row the kernel already holds in registers). The new ple_conv_res does R += conv. (rb + gated) + conv is exactly ple_conv_add's r + gated + conv (left-assoc; the runtime compiles MSL with safe math, mtl_rt.mm:80, so no reassociation).
• ple_query_gate_norm (decode): x -> device float* R (read and written in place), the gated arg is gone (buffer indices shift: norm_w 5, args 6).
• ple_moe_gate_norm (prompt, non-comb_fold / DUMP_MOE path): gated gone (norm_w 10, args 11); R is now written once (before: written with R' and then RMW'd again by conv_add with gated).
• new ple_moe_gate_norm_c: mm_combine folded in (same fma chain as prompt.metal's moe_comb), grid (1, T): a token's 4 streams in turn, so YE (topk gathered rows), ys, value, the slots are read once per token.
• ple_conv_add kept unchanged (prompt_chain.inc and the test reference chain still use it); conv math factored into a shared ple_conv() helper.
• ple_gate_norm, ple_gnorm untouched.

window.cpp #

• decode (Window::run, ~1566): ple_query_gate_norm without gated, ple_conv_add -> ple_conv_res. gated_ member + its A(4*MAXT*10240) allocation removed (lines 374/719 — outside the listed line ranges but only a one-token removal each; flagging it).
• PromptGemm struct: gated member removed (line 2928, same note).
• prompt_gemm_init make(): with prompt_bo_alias(): yr = F32(16) (never written now), bo = p.ys (ys lives shared expert -> next layer's HC read / PLE kernel / final gr_apply_moec; bo lives attention GEMM -> the FFN's HC read, which is after the previous ys read and before the next ys write; all with barriers between). PLE: embh = ygh, pkey = qkv, pval = z, Wpk = Wqkv, Wpv = Wz; emb and normW keep their own allocations.
• prompt_gemm_need_bytes() updated to match (yr 16 floats, if (!alias) bo, PLE only emb + normW). The order of allocations is the same as make(); sum with per-buffer 16 KiB rounding matches by construction. The runtime need != prompt_gemm_need_bytes() check in prompt_gemm_init is still there.
• ple_proj() now always runs at l == 0 right before the PLE kernel (both ple_late and !ple_late), i.e. after layer 0's attention + MoE + shared expert, so the aliases are dead/idle there. The if (ple_on_ && !ple_late) ple_proj(); before the layer loop is gone.
• mm_combine condition: if (!comb_fold) -> zero mm_combine dispatches across all 48 layers when comb_fold.
• PLE dispatch at l == 0: ple_moe_gate_norm_c (1, T) when comb_fold, else ple_moe_gate_norm (HC, T) with yr; then ple_conv_res.

REJECTED from the brief: normW = p.XS #

p.normW is read on the CPU after the graph (window.cpp ~3648, the PLE history advance: const float* nw = p.normW.as<float>()), and mm_gather_h writes p.XS in every one of the 48 layers. Aliasing it would corrupt the PLE conv history for every following chunk. (Follow-up if wanted: do the history advance on the GPU right after ple_conv_res — then normW could alias layer-0-idle scratch; that touches lines 3646-3660, outside my range.) Also emb kept separate, as the brief said (CPU gather writes it during layer 0 in the streaming case).

RAM saved (prompt buffers, 16 KiB rounding) #

yr 4MTN + embh 2MTN + pkey 4MT10240 + pval 4MTN + gated 4MT10240 + Wpk 52.4 MB + Wpv 13.1 MB:

• C = 1024 (stream default on <96 GB, i.e. the 24 GB laptop): -165 MB
• C = 2048 (non-stream default): -286 MB
• C = 4096: -506 MB (the brief's -390 MB at 2048 assumed normW->XS, which is unsafe; see above.) Decode: -4MAXT10240 B (gated_).

verification (all short, no :8080, no full-model runs) #

• ple_moe_parity (rewritten): BITWISE vs the unfused primitive chains (moe_out + gr_norm + ple_gnorm x2 + ple_gate_norm + ple_conv_add; mm_combine prepended for _c), comparing final R and normalized rows:
  • ple_moe_gate_norm + ple_conv_res, ple_moe_gate_norm_c + ple_conv_res, ple_query_gate_norm + ple_conv_res: PASS at T = 1, 7, 64 (default run, 0.23 s) and T = 2048 (ple_moe_parity 2048 bench, 1.6 s). gr_norm_h_moe / gr_apply_moe still PASS.
  • mutation checks: R = rb (gated dropped) -> all three FAIL; perturbing the folded combine -> _c FAIL. So the test bites.
  • pos has dropped (-1) slots and scattered rows, like comb_fold_parity.
• comb_fold_parity: ALL PASS. arena_budget_test: all passed.
• builds clean: ple_moe_parity, arena_budget_test, comb_fold_parity, strata_window_cli (ninja, scratch tree).
• NOT run: full model / strata_window_cli (forbidden while :8080 is live). The need==need_bytes check and the alias lifetimes are verified by reading + the runtime check; parent should do one prompt-chunk run (stream and non-stream, PLE on) before merging, comparing logits/tokens to pre-patch (expected bitwise-identical).

GPU timing (M5, scratch bench /tmp/opus_r6_ple/zz_ple_bench.cpp, best of 7 x 8 reps, live server may add noise) #

T = 2048, the PLE stage once per chunk: