# Round 22 Track C Audit Report: `moe_down` (`T_Q2_0R`/`T_Q2_0`) Compressed-Register Weight Reuse & `moe_route` / `plan_body` Optimization

- **Agent**: `argon-moe-down-q2r-and-route` (`gemini-4-argon`)
- **Scratch workspace**: `/tmp/argon_r22_moe/strata`
- **Unified patch**: `/tmp/argon_r22_moe/argon_r22_moe.patch` (`git -C vendor/strata apply --check` clean, exit code `0`)
- **Files modified**: `metal/kernels/moe.metal`, `metal/tests/expert_parity.cpp`

---

## 1. Key Findings & Design

### 1.1 `dn_body` & `gu_body`: Compressed-Register Weight Reuse (`q2_combine<F>`) Across `NR` & `NE`
Previously, in `dn_body` (`metal/kernels/moe.metal`):
- Holding `float q[NR][16]; float4 sc[NR];` (`40` registers for `NR = 2`) across `r = 0..NR-1` spilled registers on Apple M5 when `NE >= 2`, which forced `dn_body` for `F == T_Q2_0 || F == T_Q2_0R` to place `r` as the outer loop and reload `xs[16]` from `h` `NR = 2` separate times per token `t`.
- In `gu_body`, `F == T_Q2_0` at `NE >= 2` reloaded `xs[16]` `2 * NR = 4` times per token `t` (`g0` and `u0` for each `r`).

By introducing `q2_combine<F>(ushort2 pk, float d, thread const float* xs, float xsum0)`:
1. **Compressed register staging (`4` registers total for `NR = 2` in `dn_body`, `8` in `gu_body`)**:
   - In `dn_body` (`F == T_Q2_0 || F == T_Q2_0R`, across all `NE = 1..4`), we pre-load `ushort2 pk[NR]; float d[NR];` for all `r = 0..NR-1` first (sharing `(u >> 2) * 18` and `2 + 4 * (u & 3)` across rows), then loop `for (int t = 0; t < NE; ++t)`, loading `X<F>::load(h + (c + t) * int(a.n_ff), ...)` **once per token `t`** (`-50%` `h` device loads at `NR = 2, NE >= 2`) and accumulating `ac[r][t] += q2_combine<F>(pk[r], d[r], xs, xsum[0])` without ever materializing `float q[16]`.
   - In `gu_body` (`F == T_Q2_0 || F == T_Q2_0R`, `NE >= 2`), we pre-load `pkg[NR], pku[NR], dg[NR], du[NR]` (`8` registers for `NR = 2`), load `X<F>::load` **once per token `t`** (`-75%` `x` loads at `NR = 2, NE >= 2`), and accumulate `ag[r][t]` and `au[r][t]` via `q2_combine<F>`.
   - Added `#pragma clang loop unroll(full)` to the `r < NR` and `t < NE` loops in `dn_body` and `gu_body` (keeping `r < NR` un-pragmad in `gu_body` when `UNIT == 32` to avoid `qg[32]/qu[32]` register pressure in `Q4_K`).
2. **Exact IEEE-754 safe-math bitwise parity**:
   - `q2_combine<F>` preserves the exact 4-way FMA accumulation order (`s0, s1, s2, s3`) and scale/offset expression of `Dec<F>::run` + `unit_combine<16>`, achieving `std::memcmp == 0` across every buffer (`plan`, `h`, `ye`, `y`).

### 1.2 `moe_route` & `plan_body`: `v[16]` Indexing Hypothesis Falsified, `sg == 0` Parallel SIMD Prefix Sum Shipped
1. **Falsified hypothesis (`v[ex / 32] = -INFINITY` -> 16-way `if (i == wi)` loop)**:
   - Tested in isolation on Apple M5 (`route_parity --ab`): replacing `v[ex / 32] = -INFINITY;` with `#pragma clang loop unroll(full) for (int i = 0; i < 16; ++i) if (i == wi) v[i] = -INFINITY;` **regressed `moe_route` by +22% to +54%** (`5.12 us -> 7.89 us` at `ntok=1 split=0`, `6.95 us -> 9.76 us` at `ntok=1 split=1`).
   - Root cause: Apple Silicon GPU hardware supports native register-indirect indexing on `v[ex / 32] = -INFINITY;` (1 instruction in the winning lane), whereas the 16-way `if (i == wi)` loop emits 16 compare/select instructions per rank (160 extra instructions across 10 ranks) and increases live scalar register pressure in the 512-thread (`16`-simdgroup) threadgroup. Rejected per rule ("Only ship changes that are proven net wins on M5").
