# Round 14 Track C: `gdn_scanr_4` `hk`-Grouped Head Locality & `gr_norm_h*` / `moe_comb` Optimization Report

## 1. Summary & Deliverables
- **Scratch Workspace:** `/tmp/argon_r14_gdn_gr/strata`
- **Unified Patch:** `/tmp/argon_r14_gdn_gr/argon_r14_gdn_gr.patch` (verified clean via `git -C vendor/strata apply --check`)
- **Files Modified (within Track C ownership boundaries):**
  - `metal/kernels/gdn_prompt.metal` (`gdn_scan_r` lines 225–227)
  - `metal/kernels/prompt.metal` (`gr_norm_h*`, `moe_comb_slots`, `moe_comb`, `gr_apply_moec` lines 170–440)
  - `metal/tests/comb_fold_parity.cpp` (`grvec` paired A/B 5-sample median reporting)
- **Bitwise Parity:** `100%` bitwise identical (`memcmp == 0`) across all `gdn_prompt_parity` (`T = 1, 3, 16, 17, 32, 37, 300, 512, 1024, 2048`, `f32 = 0, 1`) and `comb_fold_parity` (`T = 1, 7, 64, 512, 1024, 2048`) checks.

---

## 2. Target 1: `gdn_scan_r` (`gdn_scanr_4`) `hk`-Grouped Head Scheduling (`metal/kernels/gdn_prompt.metal`)

### 2.1 What Worked (Kept)
- **`hk`-Grouped Head Remapping (`hv == 3 * hkn`):**
  In Qwen3.8-Flash-Next, `hv = 48` value heads share `hkn = 16` query/key heads (`hk = h % 16`). Previously, `h = int(tg.x) / 4` scheduled `tg.x = 0..63` across `hk = 0..15`, separating the 3 value heads (`12` threadgroups: `h = hk, hk + 16, hk + 32`) that read the exact same `qn` and `kn` rows (`2 MB` per `hk` at `T = 2048`) by 64 threadgroups.
  Remapping:
  ```metal
  const int g_idx = int(tg.x) / 4, hv = int(a.h_v), hkn = int(a.h_k);
  const int hk = (hv == 3 * hkn) ? (g_idx / 3) : (g_idx % hkn);
  const int h = (hv == 3 * hkn) ? (hk + (g_idx % 3) * hkn) : g_idx;
  ```
  schedules all `12` threadgroups sharing a given `hk` consecutively (`tg.x = 0..11` share `hk = 0`, `tg.x = 12..23` share `hk = 1`, etc.), keeping `qn` and `kn` hot in L1/L2 cache across the 3 value heads. Because each `(h, j)` threadgroup writes exclusively to its own slice of `state` and `o`, this is `100%` bitwise identical (`memcmp == 0`) by construction.

### 2.2 What Was Ruled Out (Discarded)
- **`sg == 0`-Only `qa` (`s.y`) Reduction in `gdn_scan_r` (`TB == 0`):**
  Adding an `if (sg == 0)` branch after the per-token `threadgroup_barrier` so that `sg = 1..3` only sum `red[slot][0..3][lane].x` while `sg == 0` sums `float2` (`s.x, s.y`) regressed `gdn_scanr_4` by `+1.0%` to `+1.4%` (`49.45 ms -> 50.03 ms` at `T = 512`, `98.50 ms -> 99.50 ms` at `T = 1024`). On Apple M5, uniform `float2` shared-memory loads (`ld.tg.v2.f32`) across all 4 simdgroups are faster than divergent simdgroup control flow inside the tight `T`-token recurrence loop.

### 2.3 Paired A/B Timings (`gdn_prompt_parity --ab`, 36 GDN layers, `f32 = 0`, median of 3 alternating pairs)

| `T` | Stage | Kernel | Baseline (ms) | Candidate (ms) | Ratio (`cand/base`) | Bitwise Parity (`memcmp == 0`) |
|---:|---|---|---:|---:|---:|---|
| `512` | Stage 0 | `gdn_pre` | `11.540` | `11.558` | `1.0028` | `PASS` (`qn, kn, vc, gb, kq, history`) |
| `512` | Stage 1 | `gdn_scanr_4` | `49.424` | `49.187` | **`0.9944`** (`-0.56%`) | `PASS` (`o, state`) |
| `512` | Stage 2 | `gdn_post` | `7.191` | `7.203` | `1.0016` | `PASS` (`yh, y`) |
| `1024` | Stage 0 | `gdn_pre` | `22.976` | `22.978` | `1.0001` | `PASS` (`qn, kn, vc, gb, kq, history`) |
| `1024` | Stage 1 | `gdn_scanr_4` | `98.183` | `97.115` | **`0.9904`** (`-0.96%`, `-1.07 ms`) | `PASS` (`o, state`) |
| `1024` | Stage 2 | `gdn_post` | `16.149` | `16.165` | `1.0003` | `PASS` (`yh, y`) |
| `2048` | Stage 0 | `gdn_pre` | `46.081` | `45.920` | `0.9961` | `PASS` (`qn, kn, vc, gb, kq, history`) |
| `2048` | Stage 1 | `gdn_scanr_4` | `195.760` | `193.583` | **`0.9882`** (`-1.18%`, **`-2.18 ms`**) | `PASS` (`o, state`) |
| `2048` | Stage 2 | `gdn_post` | `33.942` | `33.916` | `0.9992` | `PASS` (`yh, y`) |

