9.4 kB 120 lines
b3decf33
Round 3 blitz: GDN decode/prefill kernel optimization (gdn_step -12%, gdn_scanr_4 -8.4%), r1 MoCE unscaling fix, spec_min_p schedule (32.3-40.3 tok/s decode, 325 tok/s prefill)
dawn
⋯
audits/opus_gdn_qsa.md
sourcerendered
download raw history
at main Markdown

opus_gdn_qsa — round 3 track A: gemv_cat for QSA attn_k + attn_v (+ MTP k/v), GDN decode/prefill kernels #

patch: /tmp/opus_gdn_qsa/opus_gdn_qsa.patch (639 lines; git -C vendor/strata apply --check passes on the current tree). scratch tree: /tmp/opus_gdn_qsa/strata (baseline commit = the live working tree at start), A/B harnesses in /tmp/opus_gdn_qsa/ab/. nothing touched :8080, vendor/strata/, or ran a model benchmark. the bench numbers below came from a box that also had the live server on it, so trust the interleaved A/B harness numbers (base and new kernels alternated in one process) over the standalone --bench runs.

1. gemv_cat: attn_k + attn_v in all 12 QSA layers, MTP k_proj + v_proj #

what the GGUFs actually hold (blk.{3,7,..,47}, 2560 -> 512) #

model attn_k / attn_v types
K512-iq2_xs-metal 3: Q6_K/IQ4_XS, 7: Q6_K/Q6_K, 11: IQ4_XS/IQ4_XS, 15: IQ3_S/IQ4_XS, 19,27,35,43: IQ4_XS/IQ4_XS, 23,31,39: IQ4_XS/Q6_K, 47: Q6_K/Q6_K
K256-q2_0-healed / K256-metal Q6_K/Q6_K except 7: IQ4_XS/Q6_K, 11 & 19: Q6_K/Q5_K
mtp-q2_0.gguf k_proj/v_proj BF16 (Q8_0 in mq8_ with STRATA_METAL_MTP_Q8)

so a same-format-only gemv_cat would have fused only 7/12 K512 layers. gemv_tile returns 32_1_4_1 (KS=1) at 2560->512 for all of Q6_K / IQ4_XS / IQ3_S / Q5_K at every ntok (the m5 special cases only fire at n_out == 640). MTP's own tile logic gives BF16 32_1_4_1, Q8_0 32_1_4_4 at one row and 16_1_4_1 above.

design #

  • gemv.metal: gemv's body moved verbatim into gemv_rows<F,LPR,NR,SG,KS>(); gemv is now load_tab + barrier + gemv_rows. gemv_cat<F,F2,LPR,NR,SG,KS> splits the grid instead of the rows: threadgroups [0, groups) run gemv_rows on W (format F), the rest on W2 (format F2, its own prepared-x plane xp2 and row bytes). every threadgroup runs exactly gemv's code, so it's bitwise for any KS (including the KS=4 Q8_0 MTP tile), and mixed formats are fine because the branch is uniform per threadgroup. the old row-interleaved same-format-only kernel is gone; alpha+beta now goes through the same kernel.
  • instances: BF16/F32 same-format at the 4 old tiles; {Q6_K, IQ4_XS, IQ3_S, Q5_K}² at 32_1_4_1; Q2_0 pair; Q8_0 at 32_1_4_4 / 16_1_4_1 / 32_1_4_1.
  • window.cpp: gemv_cat() takes the pair's formats (instance table mirrors the macros), requires gemv_tile(ta) == gemv_tile(tb), and optionally an explicit tile + weight buffers (for MTP's mq8_/Wm_). QSA layer: attn_k+attn_v → one dispatch, fallback to the two gemvs. mtp_encode: gm's tile code pulled out into mtile(), k_proj+v_proj → gemv_cat when both tiles agree.
  • token_chain.cpp (CHAIN_ALBE path) updated to the new kernel name/args (built, not run: it needs the model).

