opus_r17_mtp — Round 17 Track A: MTP encode / prefill (gemv_cat Q8_0, router, wave zipping, mtp.metal vectorization) #

patch: /tmp/opus_r17_mtp/opus_r17_mtp.patch (git -C vendor/strata apply --check: clean; 4 files, +~170/−31) test: /tmp/opus_r17_mtp/r17_mtp_test.cpp (parity + paired A/B, 4.9 s), /tmp/opus_r17_mtp/r17_gemv_cold.cpp (cold-cache GEMV A/B, 3.5 s); full output /tmp/opus_r17_mtp/r17_mtp_test.out. Built in /tmp/opus_r17_mtp/strata/build-metal (strata_window_cli, gemv_parity, route_parity, fuse_parity all build). device: Apple M5 (m5_layout). No :8080, no full-model runs.

TL;DR #

• All 151 bitwise checks PASS (memcmp == 0): gemv_cat/gemv_catk Q8_0 at every tile × ntok 1..8 × {512, 640}, mtp_xprep ×8, mtp_rms_xprep (embed ×8, hidden-split ×1, hidden-keep ×8), mtp_combine_gnorm ×8, mtp_pick ×6. Existing gemv_parity / route_parity / fuse_parity: ALL PASS.
• mtp.metal vectorization is the real win: mtp_xprep −41/−45 %, mtp_rms_xprep −21…−32 %, mtp_combine_gnorm −9 %, mtp_pick −2 %. ≈ 5 µs GPU per draft step (rms embed −1.2, rms hidden-split −3.7, combine −0.3, pick −0.1) — small next to a draft step, but free and bitwise.
• Q8_0 gemv_cat is a wash on GPU time, so it's wired for the draft step only (bitwise, −2 dispatches/step), not prefill.
• router_route: not applicable — the brief's premise (BF16 router) doesn't hold for the default MTP path.
• Wave zipping: nothing left to zip without changing bits.

1. Q8_0 gemv_cat (k+v, shared gate+up) #

a correction to the brief #

"gemv_cat runs KS = 1, so bitwise" is false for the draft step on M5: mtp_encode's gm tile for n_out ≤ 640, Q8_0, rows == 1 is 32,1,4,K=4 (K-split). A KS = 1 cat there would change the reduction order. So I added gemv_catk<F, LPR, NR, SG, KS> (gemv's exact body + gemv_cat's row → W/W2 select, same red[] fixed-order sum), instances gemvck_8_{32_1_4_4, 32_2_4_2, 32_1_8_8}, next to the GVC block. This is new kernel code inserted at the GVC region (gemv.metal 215+), slightly past the "instantiations only" boundary; the existing gemv_cat body is untouched. Drop the GVCK block if that's out of bounds; then the draft can't use cat (prefill could, see below).

Also added GVC(T_Q8_0, 32,1,4 / 32,1,8 / 16,1,4 / 32,2,4 / 8,1,8), and Window::gemv_cat now accepts Q8_0 (+ the 8,1,8 instance in its inst check). Note: that one also affects the main window's GDN alpha+beta if a model stores them as Q8_0 (bitwise; the shipped models use BF16 there).

window.cpp (mtp_encode) #

• gm's tile choice pulled into a tile(ty, n_in, n_out, rows) lambda (unchanged logic), so gm_cat uses exactly gm's tile → bitwise by construction (KS = 1 → gemvc_, KS > 1 → gemvck_).
• gm_cat(a, b, x, ya, yb, stride, rows): same type + shape, Q8_0/BF16/F32, instance exists, rows == 1, STRATA_METAL_GEMV_CAT ≠ 0; else falls back to the old gm(a); nb_ = true; gm(b) (same wave, same barrier state).
• used for k_proj+v_proj and shared gate_proj+up_proj.

measured (paired A/B, medians, 11 rounds) #

The two GEMVs already share a wave (nb_ = true), run concurrently and saturate bandwidth (~135–139 GB/s cold), so merging saves nothing on the GPU. Prefill got slower → gm_cat is gated to rows == 1. Honestly you could drop 1.x entirely and lose nothing measurable; I kept the draft path for the 2 fewer dispatches encoded per step. (The hot-cache 640+640 case read ~142 µs in both arms; it's a cache artifact of looping the same 3.5 MB, gone in the cold run. Equal across arms either way.)

2. router_route + wave zipping #

• models/K512_iq2_xs/mtp-q2_0.gguf: mtp.layers.0.mlp.gate.weight is BF16 [2560, 512] on disk, but mtp_load quantizes every 2-D BF16 tensor with n_in % 32 == 0 into mq8_ (no router exception), so TQ(gate) is Q8_0 by default. router_route only has BF16/F32 instances (router.metal, not mine), and routef_ is opt-in because it measured no faster in the main window. → kept gm + moe_route, as the brief says to for Q8_0.
• MLP wave structure today: [router GEMV ‖ shared gate/up ‖ shared_expert_gate] → [moe_route ‖ swiglu_prep] → [moe_gu ‖ shared down] → [moe_down] → comb_defer. Each barrier is a true dependency (route needs logits; gu needs plan + XH_; down needs h_); no bitwise-safe zip is left.

3. mtp.metal / mtp_chain.metal #

• mtp_prepare: unrolled; 16 float4 stores (4 planes × 4) instead of 64 scalar. The plane-3 even lanes stay unmultiplied (float4(x.x, x.y*2^-8, x.z, x.w*2^-8)), so a denormal can't get flushed by a *1.0f the original never did. Alignment: raw = xp + t·n + 16i, plane = plane_tok·n (multiples of 16 floats), buffer offsets are 16-B multiples.
• mtp_xprep: 4 float4 loads.
• mtp_rms_xprep: output loop float4 loads of x/w + float4 store of y (x*rs*w, same op order). The sum-of-squares loop is untouched on purpose: vectorizing it changes which elements each thread sums → not bitwise.
• mtp_combine_gnorm: n == 2560 specialization, fully unrolled 10-step loops with hoisted pointers; same fma order; generic path kept.
• mtp_pick: #pragma clang loop unroll(full) on the sequential 32-term sum (order kept).

(400 serialized copies per graph, barrier each, 9 interleaved rounds; medians.)

Caveats / not done #

• No end-to-end draft/acceptance run (rules). Host-side wiring of gm_cat (MB(name)+offset, y_off 0, nb_ state) mirrors gm line for line; kernel arithmetic is covered by the parity test at the exact tiles gm picks.
• gemv_parity took 12 s (I ran it in parallel with route/fuse parity; it's the existing all-formats test).
• Suggested next check once the server is free: STRATA_METAL_GEMV_CAT=0 vs default draft tokens identical.