2. **Proven win in `plan_body` (used by `moe_route`, `moe_plan`, and `router_route`)**:
   - Previously, after the first `threadgroup_barrier`, thread `tid == 0` executed a serial `2 x 16 = 32`-iteration loop reading and writing `s_grp` and `s_ent` in threadgroup memory while the other 511 threads waited at the barrier, and unconditionally ran `simd_prefix_exclusive_sum(miss)` / `c1` even when `!split`.
   - Replaced the `tid == 0` serial loop with a single-instruction parallel `simd_prefix_exclusive_sum` across `sg == 0` (`lane < nsg`) and gated the `split` second-half prefix sums on `split`, cutting `moe_route` time by **-8.8% to -22.7%** and `moe_plan` time by **-14.6% to -35.4%** across `ntok = 1..8`.

---

## 2. Verification & Paired A/B Benchmarks on Apple M5

### 2.1 Bitwise Parity (`std::memcmp == 0`)
- `route_parity`: K=512 (200 windows) `0` mismatches, K=256+MoCE=0.85 (200 windows) `0` mismatches (`PASS`).
- `route_parity --ab /tmp/argon_r22_moe/moe_ref.metal`: 200 trials (`moe_route` + `moe_plan`, `split = 0, 1`, `ntok = 1..8`): **`0 fail(s) ALL PASS`**.
- `expert_parity --ab /tmp/argon_r22_moe/moe_ref.metal`: `Q2_0/Q2_0(142)`, `IQ2_XS/Q2_0(prep)`, `IQ3_XXS/Q5_1`, `Q4_K/IQ4_NL` across `NE = 1, 2, 3, 4, 8 (same)` and `ntok = 1, 2, 3, 4, 8 (random)`: **`0 fail(s) ALL PASS`** (`memcmp == 0` on `plan`, `h`, `ye`, `y`).

### 2.2 Paired Interleaved A/B Benchmarks (`expert_parity --ab ... --bench`)

#### A. `Q2_0` Gate/Up (`32,2,4`) + `T_Q2_0R` (`142`) Down (`moe_down_142_8_2_4`)
| Case | Groups | `dn` Base (us) | `dn` Cand (us) | `dn` Ratio | `gu` Base (us) | `gu` Cand (us) | `gu` Ratio |
|---|---:|---:|---:|---:|---:|---:|---:|
| `NE=1 (same)` | 10.0 | 47.2 | **46.0** | **0.9735 (-2.6%)** | 74.9 | **73.7** | **0.9907 (-0.9%)** |
| `NE=2 (same)` | 10.0 | 64.1 | **59.1** | **0.9229 (-7.7%)** | 110.9 | **97.3** | **0.8777 (-12.2%)** |
| `NE=3 (same)` | 10.0 | 82.8 | **75.7** | **0.9143 (-8.6%)** | 174.0 | **134.8** | **0.7753 (-22.5%)** |
| `NE=4 (same)` | 10.0 | 106.0 | **96.9** | **0.9151 (-8.5%)** | 219.1 | **174.1** | **0.7940 (-20.6%)** |
| `ntok=1 sh=0.3` | 10.0 | 47.2 | **46.0** | **0.9755 (-2.5%)** | 73.0 | **73.0** | **0.9995 (-0.1%)** |
| `ntok=2 sh=0.3` | 17.5 | 82.7 | **80.6** | **0.9746 (-2.5%)** | 127.3 | **126.3** | **0.9899 (-1.0%)** |
| `ntok=4 sh=0.3` | 31.0 | 151.1 | **144.3** | **0.9539 (-4.6%)** | 234.9 | **225.9** | **0.9618 (-3.8%)** |
| `ntok=8 sh=0.3` | 59.2 | 290.3 | **274.6** | **0.9452 (-5.5%)** | 461.8 | **435.0** | **0.9424 (-5.8%)** |

