regolith
strata but optimized for m5 (and maybe more later)
Star
Fork
Watch
regolith audits opus_decode_fusion.md
4.2 kB · 44 lines
Markdown
at main

opus_decode_fusion — Round 2 Track A: gemv_cat + gr_norm_comb wiring, mtp_encode barrier audit #

Patch: /tmp/opus_decode_fusion/opus_decode_fusion.patch (applies cleanly with git apply at repo root against the current working tree of vendor/strata/metal/src/engine/window.cpp; also extends tests/fuse_parity.cpp to ntok 1..8). Work tree: /tmp/opus_decode_fusion/strata (patched), /tmp/opus_decode_fusion/strata_orig (baseline), both built.

What changed (window.cpp) #

  1. gemv_cat for ssm_alpha + ssm_beta (36 GDN layers → −36 dispatches/window). gemv()'s tile choice factored out into gemv_tile() (pure code motion); Window::gemv_cat() uses gemvc_ only when the tile gemv() would pick has KS = 1 and is instantiated (BF16/F32; a row's sums depend only on LPR, so bitwise). For the shipped GGUFs (alpha/beta BF16, 2560→48) the tile is 32_1_4_1 → gemvc_30_32_1_4. Writes alL/beL_ at o_ab with y_stride 48: same layout fold() reads. Falls back to two gemvs otherwise. STRATA_METAL_GEMV_CAT=0 disables. Prewarm needs no change: prewarm_pipelines() goes through layer(), so the gemvc#t pipelines are made there.
  2. gr_norm_comb: moe_combine_sh is deferred (comb_defer) and folded into the next gr_norm that reads the same bo buffer with apply=1 (new Window::gr_norm() helper used by gr_read, the PLE gr_norm after layer 0, and gr_read_m). Any other gr_norm flushes the pending combine first (safety net; never hit on current paths). Audit of bo_: written only by ssm_out/attn_output gemv and the combine; read only by gr_norm. After every layer's combine the very next dispatch is that gr_norm (layer l+1 hc_attn, the PLE gr_norm for l=0, output_hc after l=47); nothing else reads bo_ (ydump_ is only fold()'s gdn output, DUMP_MOE is prompt-chain only). Same for mbo_ in mtp_encode → mixer's gr_norm. Saves 48 dispatches + 48 barriers + bo_ round trip per window, 1 + 1 per MTP draft step. STRATA_METAL_GR_COMB=0 disables.
  3. mtp_encode barriers: the embedding chain (rms→xp→fc_embedding) and hidden chain (rms→xp→fc_hidden) now run side by side: 3 barriers instead of 6 (n==1); for n>1 stream 0 joins the pairs, streams 1..3 keep their barriers (XHn_ is reused). Buffers disjoint (men_/mhn_, XEn_/XHn_, me2_/mh2_).
  4. PLE key norm (ple_gnorm on pkey_) moved ahead of layer 0, beside layer 0's first gr_norm: −1 barrier/window. Other nb_ sites in layer()/mtp_encode/mtp_pick audited: remaining barriers are true RAW/WAR deps.

Verification #

  • fuse_parity (M5, extended to ntok 1..8): gemv_cat, router_route, gr_norm_comb all bitwise PASS (~6 s incl. compile).
  • Engine builds clean (ninja strata_window_cli).
  • End to end K512 IQ2_XS + PLE + MTP, deterministic config (see finding below): teacher-forced logits (--logits-out, 22-token prompt + 24 T=1 steps) byte-identical orig vs patched (defaults) vs patched with both knobs =0. Greedy generation --spec 3 (T=3 windows + MTP chain, catch-up n=1,2) and --spec 4 --suffix-draft 3 --n-prompt 100 (prompt GEMM + T≤4 + MTP): identical GEN_IDS and identical draft/accept counts, orig vs patched, 2 runs each.
  • Speed: with the served env the T=3 window time is within noise (orig 65.9–67.7 ms, patched 66.7–74.7 ms; I/O- dominated streaming). Expected gain ~84 fewer dispatches + ~50 fewer barriers per window, i.e. sub-ms; needs a quiet-machine benchmark (parent to run).

Finding (not fixed, other track): served config is run-to-run nondeterministic #

With STRATA_METAL_EXPERT_SPLIT=1 (the served default with SHEXP_OVERLAP) the baseline binary gives different logits on every run (all rows, |Δ| up to ~0.8 already at row 0), even with MOCE_BONUS=0 and PROMPT_GEMM=0. With STRATA_METAL_EXPERT_SPLIT=0 it is bitwise deterministic. The moe.metal comment claims the split is bit-for-bit the unsplit plan, so this looks like a race (phase-1 resident experts running beside the host's turn — slot reuse or plan/ye hazard). MOCE_BONUS>0 adds residency-dependent routing on top. Any A/B on the served config is therefore noise-limited; use EXPERT_SPLIT=0, MOCE_BONUS=0 for bitwise checks.