Round 35 Track C: MTP Draft Kernels (mtp.metal / mtp_chain.metal) & router.metal Optimization Report #

• Agent: argon-r35-mtp-and-router-opt (gemini-4-argon)
• Patch: /tmp/argon_r35_mtp/argon_r35_mtp.patch (verified clean with git -C vendor/strata apply --check)
• Owned files:
  • metal/src/engine/msl/mtp.metal
  • metal/src/engine/msl/mtp_chain.metal
  • metal/kernels/router.metal
  • metal/tests/route_parity.cpp

1. Summary of Changes & Root Causes #

A. metal/src/engine/msl/mtp.metal (mtp_rms_xprep, mtp_combine_gnorm, mtp_xprep) #

1. mtp_rms_xprep:
   • Pre-offset + unrolled reduction (a.n == 2560 and a.n == 10240): Pre-offset xrp = xr + tid and specialized the two MTP widths (2560 for mtp.pre_fc_norm_embedding.weight, 10 steps of 256; 10240 for mtp.pre_fc_norm_hidden.weight, 40 steps of 256) with #pragma clang loop unroll(full) so all loads use constant immediate offsets [256 * i].
   • Prepared-output loop specialization & division elimination:
     • For a.n == 2560 && a.xp_n == 2560, a.n / 16 == 160 < 256, so each thread tid < 160 executes at most 1 iteration (if (tid < 160)), eliminating the loop back-edge and all three runtime integer divisions (a.n / a.xp_n, i / per, i % per).
     • For a.n == 10240 && a.xp_n == 2560, per == 160 and a.n / a.xp_n == 4 are compile-time constants (q = i / 160u, k = i - q * 160u), replacing expensive 32-bit runtime udiv/urem sequences on AGX with a single constant magic-multiply.
2. mtp_combine_gnorm:
   • Pre-offset rp = r + tid, o += tid, and w += tid in the n == 2560 fast path so all 40 unrolled loads/stores (hp, ep, rp, w, o) use constant immediate offsets [256 * i].
3. mtp_xprep:
   • Specialized per == 160u ? gid / 160u : gid / per and i = gid - t * per to avoid runtime integer division/modulo for n == 2560.

B. metal/src/engine/msl/mtp_chain.metal (mtp_embed, mtp_slop_pen, mtp_pick) #

1. mtp_embed<F> (-15.8% to -31.6% across all 13 GGUF formats):
   • Previously, mtp_embed<F> called X<F>::load(posv, posv, u, ps, xsum) on every unit to read 16 or 32 float indices from global buffer posv (0.0f .. 2559.0f), computed unused xsum[0..1], converted int(ps[j]) back to int, and performed 16 or 32 dynamic-index scalar scatter stores y[pos] = ... plus runtime-indexed kind_scale<F>(pos).
   • Because unit base offset p0 (and p1 for UNIT == 32) is always a multiple of 16 (16 * u, 32 * u, or the K-quant block/subblock formula), (p0 + j) & 7 == j & 7, (p0 + j) & 3 == j & 3, and (p0 + j) & 1 == j & 1.
   • Replaced X<F>::load(posv, ...) and scalar scatter stores with direct analytical unit offset p0/p1, compile-time constant kind_scale<F>(j), and coalesced float4 stores to y + p0 (and y + p1), plus if (U <= 256) { if (int(tid) < U) ... } loop elimination for n == 2560.
2. mtp_slop_pen (-31.0% to -96.7%):
   • Fast-exit when a.j == 0 && !any_knob (mode=0): Immediately writes state[0] = a.state on tid == 0 and returns without barriers or table scans (24.90 us -> 0.81 us, -96.7%).
   • Parallel token classification (kNewline / kBlank): Both slop4::kNewline (200 tokens) and slop4::kBlank (222 tokens) fit within the 256-thread threadgroup (tid < 200, tid < 222), replacing 6 serial 422-element constant-memory linear scans on tid == 0 with a single parallel compare across the threadgroup.
   • Zero thread stack frame (mtp_slop_penalties): Emits penalties directly into threadgroup int ids[80] and threadgroup float vals[80] using file-scope constant int32_t tables and #pragma clang loop unroll(full) instead of allocating slop4::Pen pen[80] (640B) + 20 local arrays on the per-thread DRAM stack.
   • Parallel inverse[] lookup & lg[ids[tid]] hold/removed evaluation: Replaced serial tid == 0 global memory loops over inverse[pen[i].id] and lg[ids[i]] with parallel evaluation across tid < count.

C. metal/kernels/router.metal & metal/tests/route_parity.cpp #

1. route_body in router.metal (-0.2% to -2.1%):
   • Unrolled the per-rank argmax and butterfly reduction loops and replaced if ((ex & 31) == lane) v[ex / 32] = -INFINITY; (dynamic register-array indexing that forced v[16] to spill to stack memory) with #pragma clang loop unroll(full) for (int i = 0; i < 16; ++i) v[i] = (ex == lane + 32 * i) ? -INFINITY : v[i]; so v[0..15] remains in scalar registers while keeping the initial threadgroup barrier that synchronizes X<F>::load across simdgroups.
2. route_parity.cpp:
   • Added --mtp-ab BASE_MTP_METAL for paired bitwise parity and timing of mtp_rms_xprep (n=2560 and n=10240), mtp_combine_gnorm, and mtp_xprep, and aligned buffer placement during paired timing in --router-ab.

2. Quiet-GPU Benchmarks & Bitwise Parity Verification (Apple M5) #

A. route_parity --mtp-ab (mtp.metal) #

In comb_fold_parity 0 mtp bench (cold 16-copy DRAM walk of the full hidden branch mtp_rms_xprep + fc_hidden GEMV + mtp_combine_gnorm):

• All n = 1..8 pass ALL PASS.
• Cold batched hidden branch time improved from 94.1 us -> 65.5 us (n=1), 102.0 us -> 96.3 us (n=2), 156.3 us -> 150.7 us (n=3), 190.4 us -> 184.9 us (n=4), 366.6 us -> 361.2 us (n=8).

B. sample_parity --mtp (mtp_chain.metal) #

C. route_parity --router-ab (router.metal) #

• type=30 (BF16): nt=1..8, split=0/1 all PASS with median ratios 0.9874 .. 0.9977 (-0.2% to -1.3%).
• type=8 (Q8_0): nt=1..8, split=0/1 all PASS with median ratios 0.9787 .. 0.9981 (-0.2% to -2.1%).
• Standalone route_parity (400 windows, K=512 + K=256 MoCE) and prompt_parity: ALL PASS.