# 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`)
| Kernel | Shape | Baseline (us) | Candidate (us) | Ratio | Delta | Parity |
|---|---|---:|---:|---:|---:|---|
| `mtp_rms_xprep` | `n=2560, nt=1` | 2.895 | 2.223 | 0.7676 | **-23.24%** | BITWISE PASS |
| `mtp_rms_xprep` | `n=2560, nt=4` | 2.932 | 2.338 | 0.7945 | **-20.55%** | BITWISE PASS |
| `mtp_rms_xprep` | `n=2560, nt=8` | 2.959 | 2.370 | 0.7981 | **-20.19%** | BITWISE PASS |
| `mtp_rms_xprep` | `n=10240, nt=1` | 7.893 | 5.060 | 0.6390 | **-36.10%** | BITWISE PASS |
| `mtp_rms_xprep` | `n=10240, nt=4` | 8.352 | 5.626 | 0.6745 | **-32.55%** | BITWISE PASS |
| `mtp_rms_xprep` | `n=10240, nt=8` | 8.398 | 5.801 | 0.6879 | **-31.21%** | BITWISE PASS |
| `mtp_combine_gnorm` | `hc=4, nt=1` | 2.901 | 2.888 | 0.9962 | **-0.38%** | BITWISE PASS |
| `mtp_combine_gnorm` | `hc=4, nt=4` | 3.097 | 3.067 | 0.9896 | **-1.04%** | BITWISE PASS |
| `mtp_combine_gnorm` | `hc=4, nt=8` | 3.804 | 3.798 | 0.9995 | **-0.05%** | BITWISE PASS |
| `mtp_xprep` | `rows=8, n=2560` | 2.154 | 2.133 | 0.9899 | **-1.01%** | BITWISE PASS |

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`)
| Kernel | Config | Baseline (us) | Candidate (us) | Ratio | Delta | Parity |
|---|---|---:|---:|---:|---:|---|
| `mtp_pick` | `nv=40525` | 3.501 | 3.508 | 0.9820 | **-1.80%** | BITWISE PASS |
| `mtp_pick` | `nv=248320` | 15.084 | 15.092 | 0.9996 | **-0.04%** | BITWISE PASS |
| `mtp_slop_pen` | `nv=40525, mode=0` | 24.895 | 0.876 | 0.0351 | **-96.49%** | BITWISE PASS |
| `mtp_slop_pen` | `nv=40525, mode=1` | 34.803 | 10.254 | 0.2944 | **-70.56%** | BITWISE PASS |
| `mtp_slop_pen` | `nv=40525, mode=2` | 41.072 | 14.822 | 0.3607 | **-63.93%** | BITWISE PASS |
| `mtp_slop_pen` | `nv=248320, mode=0` | 24.904 | 0.814 | 0.0327 | **-96.73%** | BITWISE PASS |
| `mtp_slop_pen` | `nv=248320, mode=1` | 57.137 | 31.043 | 0.5435 | **-45.65%** | BITWISE PASS |
| `mtp_slop_pen` | `nv=248320, mode=2` | 87.128 | 60.116 | 0.6903 | **-30.97%** | BITWISE PASS |
| `mtp_embed` | `format=30 (BF16)` | 1.618 | 1.112 | 0.6842 | **-31.58%** | BITWISE PASS |
| `mtp_embed` | `format=8 (Q8_0)` | 1.616 | 1.169 | 0.7225 | **-27.75%** | BITWISE PASS |
| `mtp_embed` | `format=2 (Q4_0)` | 1.592 | 1.177 | 0.7409 | **-25.91%** | BITWISE PASS |
| `mtp_embed` | `format=6 (Q5_0)` | 1.614 | 1.264 | 0.7822 | **-21.78%** | BITWISE PASS |
| `mtp_embed` | `format=20 (IQ4_NL)` | 1.714 | 1.289 | 0.7530 | **-24.70%** | BITWISE PASS |
| `mtp_embed` | `format=23 (IQ4_XS)` | 1.599 | 1.300 | 0.8144 | **-18.56%** | BITWISE PASS |
| `mtp_embed` | `format=42 (Q2_0)` | 1.639 | 1.182 | 0.7229 | **-27.71%** | BITWISE PASS |
| `mtp_embed` | `format=21 (IQ3_S)` | 1.730 | 1.451 | 0.8361 | **-16.39%** | BITWISE PASS |
| `mtp_embed` | `format=11 (Q3_K)` | 1.971 | 1.437 | 0.7294 | **-27.06%** | BITWISE PASS |
| `mtp_embed` | `format=12 (Q4_K)` | 1.923 | 1.440 | 0.7505 | **-24.95%** | BITWISE PASS |
| `mtp_embed` | `format=13 (Q5_K)` | 1.941 | 1.529 | 0.7855 | **-21.45%** | BITWISE PASS |
| `mtp_embed` | `format=14 (Q6_K)` | 1.719 | 1.346 | 0.7830 | **-21.70%** | BITWISE PASS |
| `mtp_embed` | `format=18 (IQ3_XXS)` | 1.622 | 1.366 | 0.8423 | **-15.77%** | BITWISE PASS |

### 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`.