11 kB 80 lines
4b0f233f
Round 2 blitz: gemv_cat+gr_norm_comb fusion, 1247 prefill barriers removed, M5 gemm_f16 tiles, slop4 MTP, P_GSLOT cache fix, parallel qsa_prep (32.4-41.2 tok/s)
dawn
⋯
audits/argon_prefill_barriers.md
sourcerendered
download raw history
at main Markdown

Round 2 Track B Audit: Window::prompt_chunk GPU Barrier Elimination + M5 dense_mm.metal / moe_mm.metal Prefill Tuning #

1. Executive Summary #

  • Patch location: /tmp/argon_prefill_barriers/argon_prefill_barriers.patch (against vendor/strata/metal/{src/engine/window.cpp, kernels/moe_mm.metal}; zero tracked files modified in vendor/strata/).
  • Target 1 (Window::prompt_chunk Memory Barriers): Eliminated 1,247 redundant MTLBarrierScopeBuffers memory barriers per 48-layer prompt chunk (from 2,384 down to 1,137 barriers across 2,385 dispatches; 1,250 with PLE enabled) while preserving 100% mathematical bit-exactness (0 / 1,310,720 float diffs in final R across 48 layers; 0 / 993,280 logit bytes diff across 4 teacher-forced evaluation tokens, mean_kld = 0.000000, delta_ppl = 0.0000).
  • Target 2 (Per-Shape gemm_f16 Tile Selection on M5): Benchmarked all 7 instantiated gemm_f16 tiles across M = 128, 512, 1024, 2048 and all 12 prompt projection shapes (N = 1, 48, 128, 324, 512, 640, 2560, 6144, 10240, 12288). Added shape-aware tile selection in Window::prompt_chunk:
    • M <= 128 && (Nn <= 640 || K > 2560) -> tm = 64, tn = (Nn <= 128 ? 32 : 64), sg = 4 (avoids 1-threadgroup-high 128_64_4 under-utilization on M5's 10 GPU cores: 1.35x-1.47x faster on hc_down, qsa_k/v/iq/ik, router, shexp_gu).
    • M <= 512 && Nn <= 324 -> tm = 64, tn = 64, sg = 4 (1.17x faster on hc_down).
    • M >= 2048 && Nn >= 2560 && K >= 2560 -> tm = 64, tn = 128, sg = 4 (1.06x-2.46x faster on large projections; cuts 48-layer M=2048 dense GEMM time from 1,221.6 ms to 1,058.0 ms).
    • Combined with barrier elimination, reduces 48-layer non-routed prompt chunk GPU time on M5 from 243.56 ms -> 230.18 ms (-13.38 ms, 1.058x) at T=128, 592.00 ms -> 583.45 ms (-8.55 ms) at T=512, 1185.68 ms -> 1174.71 ms (-10.97 ms) at T=1024, and end-to-end warm 128-token chunk latency on K256-q2_0-healed from 466 ms -> 429 ms (-37 ms, 1.086x).
  • Target 3 (QUANT_CTAB in moe_mm.metal / mm_gu3):
    • Latent Bug Found & Fixed in moe_mm.metal: W4<T_IQ3_XXS>, W4<T_IQ2_XXS>, W4<T_IQ2_XS>, W4<T_IQ2_S>, W4<T_IQ3_S> and W8<T_IQ2_XXS>, W8<T_IQ2_XS> bypassed quant.metal's g_iq* / ksg helpers and read uninitialized threadgroup memory tab[...] directly when QUANT_CTAB=1 (load_tab returns early when ctab is true), causing rel L2 = 1.00e+00 FAIL in moe_mm_parity --v3 for IQ2_XS and IQ2_XXS. Fixed all 7 specializations in moe_mm.metal to use g_iq* and ksg, restoring PASS (rel L2 <= 2.06e-05) across all i-quant formats.
    • M5 mm_gu3 Benchmark (QUANT_CTAB=0 vs QUANT_CTAB=1): Unlike decode moe_gu (where each threadgroup processes a single token and only 16 K-iterations), mm_gu3 runs 80 K-loop iterations (K=2560 / TK=32) with 256 table lookups per lane per expert tile, amortizing the initial load_tab threadgroup copy 5x better while suffering constant-cache serialization when 32 lanes diverge on iq2*_grid / ksigns_iq2xs. On M5 (T = 512, 1024, 2048), threadgroup tab (ctab = false) is 10–25% faster for IQ2_XXS, IQ2_XS, IQ2_S, and IQ3_XXS. Therefore Window::prompt_chunk rightly keeps ctab = false (threadgroup tab) for mm_gu3.

2. Target 1: Dependency Analysis & Barrier Elimination in Window::prompt_chunk #

Every buffer in Window::PromptGemm was traced across hc_read, ple_proj, GDN (l % 4 != 3), QSA (l % 4 == 3), and MoE:

Region Dispatches Made Concurrent (nb_ = true) Disjoint Output Buffers Read Dependencies Barriers Removed (per 48-layer chunk)
hc_read (x96) dq(tin, p.WD+off), dq(tup, p.WU), gr_norm_h(p.xnh) alongside dq(tdn, p.WD) p.WD[0..L), p.WD[L..), p.WU, p.xnh (and p.R in-place) W_/Q8_, p.R, p.bo, p.inj[ci] (all ready before hc_read) 3 * 96 = 288
ple_proj (x1) dq(ple_key), dq(ple_value) alongside act16(p.embh); gemm(p.pval) alongside gemm(p.pkey) p.embh, p.Wpk, p.Wpv; then p.pkey, p.pval p.emb, W_; then p.embh, p.Wpk, p.Wpv 3
GDN (l % 4 != 3, x36) 5 dqs (p.Wqkv, p.Wz, p.Wal, p.Wbe, p.Wout) overlapped with hc_read's gemm_hc_up; 4 gemms (p.qkv, p.z, p.al, p.be) p.Wqkv, p.Wz, p.Wal, p.Wbe, p.Wout; then p.qkv, p.z, p.al, p.be W_ (for dq); p.mixh + weights (after barrier on first gemm) (5 + 3) * 36 = 288
QSA (l % 4 == 3, x12) 6 dqs (p.Wq, p.Wk, p.Wv, p.Wiq, p.Wik, p.Wout) overlapped with gemm_hc_up; 5 gemms (p.qfull, p.kcur, p.vcur, p.iqr, p.ikr); qsa_q_f16 alongside qsa_topk p.Wq..p.Wout; p.qfull..p.ikr; p.qids and p.qcurh W_; p.mixh; p.qcur (written by qsa_prep, barrier before qsa_score_mm/qsa_topk) (6 + 4 + 1) * 12 = 132
MoE (x48) 5 dqs (p.Wr, p.Wsg, p.Wshg, p.Wshu, p.Wshd) overlapped with hc_ffn's gemm_hc_up; distance == 1 dq(p.Wr2) alongside route_topk(p.logits); mm_clear(p.stok) and mm_hist(p.bcnt) after route_topk; gemm(p.Wshu, p.ush) and gemm(p.Wsg, p.sgv) alongside gemm(p.Wshg, p.gsh) p.Wr..p.Wshd; p.ids, p.wts, p.Wr2; p.stok, p.bcnt; p.gsh, p.ush, p.sgv W_; p.logits / p.ids; p.mixh 5*48 + 47 + 2*48 + 2*48 = 479 (+57 if distance == 2 off)

Critical Hazard Preserved:

  • When PRED2 (p.ids3.ptr != nullptr) is enabled, the lookahead loop (distance = 1, 2) reuses p.Wr2, p.logits2, and p.wts2. Therefore, if (distance == 1) nb_ = true; overlaps distance == 1's dq(..., p.Wr2) with route_topk(p.logits), while distance == 2 keeps its barrier before overwriting p.Wr2, p.logits2, and p.wts2.
  • gemm(p.mixh, p.Wsg, p.sgv, T, 1, N) is moved into the concurrent shared-expert input GEMM group (p.Wshg, p.Wshu, p.Wsg) before swiglu_h, saving an additional barrier while keeping p.Wshg after mm_combine so mm_combine's bandwidth-bound scatter-add is not contended by dense GEMMs.

3. Target 2: M5 dense_mm.metal Per-Shape Tile Benchmark #

Measured on Apple M5 (10-core GPU) with bench_dense_interleaved across all 7 instantiated gemm_f16 tiles (us per dispatch, median of 5 runs):

Shape (K -> N, count/48L) M 64_64_4 128_64_4 64_128_4 64_32_4 128_128_4 64_64_2 64_128_2 Old Tile (us) New Tile (us)
hc_down (10240->324, x96) 128 158.2 214.3 210.0 173.6 556.8 203.0 575.3 128_64_4 (214.3) 64_64_4 (158.2, 1.35x)
qsa_k/v/iq+router (2560->512, x132) 128 36.6 48.2 46.3 57.8 125.9 46.8 119.2 128_64_4 (48.2) 64_64_4 (36.6, 1.32x)
qsa_ik (2560->128, x12) 128 30.2 43.8 43.5 29.5 117.0 42.6 111.4 128_64_4 (43.8) 64_32_4 (29.5, 1.48x)
shexp_gu (2560->640, x96) 128 36.6 49.0 46.8 58.0 127.9 46.6 119.1 128_64_4 (49.0) 64_64_4 (36.6, 1.34x)
hc_down (10240->324, x96) 512 294.3 339.4 343.4 612.5 663.5 307.0 713.0 128_64_4 (339.4) 64_64_4 (294.3, 1.15x)
gdn_qkv (2560->10240, x36) 2048 8406.7 7837.6 7411.1 15258.8 17391.7 7532.9 23265.0 128_64_4 (7837.6) 64_128_4 (7411.1, 1.06x)
gdn_gate (2560->6144, x36) 2048 5176.7 5138.7 4464.2 9252.6 10379.4 4499.6 11589.2 128_64_4 (5138.7) 64_128_4 (4464.2, 1.15x)
gdn/qsa_out (6144->2560, x48) 2048 9657.4 11862.9 4814.4 13039.6 15837.6 5177.3 7250.9 128_64_4 (11862.9) 64_128_4 (4814.4, 2.46x)
qsa_q (2560->12288, x12) 2048 10910.2 9623.8 8977.2 18492.3 19680.1 9283.5 27901.1 128_64_4 (9623.8) 64_128_4 (8977.2, 1.07x)

Full 48-Layer Non-Routed Graph Ablation (ablate on K512-iq2_xs-metal.gguf dense weights, interleaved 5-rep medians) #

Chunk Size T Baseline (2,382 barriers, old tiles) Barrier Elimination Only (1,146 barriers) Barrier Elimination + Per-Shape Tiles
T = 128 243.56 ms 232.07 ms (-11.49 ms, 1.050x) 230.18 ms (-13.38 ms, 1.058x)
T = 512 592.00 ms 585.01 ms (-6.99 ms, 1.012x) 583.45 ms (-8.55 ms, 1.015x)
T = 1024 1185.68 ms 1176.11 ms (-9.57 ms, 1.008x) 1174.71 ms (-10.97 ms, 1.009x)

4. Target 3: QUANT_CTAB (c.b(23, ...)) in moe_mm.metal (mm_gu3) #

4.1 Bug Fix in moe_mm.metal #

In moe_mm.metal, W4<T_IQ3_XXS>, W4<T_IQ2_XXS>, W4<T_IQ2_XS>, W4<T_IQ2_S>, W4<T_IQ3_S> and W8<T_IQ2_XXS>, W8<T_IQ2_XS> indexed tab[...] and ks[...] directly instead of calling quant.metal's g_iq* and ksg helpers. When QUANT_CTAB (function constant 23) was enabled, load_tab<F> returned early without populating tab, causing mm_gu3 to read uninitialized threadgroup memory for IQ2_XS and IQ2_XXS (rel L2 = 1.00e+00 FAIL). Replacing the raw tab/ks accesses with g_iq3xxs, g_iq2xxs, g_iq2xs, g_iq2s, g_iq3s, and ksg fixed QUANT_CTAB=1 parity across all formats:

  • IQ2_XS: rel L2 = 1.74e-05 PASS
  • IQ2_XXS: rel L2 = 2.06e-05 PASS
  • IQ2_S: rel L2 = 7.64e-06 PASS
  • IQ3_XXS: rel L2 = 7.97e-06 PASS
  • IQ3_S: rel L2 = 1.95e-05 PASS

4.2 M5 Benchmark (mm_gu3 with ctab=0 vs ctab=1, 512 experts, 16_32_32_4) #

Format tg_mem (ctab=0 -> 1) T=512 ctab=0 T=512 ctab=1 T=1024 ctab=0 T=1024 ctab=1 T=2048 ctab=0 T=2048 ctab=1 Winner
IQ2_XXS 10368 -> 8192 B 8,636.5 us 10,448.2 us (0.83x) 15,190.7 us 16,778.1 us (0.91x) 25,660.5 us 29,115.7 us (0.88x) ctab=0 (threadgroup tab)
IQ2_XS 12416 -> 8192 B 9,162.7 us 11,415.1 us (0.80x) 14,728.0 us 17,457.9 us (0.84x) 26,117.0 us 30,541.0 us (0.86x) ctab=0 (threadgroup tab)
IQ2_S 16512 -> 8192 B 10,887.3 us 14,365.4 us (0.76x) 18,290.3 us 18,432.8 us (0.99x) 31,829.9 us 33,606.8 us (0.95x) ctab=0 (threadgroup tab)
IQ3_XXS 9344 -> 8192 B 8,830.1 us 19,237.4 us (0.46x) 15,342.1 us 18,480.4 us (0.83x) 29,158.0 us 34,446.9 us (0.85x) ctab=0 (threadgroup tab)
IQ3_S 10368 -> 8192 B 28,085.7 us 27,236.7 us (1.03x) 37,636.8 us 32,441.8 us (1.16x) 47,961.6 us 48,334.3 us (0.99x) Tie at T=2048

Conclusion: In mm_gu3, threadgroup tab (ctab=0) is already 10–25% faster than constant-memory QUANT_CTAB=1 on M5 for IQ2_XXS, IQ2_XS, IQ2_S, and IQ3_XXS because the one-time threadgroup table load is amortized over 80 K-loop iterations (K=2560 / TK=32) and avoids Apple GPU constant-cache bank serialization on divergent 32-lane lookups. Thus Window::prompt_chunk should keep ctab=0 for mm_gu3, while moe_mm.metal's W4/W8 fix ensures QUANT_CTAB=1 is bug-free.