opus R19 Track A: fused ple_comb_query_gate_norm, register-cached ple_gnorm, load-time PSO pre-warm #

patch: /tmp/opus_r19_prewarm/opus_r19_prewarm.patch (git -C vendor/strata apply --check passes against the current working tree) scratch tree: /tmp/opus_r19_prewarm/strata (build: build-metal/), A/B bench: /tmp/opus_r19_prewarm/ab_bench.cpp (+ binary), baseline kernel copy: /tmp/opus_r19_prewarm/ple_base.metal files touched: metal/src/engine/msl/ple.metal, metal/src/engine/window.cpp (prewarm region ~475-570, Window::run PLE block, window_init one line)

what changed #

ple.metal #

• ple_comb_query_gate_norm (grid (hc, T), 256 thr): reads R once, forms layer 0's deferred MoE combine on the fly (fma(fma(gs, ysh, sum_k ye[plan[P_TOKENT + t*10 + j]]), g, R), gr_norm_comb's exact expression and k order), then gnorm_regs + gate_norm_regs. R written once; gr_norm's xn (dead: layer 1's gr_read runs with apply=false and re-norms R after the conv) never written. Unrolled fast path for n=2560 && topk=10, generic loop otherwise. P_TOKENT defaults to 532 under #ifndef (same as gr2.metal).
• ple_gnorm: n=2560 path keeps the row in xv[NPT] and reuses gnorm_regs, so x is read once (was twice, un-unrolled loops). Generic path unchanged. Moved below gnorm_regs (needs NPT / the helper).

window.cpp #

• new members in the prewarm region: ple_key_norm(ds, Tn) and ple_after_layer0(ds, Tn) carry the PLE dispatches, shared by Window::run and the prewarm (so prewarm now builds exactly run's graph, incl. the PLE kernels; the old prewarm skipped the PLE block and so built a moe_combine_sh flush that run never uses).
• ple_after_layer0: when comb_bo_ == &bo_ && comb_ntok_ == Tn it takes the deferred combine and dispatches ple_comb_query_gate_norm; else (or STRATA_METAL_PLE_COMB=0, or STRATA_METAL_GR_COMB=0) the old gr_norm + ple_query_gate_norm. Window graph: -1 dispatch, -1 barrier, -120 KB/token DRAM (xn write + R write + R re-read).
• prewarm_pipelines: bool prewarmed_ -> int prewarmed_t_ (sizes 1..prewarmed_t_ done; a later larger max_t warms only the new sizes) + mtp_prewarmed_. Loop now t = 1..max_t_ (was 2..), with fold(ds, t, 0) (gdn_step#t), the PLE block, argmax_part/final, and once the MTP layer: mtp_encode(k = 1..MAXT, full=false) (mtp_prefill chunks / draft catch-up), mtp_encode(1, full=true) (draft step), mtp_pick(0, true) when chain_. Resets cuts_/sigs_/rec_/cur_inj_/nb_/comb_bo_/comb_ntok_.
• prewarm(): env gate (STRATA_METAL_PREWARM=0 off), skips while a streaming arena is not made yet (stream_ && !st_->ready(): the MoE dispatches bind st_->chunk(), which indexes an empty vector before the arena), then warms what is missing. Called from window_init right after set_max_t (default path: prompt_gemm_init has made the arena by then), and from Window::run as the lazy fallback (STRATA_METAL_LAZY_ARENA=1, or MTP loaded later).

deliberate deviations from the brief #

• no extra layer(ds, l, 1, 0) (pos0=0, sel=false) pass: in layer() pos0 only feeds sel, and sel=false only skips pushing qsa_score/qsa_topk + the indexer q gemv; P() for qsa_score/qsa_topk is called either way and the indexer gemv's PSO is the same 512-out gemv as attn_k/attn_v. pos0=4096 is a strict PSO superset; the pass would only cost build time.
• no membp_ = &memb_ swap: membp_ is nullptr outside mtp_prefill (which resets it), and with tokens == nullptr mtp_encode only binds (membp_ ? *membp_ : memb_) as an argument, so the swap is a no-op.
• the max_t_ > 1 condition is gone: at max_t=1 the first T=1 window's compile now also moves to load time.

verification (M5 Air; NOTE the live :8080 server shares the GPU, so absolute us are noisy; ratios are paired medians) #

bitwise, ab_bench BASE NEW 128 T (memcmp on R, norm_w; on y for ple_gnorm; R checked to have moved):

(3 sweeps, every run ALL PASS; the quiet-sweep numbers above, contended sweeps had the same ratios +-0.05.)

existing tests in the scratch build: ple_moe_parity ALL PASS (its chain uses the new ple_gnorm against the fused register kernels, so it cross-checks the new ple_gnorm too), ple_moe_parity --ab ple_base.metal 7 ALL PASS, fuse_parity ALL PASS, arena_budget_test all passed. strata_window_cli builds with no new warnings in window.cpp.

honest scale / not verified #

• per decode window the two kernel changes save roughly 1.2 + 1.5 us of GPU time (~3 us vs ~47 ms/token): real but ~0.006%; the 120 KB/token DRAM saving is the same order. Not a tok/s lever on its own.
• the pre-warm is the bigger user-visible piece (moves ~280 ms + MTP first-use compiles out of the first request), but it was NOT measured end-to-end: per the brief no full-model strata_window_cli run and nothing on :8080. Needs one load + first-request run to confirm (a) first-request prefill/decode/mtp ms drop, (b) the load-time cost of warming (window_ensure passes MAXT=8, so all 8 sizes are compiled at load even if the server runs a smaller max_t), and (c) output unchanged (greedy tokens identical with STRATA_METAL_PLE_COMB=0 vs default).
• static review points for the auditor: prewarm at init mutates only host mqids_ slot 0 (via mtp_encode full=true at cell0=0), which every real full mtp_encode rewrites before use; all other state is restored at the end of prewarm_pipelines.