# Round 28 Track C: `ple.metal` `tg_sum` Barrier Halving & `sample.metal` `tg_argmax` Optimization

## Summary of Changes

### 1. `metal/src/engine/msl/ple.metal`
- **Halved `tg_sum` barriers (`2 -> 1` per reduction)** by removing the trailing `threadgroup_barrier(mem_flags::mem_threadgroup)` inside `tg_sum` and alternating `part[0..7]` and `part[8..15]` across back-to-back reductions (`threadgroup float part[16]` = 64 B vs 32 B):
  - `ple_gnorm`: `2 -> 1` barrier (`-50%`), using `threadgroup float part[8]`.
  - `ple_gate_norm`: `4 -> 2` barriers (`-50%`), call 1 (`acc`) on `part[0..7]`, call 2 (`ss`) on `part + 8`.
  - `ple_query_gate_norm` & `ple_comb_query_gate_norm`: `6 -> 3` barriers (`-50%`), `gnorm_regs(q)` on `part[0..7]`, `gate_norm_regs` first reduction (`acc`) on `part + 8` (whose barrier synchronizes prior readers of `part[0..7]`), and `gate_norm_regs` second reduction (`ss`) on `part[0..7]`.
  - `ple_moe_gate_norm` & `ple_moe_gate_norm_c` (`moe_gate_norm_row`): `8 -> 4` barriers per stream (`-50%`), `gnorm_regs(k)` on `part + 8`, `gnorm_regs(q)` on `part[0..7]`, and `gate_norm_regs` on `part + 8` then `part[0..7]`. Across `hc = 4` streams in `ple_moe_gate_norm_c`, total `tg_sum` barriers drop from `32 -> 16` (`33 -> 17` including the initial `roff` setup barrier).
- **Added `#pragma clang loop unroll(full)`** to both `for (int r = 0; r < 9; ++r)` loops in `ple_hist_adv` so `float col[9]` stays in registers without dynamic indexing.

### 2. `metal/kernels/sample.metal`
- **Added `#pragma clang loop unroll(full)`** to both fixed-trip loops in `tg_argmax`:
  - `for (ushort m = 16; m > 0; m >>= 1)` (5-step warp butterfly reduction).
  - `for (int k = 0; k < 8; ++k)` (8-way shared-memory reduction).
- **Evaluated & Rejected (`tid == 0` pruning and `unroll_count(4)` in `argmax_part` / `argmax_final`):**
  - *`if (sg == 0 && lane == 0)` / `if (tid == 0)` around the post-barrier loop in `tg_argmax`*: Because all 8 simdgroups must wait for threadgroup completion anyway and simdgroup 0 still executes the 8-way loop when `tid == 0` is taken, adding divergent predication inside `tg_argmax` regressed `nv = 40525` (`nt = 1..8`) by `+2.0%` to `+3.5%` and `nv = 248064, nt = 1..2` by `+0.8%` to `+2.0%`. Leaving the unrolled 8-way reduction uniform across lanes avoids the branch and is faster.
  - *`#pragma clang loop unroll_count(4)` in `argmax_part` and `argmax_final`*: In `argmax_part`, each iteration contains `if (!isnan(v)) better(bv, bi, v, int(i))`. Unrolling by 4 (or 2) either via pragma or manual split-loop (`i + 1024 < b1`) bloated register/predicate state and regressed the draft-vocab path (`nv = 40525`, `per = 338`, only 1–2 iterations per thread) by `+7.5%` to `+16.4%` (`+1.0%` to `+3.8%` even with a `per >= 2048` guard). In `argmax_final`, `a.parts = 120 < 256`, so the loop executes at most once (`tid < 120`) and `unroll_count(4)` added dead duplicate checks (`+0.5%` to `+2.0%` regression). Both loops are kept unrolled-by-1.