a compile-time trap i hit #

static inline on the new template helper made the runtime MSL compile of the gemv library +0.65 s (2.40 → 3.05 s, 4 reps each, stable); plain inline brings it back to 2.43 s. kept inline. (the cat instances themselves cost ~nothing.)

verification #

  • fuse_parity (extended): BF16+BF16 48, F32+F32 48, Q6_K+Q6_K, IQ4_XS+IQ4_XS, Q6_K+IQ4_XS, IQ4_XS+Q6_K, IQ3_S+IQ4_XS, Q6_K+Q5_K, Q2_0+Q2_0, BF16+BF16 512, Q8_0+Q8_0 (KS=4 at ntok 1) — ntok 1..8 each, bitwise PASS. the mixed cases read different x planes per matrix; i mutated the kernel to read W2 through xp and the 4 mixed cases FAIL, so the test catches that bug.
  • gemv refactor A/B (base vs new gemv.metal, 7 tiles incl. KS 2/4/8, ntok 1/2/4/8): bitwise identical, timing within ±1%.
  • gemv_parity, pair_parity: ALL PASS. all 28 instances the host can pick resolve; the 4 engine kernel libraries compile together (no name clashes); strata and strata_window_cli build and link in a scratch engine build (/tmp/opus_gdn_qsa/build-eng).

expected gain: honestly about zero GPU time #

in a concurrent encoder the two gemvs already overlap (nb_ = true), and gemv_cat measured the same (12 distinct layers, run_graph):

pair, ntok 2 gemvs concurrent gemv_cat
IQ4_XS+IQ4_XS 1/2/4 16.10 / 22.62 / 32.70 us 16.11 / 22.58 / 32.59 us
Q6_K+IQ4_XS 1/2/4 17.54 / 22.73 / 34.79 us 17.72 / 22.83 / 35.03 us
Q8_0+Q8_0 1/2/4 20.25 / 21.05 / 25.95 us 20.37 / 21.31 / 26.77 us
BF16 48+48 (alpha/beta, already fused upstream) 1/2/4 4.25 / 4.28 / 5.42 us 4.22 / 4.35 / 5.09 us

what it buys: 12 fewer dispatches per decode window (+1 per MTP encode): host encode time and command-processor work, which these GPU-time numbers can't see. it only pays off end to end if the window is encode- or dispatch-bound; measure it on the live server before believing it. STRATA_METAL_GEMV_CAT=0 still turns every cat off.

2. gdn_step (kernels/gdn.metal): bitwise, faster on multi-token windows #

changes (all exact: every value comes from the same expression in the same order):

  • z[t] and the next token's conv input are issued early (z at the top of the token, qkv[t+1] during token t, qkv[0] next to the state loads); gamma[j] is loaded once.
  • the gate and beta are computed on two otherwise idle threads (384, 416) in parallel, instead of serially on tid 0, which also does a conv.
  • kv and o partials use separate buffers (red / red2), and the q/k-norm and output-norm partials use separate ones too (part / part2), which drops the WAR barrier between them. the plane-mode per-16 sums come from simd_shuffle_down within the simdgroup (the same even/odd column order), which drops the ys buffer and the end-of-token barrier.
  • k and q rows reach the FMA loops through simd_broadcast from one lane each, not 64 broadcast threadgroup loads a thread per token (the biggest single piece: per-token cost 1.5 → 1.33 us at grid 10).

A/B (base vs new, same process, 36 rotating layer states, plane mode on), per layer:

