# Round 10 Track C Audit: `mm_gather_h` 128-Bit `uint4` Vectorization + `mm_scan`/`mm_place` Plan Optimization

- **Agent:** `argon-mm-gather-and-plan` (`opencode-go/gemini-4-argon-contributor`)
- **Scratch workspace:** `/tmp/argon_r10_gather/strata`
- **Unified patch:** `/tmp/argon_r10_gather/argon_r10_gather.patch` (verified clean with `git -C vendor/strata apply --check`)
- **Files touched in patch:** `metal/kernels/moe_mm.metal`, `metal/tests/moe_mm_parity.cpp` (`metal/src/engine/window.cpp` line 3540 `g.tg[0] = 256` kept unchanged; compatible with `prompt_chain.inc` and `prompt_chain.cpp` too).

---

## 1. `mm_gather_h` Investigation (`half4` vs `uint4`, Threadgroup Sizing, Unrolling & Multi-Row-per-TG)

### 1.1 Roofline & Geometry Analysis on Apple M5 (24 GB Unified Memory, 153.6 GB/s Peak DRAM Bandwidth)
For `N_EMBD = 2560` (`5,120` bytes per row = `640` `half4`s = `320` `uint4`s):
- At `T = 512` (`512` experts, `TM = 16`): `tot_rows = 8,432` -> `XS` write footprint is `43.17 MB` per layer (`2.07 GB` across 48 layers) plus gathered reads from `xh` (`2.62 MB`).
- At `T = 1024`: `tot_rows = 14,352` -> `XS` write footprint is `73.48 MB` per layer (`3.53 GB` across 48 layers) plus gathered reads from `xh` (`5.24 MB`).
- At `T = 2048`: `tot_rows = 24,576` -> `XS` write footprint is `125.83 MB` per layer (`6.04 GB` across 48 layers) plus scattered reads from `xh` (`10.49 MB`). Because `125.83 MB` exceeds the M5 System Level Cache (SLC), every layer's `XS` write and subsequent read stream through DRAM at **~130–153 GB/s (85–99% of hardware DRAM peak)**.

### 1.2 Paired 48-Layer Sweep Results (`mm_gather_h` Variants)
Measured via paired alternating 48-dispatch command buffers (`7–9` pairs after GPU warmup) on Apple M5:

| Variant | Vector Type | `tg[0]` | Grid (`T=2048`) | `T=512` (48L ms / us/L) | `T=1024` (48L ms / us/L) | `T=2048` (48L ms / us/L) | `XS` `memcmp` |
|---|---|---:|---:|---:|---:|---:|---|
| Baseline (`half4`, fixed stride 256) | `half4` (64-bit) | 256 | 24,576 | 13.153 ms (274.0 us) | 23.455 ms (488.6 us) | 55.543 ms (1157.1 us) | Ref |
| `half4_tg32` (`nth` stride) | `half4` (64-bit) | 32 | 24,576 | 13.555 ms (282.4 us) | 24.314 ms (506.5 us) | 56.517 ms (1177.4 us) | PASS (`0`) |
| `half4_tg64` (`nth` stride) | `half4` (64-bit) | 64 | 24,576 | 13.446 ms (280.1 us) | 23.932 ms (498.6 us) | 56.172 ms (1170.3 us) | PASS (`0`) |
| `half4_tg128` (`nth` stride) | `half4` (64-bit) | 128 | 24,576 | 13.307 ms (277.2 us) | 23.611 ms (491.9 us) | 55.931 ms (1165.2 us) | PASS (`0`) |
| `uint4_tg32` (`nth` stride) | `uint4` (128-bit) | 32 | 24,576 | 13.265 ms (276.4 us) | 23.372 ms (486.9 us) | 56.978 ms (1187.1 us) | PASS (`0`) |
| `uint4_tg64` (`nth` stride, 5 iters) | `uint4` (128-bit) | 64 | 24,576 | 13.080 ms (272.5 us) | 23.038 ms (480.0 us) | 56.063 ms (1168.0 us) | PASS (`0`) |
| `uint4_u64_tg64` (unrolled 5x `uint4`) | `uint4` (128-bit) | 64 | 24,576 | 13.212 ms (275.2 us) | 23.061 ms (480.4 us) | 56.177 ms (1170.3 us) | PASS (`0`) |
| `uint4_tg128` (`nth` stride) | `uint4` (128-bit) | 128 | 24,576 | 13.022 ms (271.3 us) | 22.979 ms (478.7 us) | 55.820 ms (1162.9 us) | PASS (`0`) |
| **`uint4_tg256` (`nth` stride, selected)** | **`uint4` (128-bit)** | **256** | **24,576** | **12.992 ms (270.7 us)** | **22.955 ms (478.2 us)** | **55.508 ms (1156.4 us)** | **PASS (`0`)** |
| `mr64_R2_tg128` (2 rows/TG, 64 thr/row) | `uint4` (128-bit) | 128 | 12,288 | 13.104 ms (273.0 us) | 23.156 ms (482.4 us) | 56.628 ms (1179.8 us) | PASS (`0`) |
| `mr64_R4_tg256` (4 rows/TG, 64 thr/row) | `uint4` (128-bit) | 256 | 6,144 | 13.080 ms (272.5 us) | 23.167 ms (482.6 us) | 56.563 ms (1178.4 us) | PASS (`0`) |
| `mr32_R4_tg128` (4 rows/TG, 32 thr/row) | `uint4` (128-bit) | 128 | 6,144 | 13.189 ms (274.8 us) | 23.393 ms (487.4 us) | 57.298 ms (1193.7 us) | PASS (`0`) |

