8.3 kB 75 lines
2267f2a8
Round 28 blitz: mtp.fc_hidden [t][c] stream batching (-52.1% n=2 / -29.3% n=3), ple.metal tg_sum barrier halving (-1.5..-5.8%), sample.metal tg_argmax unroll + router_route_8
dawn
⋯
audits/argon_r28_ple.md
sourcerendered
download raw history
at main Markdown

Round 28 Track C: ple.metal tg_sum Barrier Halving & sample.metal tg_argmax Optimization #

Summary of Changes #

1. metal/src/engine/msl/ple.metal #

  • Halved tg_sum barriers (2 -> 1 per reduction) by removing the trailing threadgroup_barrier(mem_flags::mem_threadgroup) inside tg_sum and alternating part[0..7] and part[8..15] across back-to-back reductions (threadgroup float part[16] = 64 B vs 32 B):
    • ple_gnorm: 2 -> 1 barrier (-50%), using threadgroup float part[8].
    • ple_gate_norm: 4 -> 2 barriers (-50%), call 1 (acc) on part[0..7], call 2 (ss) on part + 8.
    • ple_query_gate_norm & ple_comb_query_gate_norm: 6 -> 3 barriers (-50%), gnorm_regs(q) on part[0..7], gate_norm_regs first reduction (acc) on part + 8 (whose barrier synchronizes prior readers of part[0..7]), and gate_norm_regs second reduction (ss) on part[0..7].
    • ple_moe_gate_norm & ple_moe_gate_norm_c (moe_gate_norm_row): 8 -> 4 barriers per stream (-50%), gnorm_regs(k) on part + 8, gnorm_regs(q) on part[0..7], and gate_norm_regs on part + 8 then part[0..7]. Across hc = 4 streams in ple_moe_gate_norm_c, total tg_sum barriers drop from 32 -> 16 (33 -> 17 including the initial roff setup barrier).
  • Added #pragma clang loop unroll(full) to both for (int r = 0; r < 9; ++r) loops in ple_hist_adv so float col[9] stays in registers without dynamic indexing.

2. metal/kernels/sample.metal #

  • Added #pragma clang loop unroll(full) to both fixed-trip loops in tg_argmax:
    • for (ushort m = 16; m > 0; m >>= 1) (5-step warp butterfly reduction).
    • for (int k = 0; k < 8; ++k) (8-way shared-memory reduction).
  • Evaluated & Rejected (tid == 0 pruning and unroll_count(4) in argmax_part / argmax_final):
    • if (sg == 0 && lane == 0) / if (tid == 0) around the post-barrier loop in tg_argmax: Because all 8 simdgroups must wait for threadgroup completion anyway and simdgroup 0 still executes the 8-way loop when tid == 0 is taken, adding divergent predication inside tg_argmax regressed nv = 40525 (nt = 1..8) by +2.0% to +3.5% and nv = 248064, nt = 1..2 by +0.8% to +2.0%. Leaving the unrolled 8-way reduction uniform across lanes avoids the branch and is faster.
    • #pragma clang loop unroll_count(4) in argmax_part and argmax_final: In argmax_part, each iteration contains if (!isnan(v)) better(bv, bi, v, int(i)). Unrolling by 4 (or 2) either via pragma or manual split-loop (i + 1024 < b1) bloated register/predicate state and regressed the draft-vocab path (nv = 40525, per = 338, only 1–2 iterations per thread) by +7.5% to +16.4% (+1.0% to +3.8% even with a per >= 2048 guard). In argmax_final, a.parts = 120 < 256, so the loop executes at most once (tid < 120) and unroll_count(4) added dead duplicate checks (+0.5% to +2.0% regression). Both loops are kept unrolled-by-1.

3. metal/tests/ple_moe_parity.cpp & metal/tests/sample_parity.cpp #

  • ple_moe_parity.cpp:
    • Added ple_gnorm, ple_gate_norm, ple_comb_query_gate_norm, and ple_comb_query_gate_norm + conv_res to --ab bitwise parity and paired A/B timing checks across T in {1, 3, 7, 9, 10, 64, 512, 1024, 2048}.
    • Deferred prompt.metal / moe_mm.metal compilation out of --ab mode and sliced only the needed kernels (gr_norm, moe_out, gr_norm_h*, gr_apply_moe*, mm_combine) in full-parity mode, cutting cold compilation + execution time from >20s (timeout) to 1.2s.
  • sample_parity.cpp:
    • Added nv = 248064 (nt = 1..8) and nv = 40525 (nt = 1, 2, 4, 8) to the paired A/B benchmark and bitwise parity sweep.
    • Reset lg to normal finite random floats before timing() (previously trial == 6 left -INFINITY in lg, skewing isnan/branch timing) and reported median per-pair ratios over 9 alternating pairs of 64 reps.

Verification & Measured A/B Results on Apple M5 #

1. ple_moe_parity (100% Bitwise Parity memcmp == 0 across all kernels & T) #