bitwise base vs new: PASS
ntok 1 commit: base  60.11 us  new  60.28 us  (+0.3%)
ntok 1 verify: base  28.48 us  new  27.93 us  (-2.0%)
ntok 2 commit: base  60.99 us  new  60.57 us  (-0.7%)
ntok 2 verify: base  38.36 us  new  36.11 us  (-5.9%)
ntok 3 commit: base  65.35 us  new  63.86 us  (-2.3%)
ntok 3 verify: base  49.02 us  new  44.35 us  (-9.5%)
ntok 4 commit: base  79.50 us  new  70.22 us  (-11.7%)
ntok 4 verify: base  60.23 us  new  55.98 us  (-7.1%)
ntok 8 commit: base 108.40 us  new  91.64 us  (-15.5%)
ntok 8 verify: base  89.86 us  new  73.27 us  (-18.5%)

the bitwise check covers ntok 1..8 × commit {1, 0, 0x100|k for every k} × plane {0, on}: state, hist_out and y (all planes + the 16-sums). gdn_parity ALL PASS (IB=4, and IB=8 via GDN_IB8). gdn_prompt_parity ALL PASS.

why ntok=1 doesn't move: commit is ~60 us for 6.3 MB of state read + write ≈ 105 GB/s, i.e. bandwidth. verify is 28 us for 3.1 MB (≈112 GB/s). the per-token part scales linearly with the threadgroup count (grid 10/20/30/40/48 → slope ×1/×2/×3/×4/×4.9), so the core is saturated by one 512-thread TG: it's throughput, not barriers. a probe that dropped a whole barrier saved <1%, and setting max_total_threads to 1024 (fewer registers) changed nothing. tried and dropped: float4 threadgroup reads of k/q (no gain), alpha/beta prefetch for the next token (slower, +4 us at ntok 2-4), unroll(disable) on the token loop (same), caching the broadcast k in 32 registers (same).

3. gdn_scan_r (kernels/gdn_prompt.metal): bitwise, -8% on the engine's gdn_scanr_4 #

  • a token's dots with the state (ka = S^T k, qa = S^T q) move into the previous token's update pass. the sums and their order are unchanged: st[r] is updated, then fma'd into ka/qa in the same r order. that makes one pass over the state registers a token instead of two; the block's first token still takes its dots on its own (the last token of a block does one redundant dot).
  • with R ≤ 32 (KC), the next token's k float4s loaded for the dots stay in registers for its update, which saves 8 of 24 threadgroup float4 loads a token.

A/B, T=2048, one layer:

gdn_scanr_4    T 2048: base 8.193 ms  new 7.509 ms  (-8.4%)
gdn_scanr_2    T 2048: base 8.316 ms  new 8.845 ms  (+6.4%)
gdn_scanr_8    T 2048: base 9.626 ms  new 9.170 ms  (-4.7%)
gdn_scanr_4_8  T 2048: base 7.783 ms  new 7.451 ms  (-4.3%)
bitwise base vs new: PASS

(bitwise over T ∈ {1,2,3,15,16,17,37,300,2048} × all four scanr instances: state and o.) gdn_scanr_4 is the only one the engine uses (prompt_chain.inc, window.cpp); scanr_2 is a test-only variant whose timing is noisy (base vs an identical copy showed +20% once), so i left it alone. gdn_pre/gdn_post are already ≈ bandwidth-bound (~170 MB / ~125 MB per 2048-token layer), so i didn't touch them.

files #

  • metal/kernels/gemv.metal: gemv_rows, mixed-format any-KS gemv_cat + instances
  • metal/src/engine/window.cpp: gemv_cat signature/instance table, QSA attn_k+attn_v, MTP k_proj+v_proj, mtile
  • metal/kernels/gdn.metal: gdn_step as above
  • metal/kernels/gdn_prompt.metal: gdn_scan_r fused dots
  • metal/tests/fuse_parity.cpp: gemv_cat matrix test; metal/tests/token_chain.cpp: the new gemv_cat ABI

not verified here (needs the model / :8080) #

end-to-end tok/s and output identity on the real server. kernel-level bitwise identity suggests generations come out byte-identical with STRATA_METAL_GEMV_CAT on or off, and before or after the patch, but i haven't run that.