### 3. `metal/tests/ple_moe_parity.cpp` & `metal/tests/sample_parity.cpp`
- **`ple_moe_parity.cpp`**:
  - Added `ple_gnorm`, `ple_gate_norm`, `ple_comb_query_gate_norm`, and `ple_comb_query_gate_norm + conv_res` to `--ab` bitwise parity and paired A/B timing checks across `T in {1, 3, 7, 9, 10, 64, 512, 1024, 2048}`.
  - Deferred `prompt.metal` / `moe_mm.metal` compilation out of `--ab` mode and sliced only the needed kernels (`gr_norm`, `moe_out`, ` gr_norm_h*`, `gr_apply_moe*`, `mm_combine`) in full-parity mode, cutting cold compilation + execution time from `>20s` (timeout) to `1.2s`.
- **`sample_parity.cpp`**:
  - Added `nv = 248064` (`nt = 1..8`) and `nv = 40525` (`nt = 1, 2, 4, 8`) to the paired A/B benchmark and bitwise parity sweep.
  - Reset `lg` to normal finite random floats before `timing()` (previously `trial == 6` left `-INFINITY` in `lg`, skewing `isnan`/branch timing) and reported median per-pair ratios over 9 alternating pairs of 64 reps.

---

## Verification & Measured A/B Results on Apple M5

### 1. `ple_moe_parity` (`100%` Bitwise Parity `memcmp == 0` across all kernels & `T`)
```
A/B T 1..2048: ple_gnorm PASS, ple_gate_norm PASS, ple_hist_adv PASS
A/B T 1..2048: ple_moe_gate_norm PASS, ple_moe_gate_norm_c PASS, ple_query_gate_norm PASS, ple_comb_query_gate_norm PASS
```

| Kernel | `T = 1` (base -> cand, ratio) | `T = 3` (base -> cand, ratio) | `T = 7` (base -> cand, ratio) | `T = 64` (base -> cand, ratio) | `T = 2048` (base -> cand, ratio) |
|---|---:|---:|---:|---:|---:|
| `ple_gnorm` (`2 -> 1` bar) | `0.0060 -> 0.0058 ms` (`0.9655`, `-3.5%`) | `0.0074 -> 0.0071 ms` (`0.9635`, `-3.7%`) | `0.0098 -> 0.0097 ms` (`0.9710`, `-2.9%`) | `0.0195 -> 0.0196 ms` (`1.0091`, `+0.9%`) | `1.2843 -> 1.2838 ms` (`1.0014`, `+0.1%`) |
| `ple_gate_norm` (`4 -> 2` bar) | `0.0042 -> 0.0042 ms` (`0.9855`, `-1.5%`) | `0.0043 -> 0.0043 ms` (`0.9990`, `-0.1%`) | `0.0048 -> 0.0049 ms` (`1.0143`, `+1.4%`) | `0.0514 -> 0.0513 ms` (`0.9977`, `-0.2%`) | `3.3655 -> 3.3687 ms` (`0.9999`, `-0.0%`) |
| `ple_query_gate_norm` (`6 -> 3` bar) | `0.0033 -> 0.0033 ms` (`0.9911`, `-0.9%`) | `0.0038 -> 0.0037 ms` (`0.9783`, `-2.2%`) | `0.0050 -> 0.0050 ms` (`0.9974`, `-0.3%`) | `0.0517 -> 0.0519 ms` (`1.0045`, `+0.5%`) | `3.3083 -> 3.3120 ms` (`1.0003`, `+0.0%`) |
| `ple_comb_query_gate_norm` (`6 -> 3` bar) | `0.0067 -> 0.0067 ms` (`0.9933`, `-0.7%`) | `0.0088 -> 0.0087 ms` (`0.9936`, `-0.6%`) | `0.0115 -> 0.0114 ms` (`0.9875`, `-1.3%`) | `0.2365 -> 0.2361 ms` (`0.9985`, `-0.1%`) | `10.6802 -> 10.6833 ms` (`1.0012`, `+0.1%`) |
| `ple_moe_gate_norm` (`8 -> 4` bar) | `0.0042 -> 0.0041 ms` (`0.9636`, `-3.6%`) | `0.0051 -> 0.0051 ms` (`0.9964`, `-0.4%`) | `0.0068 -> 0.0067 ms` (`0.9840`, `-1.6%`) | `0.0910 -> 0.0907 ms` (`0.9973`, `-0.3%`) | `4.6449 -> 4.6330 ms` (`0.9978`, `-0.2%`) |
| `ple_moe_gate_norm_c` (`33 -> 17` bar) | `0.0245 -> 0.0240 ms` (`0.9785`, `-2.2%`) | `0.0257 -> 0.0252 ms` (`0.9807`, `-1.9%`) | `0.0264 -> 0.0258 ms` (`0.9765`, `-2.3%`) | `0.1105 -> 0.1099 ms` (`0.9852`, `-1.5%`) | `4.6122 -> 4.6190 ms` (`0.9988`, `-0.1%`) |
| `ple_hist_adv` (unroll) | `0.0029 -> 0.0029 ms` (`0.9955`, `-0.5%`) | `0.0029 -> 0.0030 ms` (`1.0121`, `+1.2%`) | `0.0032 -> 0.0032 ms` (`1.0042`, `+0.4%`) | `0.0023 -> 0.0023 ms` (`0.9902`, `-1.0%`) | `0.0024 -> 0.0024 ms` (`1.0001`, `+0.0%`) |

