18 kB 153 lines
9d4fcca5
Add Anti-Slop v4.1b, M5 decode/prefill optimizations (22.64 tok/s decode, 5.27x prefill), and SCR pad-carry
dawn
⋯
audits/argon_k512_decode_opt.md
sourcerendered
download raw history
at main Markdown

Track 2: M5 GPU Decode & SSD Stream Optimization for Unpruned K=512 IQ2_XS #

1. Executive Summary #

We audited and benchmarked the Metal decode pipeline (gemv.metal, moe.metal, gr.metal, gdn.metal, qsa.metal, window.cpp, stream.cpp) on the 24 GB M5 MacBook Air (10-core GPU) for the unpruned 512-expert models/K512_iq2_xs/K512-iq2_xs-metal.gguf pack (3.90 GB dense + 36.23 GB .experts).

Across all 48 layers + PLE + LM head, a single decode token (ntok = 1, topk = 10) reads 3,649.78 MB of weights:

  • Dense + HC + Router weights: 2,942.20 MB/tok (80.6% of total weight traffic)
  • Routed MoE experts (48 x 10): 707.58 MB/tok (486.40 MB moe_gu + 221.18 MB moe_down)

Our isolated M5 GPU benchmarks (gemv_parity, expert_parity, gr_parity, gdn_parity, qsa_parity) and closed-loop 48-layer K=512 SSD stream simulations (k512_sim over 128 real K=512 token trajectories) uncovered 6 concrete, zero-regression optimizations in window.cpp and stream.cpp, reducing isolated per-token M5 GPU kernel latency by ~20.23 ms/tok (a 36.2% reduction on the affected projections) and cutting SSD miss/wait latency by 32.6% fewer misses/tok when unperturbed top-K read-ahead is enabled:

  1. M5 Dense GEMV Tile Selection (Window::gemv, window.cpp:790) — -19.47 ms/tok:
    • window.cpp's generic gemv() tile selector previously chose M4-era defaults (16,1,4,1 for 2560 -> >=6144, 32,1,8,2 for 6144 -> 2560, and 32,1,4,1 for <= 640) that suffer severe register pressure, constant-cache divergence, or poor occupancy on M5 for IQ4_XS, IQ3_S, Q6_K, and Q8_0.
    • Adding an m5_layout() specialization per (type, n_in, n_out) speeds up:
      • ssm_out / attn_output (6144 -> 2560, 39x IQ4_XS + 7x IQ3_S + 2x Q6_K): 7.641 ms -> 5.146 ms (-2.495 ms/tok)
      • attn_gate (2560 -> 6144, 22x IQ4_XS + 11x IQ3_S + 3x Q6_K): 5.605 ms -> 2.627 ms (-2.978 ms/tok, as IQ3_S 16,1,4,1 was running at 21.0 GB/s vs 113.1 GB/s at 32,2,4,1)
      • attn_qkv (2560 -> 10240, 35x IQ4_XS + 1x IQ3_S): 4.196 ms -> 3.989 ms (-0.206 ms/tok)
      • attn_q (2560 -> 12288, 11x IQ4_XS + 1x IQ3_S): 2.582 ms -> 1.982 ms (-0.599 ms/tok)
      • ffn_gate_shexp + ffn_up_shexp (2560 -> 640, 57x IQ4_XS + 23x IQ3_S + 16x Q6_K): 11.124 ms -> 2.204 ms (-8.920 ms/tok, fixing a 13.6 GB/s M5 pathology on IQ4_XS 2560x640 at 32,1,4,1 vs 53.9 GB/s at 16,2,4,1!)
      • pred1 + pred2 lookahead routers (2560 -> 512, 93x Q8_0): 5.856 ms -> 4.808 ms (-1.048 ms/tok)
      • output.weight LM head (2560 -> 248320, 1x IQ4_XS, 337.72 MB): previously fell through to default 16,1,4,1 (or 32,1,4,1 at 3,128.14 us); covered by m5 && ty == ref::IQ4_XS && n_out >= 6144 -> 32,2,4,1 at 2,453.32 us (137.7 GB/s).
  2. M5 Hyper-Connection gr_down Tile (Window::gr_read, window.cpp:879) — -0.68 ms/tok:
    • Switching gr_down_q8_8_8 (8 simdgroups splitting 1 row 8 ways, 324 threadgroups) to gr_down_q8_8_4 (8 simdgroups splitting 2 rows 4 ways each, 162 threadgroups) on M5 reduces per-call latency from 49.56 us (135.5 GB/s) to 42.52 us (158.0 GB/s) across 97 calls/tok (-0.683 ms/tok).
  3. M5 moe_gu QUANT_CTAB Bypass for IQ3_S (window.cpp:1136) — -0.049 ms/tok:
    • While QUANT_CTAB=1 (function constant 23, reading i-quant grids from constant address space instead of threadgroup memory) speeds up IQ2_S (93.8 -> 99.3 GB/s) and IQ2_XXS (81.1 -> 83.6 GB/s), it slows down IQ3_S on M5 (105.8 GB/s vs 118.4 GB/s with threadgroup tab) because IQ3_S decodes 2 grid indices per loop iteration with wider lane divergence that serialized the M5 constant cache. Disabling ct when tg_.type == ref::IQ3_S recovers 118.4 GB/s on layers 0 and 1.
  4. M5 moe_gu Tile 16,1,4 for IQ2_XXS (window.cpp:1071) — -0.031 ms/tok:
    • Setting gl = 16 (16,1,4) for IQ2_XXS on M5 improves moe_gu from 103.6 us (81.6 GB/s) to 101.0 us (83.6 GB/s) across the 12 IQ2_XXS layers.
  5. Unperturbed Top-K Read-Ahead (STRATA_METAL_PREFETCH_UNPERT=1 default in stream.cpp:654):
    • In ExpertStream::ahead(), unpert_all previously defaulted to 0 (only ranking resident experts plus a small margin when moce_bonus_ > 0, which starves SSD read-ahead of true top-K missing experts). Defaulting STRATA_METAL_PREFETCH_UNPERT to 1 increases combined hit + prefetch-hit rate in our 48-layer K=512 simulator from 95.61% to 97.04% and cuts synchronous SSD stalls (miss/tok) by 32.6% (21.06 -> 14.20 misses/tok).
  6. Configurable Minimum Arena Slot Reservation (STRATA_METAL_MIN_SLOTS in window.cpp:1515):
    • Exposes STRATA_METAL_MIN_SLOTS (defaulting to pg_on ? (n_exp_ + 64) : 64 = 576 so prompt_chunk never deadlocks on victim()), allowing decode-focused or small-batch runs to lower the per-class floor and shift +88 slots (+132 MB) from the 2-layer IQ3_S class (Class 0) into the 34-layer IQ2_S class (Class 1), raising pure cache hit rate from 92.02% to 92.63% (21.06 -> 19.22 misses/tok).