### 1.3 Key Findings for `mm_gather_h`
1. **128-bit `uint4` copies beat 64-bit `half4` copies across all threadgroup sizes (`32, 64, 128, 256`)** at `T=512` (`-0.16 ms` / `-1.2%`) and `T=1024` (`-0.50 ms` / `-2.1%`), and tie/edge out baseline at `T=2048` (`55.508 ms` vs `55.543 ms`), because 128-bit loads/stores halve the number of loop iterations (`320` `uint4`s vs `640` `half4`s per row).
2. **Why `tg[0] = 256` with `uint4` beats `tg[0] = 64` at `T = 2048` (`55.51 ms` vs `56.06 ms`):**
   - With `uint4` and `tg[0] = 256`, iteration 0 issues `256 * 16B = 4,096B` (`32` cache lines of `128B`) of the `5,120B` row concurrently across 8 simdgroups, and iteration 1 issues the remaining `64 * 16B = 1,024B` (`8` cache lines) on the first 2 simdgroups (`tid < 64`). Because `64` is an exact multiple of the 32-wide SIMD width, simdgroups `2..7` exit cleanly at the SIMD level without intra-warp divergence, while launching 8 simdgroups per row maximizes memory-level parallelism (MLP) when `xh` (`10.49 MB` at `T=2048`) spills out of L1/L2 into SLC/DRAM.
   - Making `mm_gather_h` take `ushort nth [[threads_per_threadgroup]]` (with scalar `uint tg [[threadgroup_position_in_grid]]` per MSL scalar/vector attribute rules) supports any `tg[0] = 32, 64, 128, 256` while keeping `g.tg[0] = 256` in `window.cpp`, `prompt_chain.inc`, and `prompt_chain.cpp`.

---

## 2. Target 2 Investigation: `mm_scan` and `mm_place` Plan Kernels

We also inspected and benchmarked `mm_hist`, `mm_scan`, and `mm_place` in `metal/kernels/moe_mm.metal`:
1. **`mm_scan` barrier + dead-slot elimination:**
   - Previously, `mm_scan` zeroed `s_rows[0..31]` and `s_tiles[0..31]`, executed an extra `threadgroup_barrier(mem_flags::mem_threadgroup)`, and then looped `for (int i = 0; i < 32; ++i)` after the prefix sum barrier because 512 threads only populate `16` of the `32` simdgroup slots.
   - By taking `ushort nsg [[simdgroups_per_threadgroup]]` and looping `for (uint i = 0; i < nsg; ++i)`, we eliminate both the initial zero-fill and the first `threadgroup_barrier` (cutting `mm_scan`'s barriers in half from 2 to 1).
2. **`mm_place` threadgroup-memory & barrier reduction (`s_id[1024]` -> register `simd_broadcast`):**
   - Previously, `mm_place` allocated `threadgroup short s_id[1024]` (`2 KiB`), wrote `s_id[tid] = short(e)` after zeroing `p_cnt`, and after the second barrier did a threadgroup-memory `short4` scan over the caller's 32-lane sub-block.
   - Because each 32-lane sub-block (`sub = tid >> 5`) is an exact 32-lane Metal simdgroup, each lane can compute its intra-simdgroup rank directly in registers via `simd_broadcast(e, l)` while waiting for the `p_cnt` atomic updates, eliminating `threadgroup short s_id[1024]` completely.

### 2.1 Paired 48-Layer Plan Timing Results (`mm_hist + mm_scan + mm_place`)

| Kernel | `T=512` Before -> After (48L) | `T=1024` Before -> After (48L) | `T=2048` Before -> After (48L) | Bitwise `memcmp` (`bcnt`, `boff`, `eoff`, `tiles`, `meta`, `stok`, `sw`, `pos`) |
|---|---:|---:|---:|---|
| `mm_hist` | 0.111 ms -> 0.109 ms | 0.049 ms -> 0.048 ms | 0.060 ms -> 0.052 ms | PASS (`0`) |
| `mm_scan` | 0.354 ms -> 0.322 ms (`-9.0%`) | 0.214 ms -> 0.200 ms (`-6.5%`) | 0.265 ms -> 0.250 ms (`-5.7%`) | PASS (`0`) |
| `mm_place` | 0.309 ms -> 0.281 ms (`-9.1%`) | 0.192 ms -> 0.165 ms (`-14.1%`) | 0.290 ms -> 0.271 ms (`-6.6%`) | PASS (`0`) |
| **Full Plan (48L)** | **1.663 ms -> 1.160 ms** | **0.445 ms -> 0.408 ms (`-8.3%`)** | **0.611 ms -> 0.573 ms (`-6.2%`)** | **PASS (`0`)** |

---

## 3. Verification (`memcmp == 0` & `moe_mm_parity`)

- Added fast (`1.40s` total runtime) `--gather-bench` mode to `metal/tests/moe_mm_parity.cpp` (using `#ifdef MM_GATHER_ONLY` to skip compiling the heavy GEMM template matrix when testing/benchmarking plan + gather):
  - Verifies **100% bitwise identity (`memcmp == 0`)** on `XS` (across all `tg[0] in {32, 64, 128, 256}` against `mm_gather`) and exact plan invariants across **`T = 1, 7, 64, 256, 512, 1024, 2048`**.
  - Benchmarks 48-layer command buffers for both `plan` and `mm_gather_h` across `T = 512, 1024, 2048`.
- Full end-to-end `moe_mm_parity --v3 --tile 16,32,32,4 --dtile 64,32,2` verified `PASS` (`rel L2 1.11e-05` at `T=1`, `3.20e-06` at `T=7`, `2.65e-06` at `T=64`).