5.2 kB 56 lines
f15273eb
Round 5 blitz: gr_norm_h_moec/gr_apply_moec, qsa_prep#qf16, -580 MB scratch aliasing, deterministic prompt lookahead (352.1 tok/s prefill, 1.58 GB scratch)
dawn
⋯
audits/opus_r5_comb.md
sourcerendered
download raw history
at main Markdown

Round 5 Track A — fold mm_combine, per-token gr_norm_h, skip dead mixed (opus-fold-mm-combine-and-hc-ffn) #

Patch: /tmp/opus_r5_comb/opus_r5_comb.patch (git -C vendor/strata apply — --check passes against the current working tree). Scratch tree + build: /tmp/opus_r5_comb/strata, /tmp/opus_r5_comb/strata/build-metal. Touches metal/kernels/prompt.metal, metal/src/engine/window.cpp (only prompt_chunk + one namespace-scope constant; not prompt_gemm_init / prompt_gemm_need_bytes / QSA lines — p.mixed, p.yr stay allocated), new metal/tests/comb_fold_parity.cpp (picked up by the tests/*.cpp glob).

What changed #

change kernel wired in prompt_chunk
per-token 4-stream HC read with apply (bo read once/token, not 4x) gr_norm_h4, grid (1,T) hc_read(..., apply, !moe) = every hc_ffn (48/chunk). apply==0 reads (layer 0 hc_attn, layer 1 with PLE) keep gr_norm_h (HC,T): the 4-stream grid is +13% slower there (no shared read, ¼ the threadgroups)
skip dead fp32 mixed store gemm_hc_up: if (mixed) mixed[...] = s; mixed bound to a namespace-scope nil kUnbound when prompt_bo_alias() (= no STRATA_METAL_DUMP_MOE); 96 calls/chunk. (A persistent empty Buffer is needed — mtl::Arg::B stores a pointer, so mtl::Buffer{} would dangle.)
mm_combine folded into the next read gr_norm_h_moec (1,T), gr_apply_moec (ceil(n/256),T) when prompt_bo_alias(): mm_combine dispatched only at l==0 && ple_on_ (ple_moe_gate_norm still reads yr; it is (HC,T) so folding there would read YE 4x — 1/48 layers, not worth it). Layers' hc_attn moe read → gr_norm_h_moec; final apply → gr_apply_moec. DUMP_MOE keeps the old path (yr written, dumped).

Fold details: threads tid < topk load pos[t*topk+k] and srt_w[pos] into threadgroup memory once per token, then each d forms yr exactly as mm_combine (same fma order k=0..topk-1, -1 slots skipped) and feeds it into fma(gs, ys, yr). gr_apply_moec was reshaped from (hc·n/256, T) to (n/256, T) with each thread looping its 4 streams, so YE is gathered once per (t,d). gr_norm_h_moe's tail was factored into gr_norm_h_out (shared by gr_norm_h4, gr_norm_h_moe, gr_norm_h_moec; bitwise unchanged — ple_moe_parity still passes). Hazards checked: YE/pos/sw are only written by the next layer's mm_place/mm_down, which come after barrier-ordered hc_attn/hc_ffn; the folded read's dispatch group starts after the barrier before dq(hc_*_down), so all of the previous layer's MoE is complete; nothing besides DUMP_MOE reads p.mixed or p.yr (the MTP source is p.R).

Parity (M5, all bitwise memcmp) #

comb_fold_parity (T = 1, 7, 64, 1024; YE rows a shuffled permutation, every 37th slot -1):

  • gr_norm_h_moec == mm_combine + gr_norm_h_moe (R', xnh): PASS ×4
  • gr_apply_moec == mm_combine + gr_apply_moe (R'): PASS ×4
  • gr_norm_h4 == gr_norm_h (R', xnh), apply 0 and 1: PASS ×8
  • gemm_hc_up with mixed nil: mixh identical to bound run: PASS ×4
  • Existing: ple_moe_parity ALL PASS, prompt_parity ALL PASS (gemm_hc_up with mixed bound still writes it). moe_mm_parity exits 3 (mm_gu_42_16_64_16_2 not found) — pre-existing, identical with the untouched vendor/strata/build-metal binary.
  • Parity run is 4.1 s wall. strata_window_cli / libstrata_metal_kernels build clean (no new warnings). No full-model run (per rules).

M5 microbench, T = 1024, ms per 48 reps in one command buffer (best of 3; comb_fold_parity 1024 bench) #

before after Δ
mm_combine + gr_norm_h_moe → gr_norm_h_moec 89.54 83.65 -6.6%
mm_combine + gr_apply_moe → gr_apply_moec 102.03 73.40 -28.1% (once/chunk)
gr_norm_h (HC,T) apply=1 → gr_norm_h4 (1,T) 53.96 44.08 -18.3%
gr_norm_h apply=0 → gr_norm_h4 22.20 25.13 +13.2% (not wired)
gemm_hc_up mixed bound → nil 35.44 31.44 -11.3%

(alone: mm_combine 40.70, gr_norm_h_moe 48.13.) The fold only wins ~6.6% because both kernels were already near M5 DRAM bandwidth (~126–141 GB/s) and the fold removes only the yr write+read (21 MB of ~262 MB per layer at T=1024); the YE gather itself is not reduced. Tried keeping R in registers for the norm tail (gr_norm_h4r/gr_norm_h_moecr): no gain (the R re-read hits cache), dropped.

Estimated per 1024-token chunk (isolated kernels, sum): fold ≈ 5.7 ms (46–47 reads) + apply 0.6 ms + gr_norm_h4 ≈ 9.9 ms + mixed skip ≈ 8.0 ms (96 calls) ≈ ~24 ms/chunk; roughly 2x at T=2048. Realized gain in the full graph may differ (those reads overlap concurrent dq dispatches) — needs the parent's full-model prompt benchmark to confirm.

Not done / notes for others #

  • prompt_chain.inc (mac-kern path) still uses gr_norm_h/mm_combine/fp32 mixed unchanged.
  • With this patch p.mixed and (except PLE layer 0 / DUMP_MOE) p.yr's MoE role are dead in the default path; Track C could drop p.mixed from prompt_gemm_init/need_bytes when prompt_bo_alias() (−MT·N·4 B) — left untouched as instructed.