5.5 kB 68 lines
2eb7dc0e
Round 15 blitz: prompt_chunk router+plan/shexp wave zipping (-190 barriers), fused Wal+Wbe N=96 GEMM (-36 dispatches), QSA PSO pre-warm (-202 ms Chunk-1 stall), qsa_score_mm causal-future skip (-19.4%), qsa_prep 320B TG + 2 barriers
dawn
⋯
audits/opus_r15_win.md
sourcerendered
download raw history
at main Markdown

opus r15 track A: prompt_chunk wave zipping, fused alpha|beta GEMM, QSA PSO pre-warm, 24 GB chunk 2048 #

patch: /tmp/opus_r15_win/opus_r15_win.patch (git -C vendor/strata apply --check clean). touches only metal/src/engine/window.cpp (prompt_gemm_init / prompt_gemm_need_bytes / prompt_chunk, original lines 3045-3606) and metal/kernels/gdn_prompt.metal (gdn_pre). scratch tree: /tmp/opus_r15_win/strata (build in build-metal). scratch-only test harnesses (NOT in the patch): /tmp/opus_r15_win/scratch_gdn_prompt_parity.cpp.diff, /tmp/opus_r15_win/scratch_dense_mm_parity.cpp.diff.

what changed #

  1. QSA PSO pre-warm. prompt_gemm_init compiles qsa_score_mm / qsa_topk (keys pg:qsa_score_mm, pg:qsa_topk, same as PP() in prompt_chunk). pcache_ outlives free_prompt_gemm, and the streaming load path calls prompt_gemm_init at load (window.cpp ~3783), so the ~200 ms compile moves to startup. caveat: it only ever hit the first >2051-token prompt of a process (pcache persists), so later prompts were already fine.
  2. stream_chunk 2048 on >= 24 GB in both prompt_gemm_init and prompt_gemm_need_bytes (kept identical, the need != prompt_gemm_need_bytes() check guards it). honest number from the need_bytes formula (k4, flash, stream, ple, alias): C=1024 0.74 GB / C=2048 1.19 GB at 64K cells (0.94 / 1.60 GB at 256K). so 2048 is NOT smaller than 1024 here, it costs ~+0.45 GB transient scratch (separate mmap while a prompt runs, unmapped after). if dawn's serve script already sets STRATA_METAL_PROMPT_CHUNK=2048 (her trace shows T=2048 chunks) this changes nothing for it.
  3. router look-ahead zip (stream && pred_on && l+1<48): dq(blk.l+1 ffn_gate_inp -> Wr2) joins wave 0, gemm(Wr2) runs beside gemm(Wr), route_topk(ids2) beside route_topk(ids). -2 barriers/layer (-94/chunk). mm_hist now always has its barrier (it used to run beside route_topk(ids2) with nb_ = true; now both top-ks are in one wave). PRED2 (ids3) kept as a serial tail after the zipped pair.
  4. MoE plan / shared-expert zip (comb_fold): mm_hist | Wshg (+Wshu when !swu) + Wsg GEMMs -> barrier -> mm_scan | gemm_swiglu_up (or swiglu_h) -> barrier -> mm_place | Wshd GEMM -> barrier -> mm_gather_h. 6 waves -> 4, -2 barriers/layer (-96/chunk); the 1-TG mm_scan and 20-TG mm_place now overlap real GEMM work. all of it stays inside [sigs, cuts) so it still runs beside the host's expert placement. buffer audit: plan writes bcnt/boff/eoff/tiles/meta/stok/sw/pos, shared writes gsh(=gqn)/ush(=gkn)/hh(=ygh)/ys(=bo)/sgv, both only read mixh/ids/wts/W* -> disjoint. !comb_fold (STRATA_METAL_DUMP_MOE) path unchanged.
  5. fused alpha|beta GEMM in GDN layers: beta dequantized into p.Wal at byte 2*48*N, one gemm(p.mixh, p.Wal, p.al, T, 96, N), gdn_pre gets p.al and p.al + 192 B. p.al = F32(MT*96), p.be no longer allocated (nothing else reads it; p.Wbe stays, it is QSA's Wiq). need_bytes updated the same way (equal bytes for MT % 256 == 0 anyway). gdn_pre: ab = beta == alpha + hv ? 2*hv : hv used for the alpha/beta reads only; separate buffers (prompt_chain.inc, the tests) keep stride hv. -36 dispatches/chunk. gemm() picks for Nn == 96: M < 768 ? 16_64_2 : 64_32_2 (measured below; 32_64_2, the old Nn<=128 pick, was slower).

verification (M5 Air, all < 3 s each) #

  • gdn_prompt_parity: ALL PASS (63 PASS, 0 FAIL), incl. scratch case: interleaved [T][96] alpha|beta with beta = alpha + 48 vs separate buffers, both grids (packed 8 and legacy 40), T in {1,3,37,300,2048}: gdn_pre outputs (qn, kn, vc, gb, kq, history) bitwise identical.

  • gdn_prompt_parity --ab <original gdn_prompt.metal> 2048 {0,1} bench: every stage output bitwise PASS vs the old kernel; gdn_pre time ratio 1.0000 / 0.9998 (the pointer compare costs nothing).

  • arena_budget_test: all passed (it does not exercise prompt_gemm_need_bytes; the in-process need != prompt_gemm_need_bytes() check is the guard there).

  • dense_mm_parity --fuse96 M (scratch mode): one N=96 GEMM vs two N=48 (window's current tiles), both halves memcmp-identical for every candidate tile at M = 256..2048. median us (two N48 serial -> N96):

    M 2 x N48 16_64_2 32_64_2 64_32_2
    256 43.0 23.1 28.9 29.3
    512 50.7 29.7 34.5 32.5
    768 55.5 39.0 40.0 35.9
    1024 74.5 54.1 51.3 41.2
    1536 102.2 91.2 74.7 58.4
    2048 143-168 121 102-106 90-103

    caveat: in the chunk these GEMMs run in the same concurrent wave as Wqkv/Wz, so the real saving is the removed work (one 10.5 MB mixh read + 32 wasted columns + a dispatch per layer), likely ~1-2 ms per 2048-row chunk, not the standalone 50 us x 36.

  • strata_window_cli builds clean. not run (rules: :8080 live, no full-model runs), so targets 2/3's effect on real chunk time and end-to-end greedy-token parity are unverified. recommended before merge: one short strata_window_cli prefill A/B (orig vs patch) on a 10K prompt with greedy token compare + STRATA_METAL_PKPROF per-chunk time, when the server is down.

risks #

  • zipping depends on the dispatch-order / nb_ hazard reasoning above (no automated race test for prompt_chunk).
  • the 24 GB 2048 default costs ~0.45 GB more transient scratch than 1024 at 64K cells, contrary to the brief's premise.