#### B. `IQ2_XS` Gate/Up (`32,2,4`) + `T_Q2_0` (`prep`) Down (`moe_down_134_8_2_4`)
| Case | Groups | `dn` Base (us) | `dn` Cand (us) | `dn` Ratio | `gu` Base (us) | `gu` Cand (us) | `gu` Ratio |
|---|---:|---:|---:|---:|---:|---:|---:|
| `NE=1 (same)` | 10.0 | 45.4 | **44.6** | **0.9843 (-1.6%)** | 101.8 | **93.4** | **0.9176 (-8.2%)** |
| `NE=2 (same)` | 10.0 | 65.6 | **55.1** | **0.8394 (-16.1%)** | 135.0 | **117.5** | **0.8702 (-13.0%)** |
| `NE=3 (same)` | 10.0 | 86.0 | **76.9** | **0.8937 (-10.6%)** | 197.5 | **163.7** | **0.8293 (-17.1%)** |
| `NE=4 (same)` | 10.0 | 112.4 | **97.7** | **0.8688 (-13.1%)** | 251.9 | **232.4** | **0.9223 (-7.8%)** |
| `ntok=1 sh=0.3` | 10.0 | 45.5 | **44.6** | **0.9805 (-1.9%)** | 102.0 | **93.6** | **0.9175 (-8.3%)** |
| `ntok=2 sh=0.3` | 17.5 | 77.2 | **74.8** | **0.9673 (-3.3%)** | 183.1 | **166.6** | **0.9104 (-9.0%)** |
| `ntok=4 sh=0.3` | 31.0 | 136.6 | **131.3** | **0.9611 (-3.9%)** | 341.5 | **310.6** | **0.9094 (-9.1%)** |
| `ntok=8 sh=0.3` | 59.2 | 260.6 | **245.3** | **0.9429 (-5.7%)** | 662.2 | **601.5** | **0.9086 (-9.1%)** |

#### C. `moe_route` & `moe_plan` Paired A/B (`route_parity --ab ... --bench`)
| Kernel | `split` | `ntok` | Base (us) | Cand (us) | Ratio | 48-Layer Savings |
|---|---:|---:|---:|---:|---:|---:|
| `moe_route` | 0 | 1 | 5.15 | **3.98** | **0.7727 (-22.7%)** | 0.247 -> 0.191 ms |
| `moe_route` | 0 | 2 | 5.68 | **4.53** | **0.7956 (-20.4%)** | 0.273 -> 0.217 ms |
| `moe_route` | 0 | 4 | 7.18 | **6.01** | **0.8377 (-16.2%)** | 0.345 -> 0.289 ms |
| `moe_route` | 0 | 8 | 9.87 | **8.79** | **0.8889 (-11.1%)** | 0.474 -> 0.422 ms |
| `moe_plan` | 0 | 1 | 3.07 | **1.99** | **0.6464 (-35.4%)** | 0.147 -> 0.095 ms |
| `moe_plan` | 0 | 2 | 3.68 | **2.60** | **0.7070 (-29.3%)** | 0.176 -> 0.125 ms |
| `moe_plan` | 0 | 4 | 5.06 | **3.98** | **0.7837 (-21.6%)** | 0.243 -> 0.191 ms |
| `moe_plan` | 0 | 8 | 7.26 | **6.21** | **0.8541 (-14.6%)** | 0.348 -> 0.298 ms |
| `moe_route` | 1 | 1 | 6.94 | **5.97** | **0.8564 (-14.4%)** | 0.333 -> 0.286 ms |
| `moe_route` | 1 | 2 | 7.43 | **6.44** | **0.8656 (-13.4%)** | 0.357 -> 0.309 ms |
| `moe_route` | 1 | 4 | 8.61 | **7.59** | **0.8842 (-11.6%)** | 0.413 -> 0.364 ms |
| `moe_route` | 1 | 8 | 11.52 | **10.50** | **0.9117 (-8.8%)** | 0.553 -> 0.504 ms |