A/B T 1..2048: ple_gnorm PASS, ple_gate_norm PASS, ple_hist_adv PASS
A/B T 1..2048: ple_moe_gate_norm PASS, ple_moe_gate_norm_c PASS, ple_query_gate_norm PASS, ple_comb_query_gate_norm PASS
Kernel T = 1 (base -> cand, ratio) T = 3 (base -> cand, ratio) T = 7 (base -> cand, ratio) T = 64 (base -> cand, ratio) T = 2048 (base -> cand, ratio)
ple_gnorm (2 -> 1 bar) 0.0060 -> 0.0058 ms (0.9655, -3.5%) 0.0074 -> 0.0071 ms (0.9635, -3.7%) 0.0098 -> 0.0097 ms (0.9710, -2.9%) 0.0195 -> 0.0196 ms (1.0091, +0.9%) 1.2843 -> 1.2838 ms (1.0014, +0.1%)
ple_gate_norm (4 -> 2 bar) 0.0042 -> 0.0042 ms (0.9855, -1.5%) 0.0043 -> 0.0043 ms (0.9990, -0.1%) 0.0048 -> 0.0049 ms (1.0143, +1.4%) 0.0514 -> 0.0513 ms (0.9977, -0.2%) 3.3655 -> 3.3687 ms (0.9999, -0.0%)
ple_query_gate_norm (6 -> 3 bar) 0.0033 -> 0.0033 ms (0.9911, -0.9%) 0.0038 -> 0.0037 ms (0.9783, -2.2%) 0.0050 -> 0.0050 ms (0.9974, -0.3%) 0.0517 -> 0.0519 ms (1.0045, +0.5%) 3.3083 -> 3.3120 ms (1.0003, +0.0%)
ple_comb_query_gate_norm (6 -> 3 bar) 0.0067 -> 0.0067 ms (0.9933, -0.7%) 0.0088 -> 0.0087 ms (0.9936, -0.6%) 0.0115 -> 0.0114 ms (0.9875, -1.3%) 0.2365 -> 0.2361 ms (0.9985, -0.1%) 10.6802 -> 10.6833 ms (1.0012, +0.1%)
ple_moe_gate_norm (8 -> 4 bar) 0.0042 -> 0.0041 ms (0.9636, -3.6%) 0.0051 -> 0.0051 ms (0.9964, -0.4%) 0.0068 -> 0.0067 ms (0.9840, -1.6%) 0.0910 -> 0.0907 ms (0.9973, -0.3%) 4.6449 -> 4.6330 ms (0.9978, -0.2%)
ple_moe_gate_norm_c (33 -> 17 bar) 0.0245 -> 0.0240 ms (0.9785, -2.2%) 0.0257 -> 0.0252 ms (0.9807, -1.9%) 0.0264 -> 0.0258 ms (0.9765, -2.3%) 0.1105 -> 0.1099 ms (0.9852, -1.5%) 4.6122 -> 4.6190 ms (0.9988, -0.1%)
ple_hist_adv (unroll) 0.0029 -> 0.0029 ms (0.9955, -0.5%) 0.0029 -> 0.0030 ms (1.0121, +1.2%) 0.0032 -> 0.0032 ms (1.0042, +0.4%) 0.0023 -> 0.0023 ms (0.9902, -1.0%) 0.0024 -> 0.0024 ms (1.0001, +0.0%)

2. sample_parity (100% Bitwise Parity memcmp == 0 across all nv, nt, and edge-case trials) #

argmax nv=248320 nt=1 baseline 25.275 us candidate 25.307 us ratio 1.0026 (+0.26%)
argmax nv=248320 nt=2 baseline 39.385 us candidate 39.453 us ratio 0.9994 (-0.06%)
argmax nv=248320 nt=3 baseline 134.310 us candidate 134.270 us ratio 0.9996 (-0.04%)
argmax nv=248320 nt=4 baseline 178.275 us candidate 178.151 us ratio 1.0011 (+0.11%)
argmax nv=248320 nt=5 baseline 222.010 us candidate 221.879 us ratio 0.9995 (-0.05%)
argmax nv=248320 nt=6 baseline 265.524 us candidate 265.486 us ratio 0.9999 (-0.01%)
argmax nv=248320 nt=8 baseline 339.676 us candidate 337.559 us ratio 0.9911 (-0.89%)
argmax nv=248064 nt=1 baseline 5.432 us candidate 5.426 us ratio 1.0014 (+0.14%)
argmax nv=248064 nt=2 baseline 8.635 us candidate 8.628 us ratio 0.9999 (-0.01%)
argmax nv=248064 nt=3 baseline 131.288 us candidate 130.807 us ratio 0.9999 (-0.01%)
argmax nv=248064 nt=4 baseline 178.063 us candidate 178.022 us ratio 0.9998 (-0.02%)
argmax nv=248064 nt=5 baseline 206.575 us candidate 202.772 us ratio 0.9957 (-0.43%)
argmax nv=248064 nt=6 baseline 265.328 us candidate 265.270 us ratio 0.9999 (-0.01%)
argmax nv=248064 nt=7 baseline 299.301 us candidate 299.935 us ratio 1.0000 (+0.00%)
argmax nv=248064 nt=8 baseline 352.632 us candidate 352.615 us ratio 0.9999 (-0.01%)
argmax nv=40525 nt=1 baseline 3.287 us candidate 3.323 us ratio 1.0016 (+0.16%)
argmax nv=40525 nt=2 baseline 4.628 us candidate 4.622 us ratio 0.9983 (-0.17%)
argmax nv=40525 nt=4 baseline 7.477 us candidate 7.443 us ratio 0.9978 (-0.22%)
argmax nv=40525 nt=8 baseline 13.010 us candidate 13.027 us ratio 1.0011 (+0.11%)
argmax nv=7 nt=1 baseline 2.857 us candidate 2.728 us ratio 0.9549 (-4.51%)
argmax pv/pi/out bitwise + cpu: 0 wrong PASS

Patch Location #

  • Clean unified diff: /tmp/argon_r28_ple/argon_r28_ple.patch (verified with git -C vendor/strata apply --check).