2. Weight Traffic Breakdown for K=512 IQ2_XS (T = 1) #

Component / Tensor Role Calls/tok Quant Format(s) Shape (n_in x n_out) Weight Bytes/tok (MB) Share (%)
attn_qkv.weight 36 35x IQ4_XS, 1x IQ3_S 2560 x 10240 498.69 13.66%
moe_gu (gate_exps + up_exps, top-10) 48x10 34x IQ2_S, 12x IQ2_XXS, 2x IQ3_S 2560 x 640 (x2) 486.40 13.33%
hc_*_down + hc_*_up + inject 97 Q8_0 (converted at load) 10240 x 324 + 320 x 10240 679.62 18.62%
output.weight (LM head) 1 IQ4_XS 2560 x 248320 337.72 9.25%
attn_gate.weight 36 22x IQ4_XS, 11x IQ3_S, 3x Q6_K 2560 x 6144 296.88 8.13%
ssm_out.weight 36 29x IQ4_XS, 6x IQ3_S, 1x Q6_K 6144 x 2560 295.77 8.10%
moe_down (down_exps, top-10) 48x10 48x Q2_0 640 x 2560 221.18 6.06%
attn_q.weight 12 11x IQ4_XS, 1x IQ3_S 2560 x 12288 197.35 5.41%
ffn_gate_inp (router + pred1 + pred2) 48+47+46 48x BF16 + 93x Q8_0 2560 x 512 255.34 7.00%
attn_output.weight 12 10x IQ4_XS, 1x IQ3_S, 1x Q6_K 6144 x 2560 103.22 2.83%
ffn_{gate,up,down}_shexp.weight 48x3 IQ4_XS, IQ3_S, Q6_K, IQ4_NL, Q8_0, Q2_0 2560 x 640 / 640 x 2560 134.98 3.70%
ple_key.weight + ple_value.weight 2 BF16 2560 x 10240, 2560 x 2560 65.54 1.80%
indexer.{q,k}_proj + attn_{k,v} + ssm_{alpha,beta} 120 BF16, IQ4_XS, Q6_K, IQ3_S 2560 x {512,128,48} 76.63 2.10%
Total 3,649.78 MB 100.00%