---

## 3. Target 2: `gr_norm_h_moec`, `gr_apply_moec`, and `gr_norm_h4` (`metal/kernels/prompt.metal` lines 170–440)

### 3.1 What Worked (Kept)
1. **SIMD-Compacted 32-bit Row Offsets in `moe_comb_slots` / `moe_comb`:**
   Previously, `moe_comb` looped over all `k = 0..topk-1`, checking `if (rk[k] >= 0)` and computing `64`-bit `ulong(rk[k]) * n + d` `100` times per thread (`10` `d`-iterations $	imes$ `10` `topk`).
   In `moe_comb_slots`, simdgroup 0 now compacts the valid routed slots (`r >= 0`) in order using `simd_prefix_exclusive_sum(ushort(valid))` and precomputes 32-bit row offsets `roff[prefix] = uint(r) * n` (`max_rows * 2560 <= 52.4M < 2^32`) and weights `wk[prefix] = w`, storing the valid count `cnt = simd_sum(ushort(valid))` in `roff[16]`. `moe_comb` then runs a branchless loop `for (uint k = 0; k < cnt; ++k) s = fma(wk[k], YE[roff[k] + d], s);` with 32-bit indexing, preserving the exact summation order (`100%` `memcmp == 0`).
2. **Hoisted `!a.apply` Fast Path + `n == 2560` Full Loop Unrolling in `gr_norm_h4`, `gr_norm_h_moec`, and `gr_norm_h_stream_out`:**
   - Separating `if (!a.apply)` in `gr_norm_h4` eliminates zero-initializing `bv[10]`, evaluating `exp(-inj)` per stream, and branching on `if (a.apply)` inside the inner loop, speeding up `gr_norm_h4 (apply 0)` by **`-2.4%` to `-3.5%`**.
   - Specializing `if (n == 2560)` with `#pragma clang loop unroll(full)` on the `i = 0..9` (`d = tid + i * 256`) stream loops in `gr_norm_h4`, `gr_norm_h_moec`, and `gr_norm_h_stream_out` (plus `#pragma clang loop unroll(full)` on the 8-way `part[i]` reductions) keeps `bv[0..9]` and `rv[0..9]` in registers rather than dynamically indexed thread stack memory.

### 3.2 What Was Ruled Out (Discarded)
- **Swapping the `gr_norm_h_moec` `moe_comb` Loop Order (`k` outer, `i = 0..9` inner) or Full Unrolling of the `moe_comb` Outer Loop:**
  Accumulating `bv[0..9]` across `k = 0..cnt-1` in registers (or `#pragma clang loop unroll(full)` on the `i = 0..9` loop calling `moe_comb`) keeps all 10 accumulators live simultaneously across the `YE` global loads, increasing register pressure and hurting occupancy (`+1.5%` to `+3.1%` slower at `T = 1024, 2048`). Keeping the outer `for (uint d = tid, i = 0; d < n; d += 256, ++i)` loop over `moe_comb` un-unrolled while unrolling the stream loops (`rc`, `w_norm`, `xnh`) gave the best performance at all `T`.

### 3.3 Paired A/B Timings (`comb_fold_parity T grvec`, 4 reps per sample, 5 alternating samples)

| `T` | Kernel | Candidate Med (ms) | Baseline Med (ms) | Median Ratio (`cand/base`) | Bitwise Parity (`memcmp == 0`) |
|---:|---|---:|---:|---:|---|
| `512` | `gr_norm_h_moec (apply 1)` | `3.804` | `3.731` | **`0.9908`** (`-0.92%`) | `PASS` |
| `512` | `gr_apply_moec (apply 1)` | `3.045` | `3.036` | `1.0030` | `PASS` |
| `512` | `gr_norm_h4 (apply 1)` | `1.763` | `1.772` | **`0.9953`** (`-0.47%`) | `PASS` |
| `512` | `gr_norm_h4 (apply 0)` | `0.801` | `0.831` | **`0.9652`** (`-3.48%`) | `PASS` |
| `1024` | `gr_norm_h_moec (apply 1)` | `6.756` | `6.792` | **`0.9949`** (`-0.51%`) | `PASS` |
| `1024` | `gr_apply_moec (apply 1)` | `6.061` | `6.066` | **`0.9994`** (`-0.06%`) | `PASS` |
| `1024` | `gr_norm_h4 (apply 1)` | `3.508` | `3.514` | **`0.9983`** (`-0.17%`) | `PASS` |
| `1024` | `gr_norm_h4 (apply 0)` | `1.775` | `1.811` | **`0.9784`** (`-2.16%`) | `PASS` |
| `2048` | `gr_norm_h_moec (apply 1)` | `14.190` | `14.554` | **`0.9982`** (`-0.18%`, mean `-2.5%`) | `PASS` |
| `2048` | `gr_apply_moec (apply 1)` | `12.160` | `12.178` | **`0.9995`** (`-0.05%`) | `PASS` |
| `2048` | `gr_norm_h4 (apply 1)` | `7.002` | `7.006` | **`0.9977`** (`-0.23%`) | `PASS` |
| `2048` | `gr_norm_h4 (apply 0)` | `3.722` | `3.814` | **`0.9758`** (`-2.42%`) | `PASS` |