5.5 kB 54 lines
f3384f69
Round 16 blitz: fused QSA Wk+Wv+Wik(+Wiq) single GEMM (-24..-36 dispatches, -5.6 MB scratch), decode resident-MoE + shared-expert wave zip (-48 barriers), gemv TAB>1 guards, ple_moe_gate_norm_c SIMD compaction + ple_conv half4/t>=9
dawn
⋯
audits/argon_r16_ple.md
sourcerendered
download raw history
at main Markdown

Round 16 Track C: ple_moe_gate_norm_c SIMD Slot Compaction + Stream-Loop Unrolling & ple_conv Vectorized half4 / Branchless t >= 9 Fast Path #

1. Summary of Changes #

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

  1. ple_moe_gate_norm_c SIMD Slot Compaction (roff[0..cnt-1]) + n == 2560 Unrolling (lines 176–263):
    • Replaced the threadgroup int rk[16] + per-iteration if (rk[k] >= 0) branch + 64-bit ulong(rk[k]) * n multiply inside the d loop with SIMD prefix compaction (simd_prefix_exclusive_sum(ushort(valid)), simd_sum(ushort(valid))) into threadgroup uint roff[17]; threadgroup float wk[16]; where roff[prefix] = uint(r) * n.
    • Reduced the inner expert accumulation loop to branchless 32-bit for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s); while preserving exact k accumulation order (memcmp == 0 under dropped slots pos == -1).
  2. Fixed NPT = 10 (n == 2560) Loop Unrolling + Pointer-Offset Hoisting & tg_sum Unrolling (lines 23–263):
    • Verified tg_sum retains its trailing threadgroup_barrier(mem_flags::mem_threadgroup) (required because gnorm_regs(k), gnorm_regs(q), and gate_norm_regs reuse threadgroup float part[8] across sequential reductions) and added #pragma clang loop unroll(full) to the 8-element reduction loop in tg_sum.
    • Added n == 2560 fast paths with #pragma clang loop unroll(full) and base-pointer hoisting (+ tid outside the loop, constant i * 256 displacements inside) across gnorm_regs, gate_norm_regs, ple_query_gate_norm, moe_gate_norm_row, ple_moe_gate_norm, and ple_moe_gate_norm_c.
  3. ple_conv Vectorized half4 Weight Load + Branchless t >= 9 Fast Path (lines 275–295):
    • Replaced the 4 scalar half tap loads (kw[k + 4 * ch]) with a single 8-byte aligned vectorized load const half4 w = *(device const half4*)(kw + 4 * ch);.
    • Added a branchless if (t >= 9) fast path (taken by 2,039 / 2,048 = 99.56% of rows in a 2048-token prompt chunk) that hoists device const float* nw = norm_w + (t - 9) * hcd + ch; const uint s3 = 3 * hcd; and accumulates nw[0], nw[s3], nw[2 * s3], nw[3 * s3] in identical sequential acc += ... order (memcmp == 0), eliminating all 4 per-tap p >= 0 ternary branches and hist address calculations.

Target 2: metal/kernels/gdn_prompt.metal (gdn_post, lines 433–460) #

  • Hoisted base row pointers (o_row = o + base, z_row = z + base, g_row = gamma + lane, yh_row = yh + base, y_row = y + base) so both the 4-channel RMSNorm preload (o_row[0], o_row[32], o_row[64], o_row[96]) and the unrolled k = 0..3 output loop use constant immediate offsets off = k * 32.
  • Verified 100% bitwise parity (memcmp == 0 on both yh and y for f32 = 0, 1 and 1-sg vs legacy grid) and up to 1.8% faster execution at T = 1024, 2048 (f32 = 1).

Target 3: metal/tests/ple_moe_parity.cpp #

  • Added --bench and --ab baseline.metal [T] modes testing bitwise parity (memcmp == 0) across T = 1, 3, 7, 9, 10, 64, 512, 1024, 2048 and interleaved 5-pair A/B timing at T = 512, 1024, 2048.

2. Bitwise Parity & Benchmark Results (Apple M5) #

ple_moe_parity Bitwise Verification (T = 1, 3, 7, 8, 9, 10, 64, 512, 1024, 2048) #

  • ple_moe_gate_norm + ple_conv_res: 100% bitwise PASS (memcmp == 0) across all tested T.
  • ple_moe_gate_norm_c + ple_conv_res (with dropped slots pos == -1): 100% bitwise PASS (memcmp == 0) across all tested T.
  • ple_query_gate_norm + ple_conv_res: 100% bitwise PASS (memcmp == 0) across all tested T.
  • ple_conv_res + ple_hist_adv: 100% bitwise PASS (memcmp == 0) across T = 1, 3, 7, 8, 9, 10, 64, 2048.

Paired A/B Timing (ple_moe_parity --ab, median of 5 interleaved pairs x 8 reps) #

Kernel / Stage T = 512 Base -> Cand (Ratio) T = 1024 Base -> Cand (Ratio) T = 2048 Base -> Cand (Ratio)
ple_moe_gate_norm_c 1.264 -> 1.133 ms (0.8966x, -10.3%) 2.455 -> 2.265 ms (0.9224x, -7.8%) 4.820 -> 4.496 ms (0.9321x, -6.8%)
ple_query_gate_norm 0.796 -> 0.777 ms (0.9832x, -1.7%) 1.614 -> 1.614 ms (1.0002x, +0.0%) 3.210 -> 3.237 ms (1.0079x, +0.8%)
ple_conv_res 0.477 -> 0.482 ms (1.0106x, +1.1%) 0.956 -> 0.957 ms (0.9994x, -0.1%) 1.907 -> 1.914 ms (1.0029x, +0.3%)
ple_moe_gate_norm_c + ple_conv_res 1.709 -> 1.608 ms (0.9426x, -5.7%) 3.391 -> 3.196 ms (0.9422x, -5.8%) 6.714 -> 6.404 ms (0.9533x, -4.7%)
ple_query_gate_norm + ple_conv_res 1.277 -> 1.273 ms (0.9993x, -0.1%) 2.567 -> 2.584 ms (1.0041x, +0.4%) 5.108 -> 5.139 ms (1.0060x, +0.6%)

gdn_prompt_parity --ab (gdn_post Stage 2, 36 layers, median of 3 interleaved pairs) #

  • T = 512, f32 = 1: 12.164 -> 12.162 ms (ratio 0.9998, yh PASS, y PASS)
  • T = 1024, f32 = 0: 16.071 -> 16.071 ms (ratio 1.0002, yh PASS, y PASS)
  • T = 1024, f32 = 1: 24.928 -> 24.410 ms (ratio 0.9823, -1.8%, yh PASS, y PASS)
  • T = 2048, f32 = 1: 50.341 -> 49.725 ms (ratio 0.9915, -0.9%, yh PASS, y PASS)

3. Artifacts #

  • Unified patch: /tmp/argon_r16_ple/argon_r16_ple.patch (verified clean with git -C vendor/strata apply --check).
  • Modified files in patch:
    • metal/src/engine/msl/ple.metal
    • metal/kernels/gdn_prompt.metal (lines 433–460 only)
    • metal/tests/ple_moe_parity.cpp