3. M5 GPU Kernel Benchmarks (gemv_parity, expert_parity, gr_parity, gdn_parity, qsa_parity) #

3.1 Dense & Shared-Expert GEMV Tile Sweep on Apple M5 (ntok = 1) #

All measurements below were collected on the 10-core Apple M5 GPU using vendor/strata/metal/tests/gemv_parity --one <TYPE> <N_IN> <N_OUT> --ntok 1 (rotating copies across a 1 GB arena so no weights hit in SLC across dispatches).

Tensor Role (Count) Type Shape (n_in x n_out) Old Tile (L,R,S,K) Old Latency (us) Old BW (GB/s) New M5 Tile (L,R,S,K) New Latency (us) New BW (GB/s) Saved per Call (us) Total Saved / tok (ms)
output.weight (1x) IQ4_XS 2560 x 248320 16,1,4,1 (or 32,1,4,1) 2475.47 (3128.14) 136.4 (108.0) 32,2,4,1 2453.32 137.7 -22.15 (-674.82) -0.022 (-0.675)
attn_qkv (35x) IQ4_XS 2560 x 10240 16,1,4,1 107.71 129.3 32,2,4,1 106.54 130.7 -1.17 -0.041
attn_qkv (1x) IQ3_S 2560 x 10240 16,1,4,1 425.56 26.5 4,2,4,1 260.47 43.2 -165.09 -0.165
attn_gate (22x) IQ4_XS 2560 x 6144 16,1,4,1 65.41 127.7 32,2,4,1 64.32 129.9 -1.09 -0.024
attn_gate (11x) IQ3_S 2560 x 6144 16,1,4,1 321.26 21.0 32,2,4,1 59.73 113.1 -261.53 -2.877
attn_gate (3x) Q6_K 2560 x 6144 16,1,4,1 195.84 65.9 16,1,4,1 195.84 65.9 0.00 0.000
ssm_out + attn_out (39x) IQ4_XS 6144 x 2560 32,1,8,2 177.99 46.9 32,2,4,1 128.43 65.1 -49.56 -1.933
ssm_out + attn_out (7x) IQ3_S 6144 x 2560 32,1,8,2 233.69 28.9 4,2,4,1 132.08 51.2 -101.61 -0.711
ssm_out + attn_out (2x) Q6_K 6144 x 2560 32,1,8,2 218.64 59.0 16,1,4,1 106.50 121.2 -112.14 -0.224
attn_q (11x) IQ4_XS 2560 x 12288 16,1,4,1 202.89 82.4 8,1,8,1 156.39 106.9 -46.50 -0.512
attn_q (1x) IQ3_S 2560 x 12288 16,1,4,1 458.31 29.5 8,2,4,1 262.21 51.6 -196.10 -0.196
shexp gate+up (57x) IQ4_XS 2560 x 640 32,1,4,1 63.94 13.6 16,2,4,1 16.15 53.9 -47.79 -2.724
shexp gate+up (23x) IQ3_S 2560 x 640 32,1,4,1 25.00 28.2 32,2,4,1 8.77 80.2 -16.23 -0.373
shexp gate+up (16x) Q6_K 2560 x 640 32,1,4,1 102.07 13.2 32,2,4,1 67.59 19.9 -34.48 -0.552
pred1 + pred2 (93x) Q8_0 2560 x 512 32,1,4,1 62.97 22.1 32,2,4,1 51.70 26.9 -11.27 -1.048
  • Why gemv_x (gemvx, staging x in threadgroup memory) was NOT enabled for ntok = 1:
    • On M5, gemv (reading x through L1) beats gemv_x on output.weight (2,453.32 us / 137.7 GB/s vs 2,503.48 us / 134.9 GB/s), attn_qkv (106.54 us / 130.7 GB/s vs 109.88 us / 126.7 GB/s), attn_gate IQ4_XS (64.32 us / 129.9 GB/s vs 67.09 us / 124.6 GB/s), attn_gate IQ3_S (59.73 us / 113.1 GB/s vs 146.30 us / 46.2 GB/s), and ssm_out IQ3_S (132.08 us / 51.2 GB/s vs 184.18 us / 36.7 GB/s).

