5.4 kB 65 lines
35f84cbe
Round 18 blitz: per-stream gr_norm_comb 5.09x (-80%), xprep float4 (-46.8%), moe_combine_sh (-49%), Q6_K moe_gu NE>=3 (-25%), qsa_topk2/qsa_score (-5..9%) -> 40.03 tok/s decode & 445.79 tok/s 10.18K prefill
dawn
⋯
audits/opus_r18_moe.md
sourcerendered
download raw history
at main Markdown

opus_r18_moe — Round 18 Track A: moe_gu / moe_down / moe_route NE>1 reuse, xprep float4, swiglu_prep, combines #

Patch: /tmp/opus_r18_moe/opus_r18_moe.patch (185 lines; git -C vendor/strata apply --check = OK). Touches only metal/kernels/{moe,units,glue}.metal and metal/tests/expert_parity.cpp. A/B harness (scratch, not in the patch since it would be a new file): /tmp/opus_r18_moe/moe_ab.cpp (compiles the untouched kernels from /tmp/opus_r18_moe/base_kernels as A and the scratch tree as B in one process, memcmp of every output buffer, paired alternating timing, median of 11 B/A ratios). Raw logs: /tmp/opus_r18_moe/ab_logs.txt, parity log /tmp/opus_r18_moe/parity_log.txt. Device: Apple M5 (the 24 GB Air), live :8080 server up the whole time, so absolute µs drift between runs; the paired ratio is the number to trust. gu/dn bench: 2560 x 640, one 512-expert arena, 48 windows, each window's NE tokens routed to the same 10 experts (so every group has exactly NE entries).

what landed (all bitwise: memcmp == 0 on every A/B run) #

change result
xprep float4 loads/stores + unroll (units.metal) -41.5 % .. -46.8 % (n 2560 ntok 8: 19.8 → 10.5 µs; ntok 1 -41.5 %; n 640 ntok 80 -44.4 %; n 5120 -45.9 %)
moe_combine_sh topk==10 unrolled (glue.metal) -32.0 % ntok 8 (17.6 → 12.0 µs), -45.1 % ntok 1
moe_combine topk==10 unrolled (moe.metal) -33.3 % ntok 8, -48.7 % ntok 1
gu_body fused gate+up per row, one x load per entry — only for UNIT==32 formats at NE≥3, IQ4_XS from NE 2 (gu_fused<F,NE>) Q6_K 32,1,4: NE3 -18.5 %, NE4 -20.5 %; Q4_K -17.6/-18.6 %; Q5_K -15/-18 %; Q8_0 -9/-19 %; IQ4_XS NE2/3/4 -31.7/-33.2/-35.8 %; NR=2 tiles up to -25 %
production formats IQ2_XS gate/up + Q2_0R down unchanged code path (A/B ±0.4 %, noise)

gr_norm_comb (gr2.metal) stays bitwise with moe_combine_sh: the unrolled sum keeps k order. fuse_parity confirms.

what was tested and rejected (numbers so nobody re-tries them) #

  1. gu NE>1 fusion for the 16-weight units (IQ2_XS, IQ3_S, Q2_0) — the brief's main target. Regresses where it matters: IQ2_XS 32,1,4 NE3 +17.0 %, 16,1,4 NE3 +21.4 %, NE4 +4..+10 %; IQ3_S 32,1,4 NE3 +14.7 %, 16,1,4 NE3 +17.3 %. Holding two decoded units live costs more than the second x load (which is an L1 hit of 4 float4s). Only NR=2 tiles gain (IQ3_S 32,2,4 NE4 -9 %, 16,2,4 NE4 -14.5 %; IQ2_XS ≈ ±2 %; Q2_0 NR2 +4 %), and no M5 production tile uses gu NR=2 (window.cpp: gr = 1; gl = 16 for IQ2_XS/IQ3_S at ntok ≥ 3).
  2. "decode every row first" (qg[NR][UNIT], qu[NR][UNIT]) — catastrophic: IQ2_XS/IQ3_S NR=2 +37..+86 %, Q6_K NR=2 +383..+623 % (spills).
  3. Hoisting tok[NE] / x row pointers before the u loop — neutral to worse: IQ2_XS +2..+5 %, IQ3_S 16,1,4 NE3/4 +6 %. 2*NE live 64-bit pointers cost more than the uniform etok[t] reload.
  4. dn_body Q2_0/Q2_0R on the shared-decode path (q[NR][UNIT]) — moe_down_142_8_2_4: NE2 +5.0 %, NE3 +11.0 %, NE4 +5.8 %; 4,2,4: +7.7/+14.2/+12.1 %; 8,1,4 = 0 (NR 1 is the same code). The existing exclusion is right; the patch only adds a one-line comment saying so. ew[NE] hoist: it is only read in the epilogue, nothing to hoist.
  5. swiglu_prep via simd_shuffle_down (0 B threadgroup memory, no barrier, bitwise) — +2.0..+2.7 % at ntok 8 (n 640 and 5120), -0.9 % at ntok 1, ±0 at ntok 4 (two repeats each). 16 shuffles on every lane cost more than one lane's 16 threadgroup reads after a barrier in a one-simdgroup threadgroup. Not landed. (The real lever is the 16-thread threadgroup itself — half a simdgroup — but that is a host change in window.cpp, out of scope.)
  6. moe_route unroll pragmas — 0 % (ntok 1: +0.6 %, ntok 8: +0.4 %): the compiler already unrolls these. Replacing v[ex / 32] = -INFINITY by a 16-way select to keep v in registers — +19 % (ntok 8) .. +56 % (ntok 1). Neither landed; moe_route is untouched.

verification #

  • A/B memcmp == 0 for every run above (gu 11 formats x tiles x NE 1..4, dn 142 x 3 tiles x NE 1..4, xprep x4 shapes, swiglu x3, combine/combine_sh ntok 1/8, route ids/w/plan ntok 1/4/8).
  • expert_parity (vs double CPU ref): 160/160 PASS — Q6_K, Q4_K, Q5_K, Q8_0, IQ4_XS, IQ2_XS, IQ3_S, Q2_0 gate/up x ntok 1,2,3,4,8 x {random 0.3-share routing, new --same} x gu tiles {32,1,4, 32,2,4}; plus IQ2_XS/Q2_0 --dn-prep. --same (new) routes every token to the same 10 experts so groups have exactly ntok entries and the NE 2/3/4/8 bodies (incl. gu_fused) are hit deterministically.
  • route_parity PASS (0 mismatches), fuse_parity ALL PASS (gr_norm_comb bitwise), prompt_parity ALL PASS.
  • No :8080 traffic, no strata_window_cli runs.

expected effect on the live model #

Our K512 model is IQ2_XS gate/up + Q2_0 (142) down: its gu/dn kernels are unchanged. What it gets is xprep (~9 µs a call saved at ntok 8 — xprep runs several times a layer) and moe_combine_sh (~5.6 µs a layer at ntok 8, if it is not folded into gr_norm_comb; the folded gr_norm_comb path is a different kernel in gr2.metal that this round does not own and that likely has the same non-unrolled topk loop — worth a follow-up there). The gu_fused win lands for the UD-Q4/5/6_K_XL and IQ4_XS packs.