### 2. `sample_parity` (`100%` Bitwise Parity `memcmp == 0` across all `nv`, `nt`, and edge-case trials)
```
argmax nv=248320 nt=1 baseline 25.275 us candidate 25.307 us ratio 1.0026 (+0.26%)
argmax nv=248320 nt=2 baseline 39.385 us candidate 39.453 us ratio 0.9994 (-0.06%)
argmax nv=248320 nt=3 baseline 134.310 us candidate 134.270 us ratio 0.9996 (-0.04%)
argmax nv=248320 nt=4 baseline 178.275 us candidate 178.151 us ratio 1.0011 (+0.11%)
argmax nv=248320 nt=5 baseline 222.010 us candidate 221.879 us ratio 0.9995 (-0.05%)
argmax nv=248320 nt=6 baseline 265.524 us candidate 265.486 us ratio 0.9999 (-0.01%)
argmax nv=248320 nt=8 baseline 339.676 us candidate 337.559 us ratio 0.9911 (-0.89%)
argmax nv=248064 nt=1 baseline 5.432 us candidate 5.426 us ratio 1.0014 (+0.14%)
argmax nv=248064 nt=2 baseline 8.635 us candidate 8.628 us ratio 0.9999 (-0.01%)
argmax nv=248064 nt=3 baseline 131.288 us candidate 130.807 us ratio 0.9999 (-0.01%)
argmax nv=248064 nt=4 baseline 178.063 us candidate 178.022 us ratio 0.9998 (-0.02%)
argmax nv=248064 nt=5 baseline 206.575 us candidate 202.772 us ratio 0.9957 (-0.43%)
argmax nv=248064 nt=6 baseline 265.328 us candidate 265.270 us ratio 0.9999 (-0.01%)
argmax nv=248064 nt=7 baseline 299.301 us candidate 299.935 us ratio 1.0000 (+0.00%)
argmax nv=248064 nt=8 baseline 352.632 us candidate 352.615 us ratio 0.9999 (-0.01%)
argmax nv=40525 nt=1 baseline 3.287 us candidate 3.323 us ratio 1.0016 (+0.16%)
argmax nv=40525 nt=2 baseline 4.628 us candidate 4.622 us ratio 0.9983 (-0.17%)
argmax nv=40525 nt=4 baseline 7.477 us candidate 7.443 us ratio 0.9978 (-0.22%)
argmax nv=40525 nt=8 baseline 13.010 us candidate 13.027 us ratio 1.0011 (+0.11%)
argmax nv=7 nt=1 baseline 2.857 us candidate 2.728 us ratio 0.9549 (-4.51%)
argmax pv/pi/out bitwise + cpu: 0 wrong PASS
```

## Patch Location
- Clean unified diff: `/tmp/argon_r28_ple/argon_r28_ple.patch` (verified with `git -C vendor/strata apply --check`).