3.2 Routed MoE Kernel Sweep on Apple M5 (expert_parity --bench, ntok = 1, topk = 10) #

Measured with expert_parity --bench --gu <TYPE> --dn Q2_0 --layers 1 --reps 48 --ntok 1:

Gate/Up Type (Layers) QUANT_CTAB (fc 23) moe_gu Tile (L,R,S) moe_gu Time (us) moe_gu BW (GB/s) moe_down (Q2_0 8,2,4) (us) moe_down BW (GB/s) Total Layer Time (us) Total Layer BW (GB/s)
IQ2_S (34 layers) 0 (tg tab) 32,1,4 111.9 93.8 47.3 97.4 164.2 92.0
IQ2_S (34 layers) 1 (const) 32,1,4 105.4 99.6 47.5 97.1 157.9 95.7
IQ2_S (34 layers) 1 (const) 16,1,4 105.8 99.2 47.3 97.4 157.5 95.9
IQ2_S (34 layers) 1 (const) 32,2,4 167.2 62.8 47.5 96.9 219.6 68.8
IQ2_XXS (12 layers) 0 (tg tab) 32,1,4 108.0 78.2 47.4 97.2 159.6 81.8
IQ2_XXS (12 layers) 1 (const) 32,1,4 103.6 81.6 47.3 97.4 155.2 84.1
IQ2_XXS (12 layers) 1 (const) 16,1,4 101.0 83.6 47.2 97.6 152.5 85.6
IQ3_S (2 layers) 0 (tg tab) 32,1,4 118.9 118.4 47.3 97.4 170.8 109.4
IQ3_S (2 layers) 1 (const) 32,1,4 121.6 115.8 47.1 97.8 173.2 107.9
IQ3_S (2 layers) 1 (const) 16,1,4 133.0 105.8 47.2 97.6 184.6 101.3
  • Finding: NR = 2 (32,2,4) spills registers badly on M5 for i-quants (62.8 GB/s vs 99.6 GB/s on IQ2_S). NR = 1 is mandatory.
  • Finding: QUANT_CTAB=1 (metal_ctab()) is a clear win for IQ2_S (+5.8 GB/s) and IQ2_XXS (+5.4 GB/s with 16,1,4), whereas IQ3_S is fastest with threadgroup table (QUANT_CTAB=0, 118.4 GB/s). We set const bool ct = metal_ctab() && tg_.type != ref::IQ3_S; and if (!m4 && tg_.type == ref::IQ2_XXS) gl = 16;.

3.3 Hyper-Connection (gr.metal), GDN (gdn.metal), and QSA (qsa.metal) on Apple M5 #

  1. gr.metal (gr_parity --bench --ntok 1, GR_Q8=1, 48 rotated weight sets):
    • GR_DOWN=8_8 (current default, 1 row/tg, 324 threadgroups): gr_down = 49.56 us (135.5 GB/s), all 3 (norm+down+up) = 102.57 us (129.4 GB/s).
    • GR_DOWN=8_4 (new M5 choice, 2 rows/tg, 162 threadgroups): gr_down = 42.52 us (158.0 GB/s), all 3 = 97.83 us (135.7 GB/s); parity passes bitwise on prepared planes (R rel L2 2.90e-08, lo 1.24e-07, inject 6.23e-08, mixed 7.32e-08). Across 97 gr_read calls per token, gr_down_q8_8_4 saves ~0.683 ms/tok.
  2. gdn.metal (gdn_parity --bench, 36 rotated layer states):
    • gdn_step (IB = 4, 512 threads/tg): ntok = 1 commit = 74.27 us/layer (84.7 GB/s), ntok = 2..4 commit = 72.93..85.66 us/layer.
    • gdn_step8 (IB = 8, 1024 threads/tg): ntok = 1 commit = 72.49..85.24 us/layer (within noise for ntok = 1), but degrades badly for ntok = 3..4 (108.18..149.19 us/layer, +74% slower). Keeping gdn_step (IB = 4, 512 threads) is strictly better.
  3. qsa.metal (qsa_parity --bench 1024):
    • qsa_attn (v1, default): 5-dispatch window (ntok = 1) = 81.87 us (ntok = 3 = 155.42 us).
    • qsa_attn2 (v2): ntok = 1 = 78.75 us, but ntok = 3 degrades to 186.58 us (+20% slower).
    • qsa_attn3 (v3): ntok = 1 = 140.92 us (+72% slower). Keeping qsa_attn (v1) is optimal.

4. K=512 SSD Expert Stream & Arena Class Analysis (stream.cpp) #

4.1 Why Class 0 (IQ3_S, Layers 0..1) Reserves min_slots = 576 When STRATA_METAL_PROMPT_GEMM=1 #

We investigated whether Class 0 (which serves only 2 IQ3_S layers: blk.0 and blk.1, 2 x 512 = 1024 total experts, 115 pages = 1.884 MB/slot) can cap its minimum reservation below min_slots = 576 when STRATA_METAL_PROMPT_GEMM=1.

  • Finding (Critical Safety Invariant): In ExpertStream::at_layer(int l, int T, const int* ids, bool prompt, bool split) (stream.cpp:817), during prompt prefill (prompt == true), layer l pins all its active experts (stamp_) and then immediately queues background read-ahead of up to all 512 experts of layer l + 1 (for (int e = 0; e < ne_; ++e)), pinning them with the same stamp_ (slots_[s].pin = stamp_).
  • When l = 0 (blk.0, Class 0) runs during a prompt chunk, up to 512 experts of blk.0 are pinned with stamp_ = 1, and up to 512 experts of blk.1 (also Class 0!) are queued into Class 0's remaining unpinned slots and pinned with stamp_ = 1.
  • Then when l = 1 (blk.1, Class 0) enters at_layer (stamp_ = 2), pass 1 updates pin = 2 on all resident/queued blk.1 experts. Any remaining blk.1 misses call victim(0) (stream.cpp:800), which requires slots_[s].state == 0 && slots_[s].pin != stamp_. Every slot in Class 0 that was filled by blk.1's background stream has pin == 2! Thus, if Class 0 has < 512 total slots, blk.1 holds all of Class 0's slots with pin == 2 and victim(0) returns -1 -> strata metal stream: no free slot (layer 1), exit(6)!
  • Even a 2-layer class (blk.0 and blk.1) needs >= 512 slots (n_exp_ + 64 = 576) whenever a prompt chunk activates all 512 experts of blk.1.
  • However, for decode benchmarks or workloads with smaller prompt chunks where fewer than 512 experts are touched per chunk, lowering min_slots via STRATA_METAL_MIN_SLOTS=192 safely shifts +88 slots into Class 1 (IQ2_S + IQ2_XXS, Layers 2..47).

4.2 Closed-Loop 48-Layer K=512 Stream Simulation (k512_sim over 128 Real Token Steps) #

Using exact router weights (ffn_gate_inp.weight BF16 + Q8_0 predictors) and real hidden states from K512-iq2_xs-metal.gguf, /tmp/argon_decode_probes/k512_sim.cpp simulates the exact ExpertStream arena layout, STRP profile initialization (vendor/strata/data/expert-profile.bin), LRU eviction, MOCE routing (bonus = 0.25, protect = 3), and 1-layer/2-layer read-ahead across 128 decode steps (6,144 layer steps, 61,440 expert selections):

Stream Configuration (budget = 6.80 GB) Total Slots Pure Hit % Combined (Hit+Pre) % Urgent Misses / tok Prefetch Hits / tok SSD Read MB / tok
0. Baseline (w=0.15, pf=6+2, unpert=0, min_slots=576) 4,326 92.02% 95.61% 21.06 17.24 97.6 MB
1. CLASS_WASTE=0.25 (2 classes: 115, 80 pages) 3,894 91.69% 95.45% 21.85 18.03 101.6 MB
2. CLASS_WASTE=0.35 (1 class: 115 pages) 3,609 91.10% 95.09% 23.56 19.15 107.2 MB
3. STRATA_METAL_PREFETCH_UNPERT=1 (pf=6, pf2=0) 4,326 92.02% 97.04% 14.20 (-32.6%) 24.12 140.8 MB
4. STRATA_METAL_MIN_SLOTS=192 (decode-focused) 4,414 92.63% 96.00% 19.22 (-8.7%) 16.15 93.0 MB

5. Verification & Deliverable Checklist #

  • Patch location: /tmp/argon_decode_probes/argon_k512_decode_opt.diff
  • Tracked repository files: unmodified (git status --short and git -C vendor/strata diff untouched).
  • Syntax check (clang++ -fsyntax-only):
    • window.cpp: exit code 0
    • stream.cpp: exit code 0
  • Patch check (git apply --check):
    • git apply --check /tmp/argon_decode_probes/argon_k512_decode_opt.diff: exit code 0