# Round 15 Track C Audit: `qsa_prep` Threadgroup-Memory & Barrier Reduction + `route_topk` `n_exp == 512` Unrolling

- **Agent:** `argon-qsa-prep-and-route-topk`
- **Scratch workspace:** `/tmp/argon_r15_prep/strata`
- **Patch artifact:** `/tmp/argon_r15_prep/argon_r15_prep.patch` (verified clean with `git -C vendor/strata apply --check`)
- **Target files & line boundaries:**
  - `metal/kernels/qsa.metal` lines 58–187 (`tg_sum256`, `norm_rope`, `qsa_prep`)
  - `metal/kernels/prompt.metal` lines 16–115 (`route_topk`)

---

## 1. Summary of Changes

### Target 1: `qsa_prep` / `norm_rope` Threadgroup Memory (`1,088 B -> 320 B`) & Barrier Reduction (`metal/kernels/qsa.metal`)
1. **Shrank `threadgroup float row[256]` (`1,024 B`) to `threadgroup float row[QROT]` (`64` floats = `256 B`):**
   - `norm_rope` only ever reads `row[p]` (`0..31`) and `row[p + 32]` (`32..63`) when `rope && tid < QROT` (`QROT = 64`).
   - Replaced `row[tid] = y;` with `if (rope) { if (tid < QROT) row[tid] = y; threadgroup_barrier(mem_flags::mem_threadgroup); ... }`.
   - Total static threadgroup memory of `qsa_prep` drops from **`1,088 B` to `320 B` (`-768 B`, `-70.6%`)** (`part[8]` = 32 B, `row[64]` = 256 B, `amax[8]` = 32 B), eliminating 192 wasted shared-memory stores per threadgroup when `rope == true` and all 256 shared-memory stores + the RoPE barrier when `rope == false` (`cell == 0` `dead` key).
2. **Eliminated Redundant Trailing Barriers in `tg_sum256` and `norm_rope` (`4 -> 2` barriers per single-call threadgroup):**
   - Removed the trailing `threadgroup_barrier(mem_flags::mem_threadgroup)` from `tg_sum256` (safe because `part[0..7]` is distinct from `row[0..63]` and `amax[0..7]`).
   - Added `bool reuse = false` to `norm_rope` so the trailing `threadgroup_barrier(mem_flags::mem_threadgroup)` is only executed when `reuse == true` (the `r == 32 && t_self == 0` loop over `t < n_pre` where `part` and `row` can be reused across `cell == 0` and `(cell & 3) == 3` or subsequent `t` iterations).
   - Every single-call threadgroup (`r < 32` and `r == 32, t_self >= 3` — i.e., `67,583` of `67,584` threadgroups at `T = 2048`) now executes **2 `threadgroup_barrier`s instead of 4 (`-50%`)**.
3. **Evaluated Per-Token Early-Exit for `r >= 28 && r < 32` (`qidx` heads, line 180):**
   - Inspected `qsa_prompt.metal` (`qsa_score_mm` lines 190–225) and `metal/tests/qsa_parity.cpp`: when `pos0 + n_tok > 2051` (e.g., a chunk crossing `QBUDGET = 2051` such as `pos0 = 1984, T = 96`), `qsa_score_mm` executes a `matmul2d` over all `16` tokens in the tile (`tQ` sliced at `t0 * IH`) and scores the tail block `dead` for every token in the tile before `qsa_topk` or `qsa_flash` runs. Leaving `qidx[t]` uninitialized (or NaN-poisoned) for tokens `t` with `pos0 + t + 1 <= 2051` when `pos0 + n_tok > 2051` would feed uninitialized memory into `qsa_score_mm` and alter `scores[t]`.
   - Kept `if (a.pos0 + a.n_tok <= uint(QBUDGET)) return;` as-is so all stages and parity tests remain 100% bitwise identical.

### Target 2: `route_topk` `n_exp == 512` Unrolled Fast Path (`metal/kernels/prompt.metal`)
- Added an `if (n_exp == 512)` fast path with `#pragma clang loop unroll(full)` for:
  1. Initial logit load + per-lane max reduction (`v[i] = l[lane + 32 * i]; mx = max(mx, v[i]);`),
  2. Softmax `exp(v[i] - mx)` + per-lane `sum` reduction,
  3. Normalization `v[i] *= rcp` + `isnan(v[i])` guard, and
  4. The MoCE `rank == p_prot && moce_scale > 1.0f` scaling loop.
- Eliminates all `idx < n_exp` (`lane + 32 * i < 512`) branch predicates on the `K = 512` path while preserving the exact floating-point operation sequence and fallback `else` path for `n_exp != 512` (`K = 256, 320`).

---

## 2. Verification & Bitwise Parity Results

All standalone parity suites and a dedicated A/B bitwise (`memcmp == 0`) & paired-dispatch benchmark (`ab_bench`) were run on Apple M5:
- `qsa_parity`: **ALL PASS** (`q rel L2 5.22e-08`, `indexer q 5.60e-08`, `pooled keys 4.36e-08`, `top-k 0/6153 diff`, `attention + gate 2.44e-07`)
- `qsa_prompt_parity`: **ALL PASS** (`qsa_prep(QSA_Q_F16_FC) vs qsa_prep+qsa_q_f16: qcurh memcmp PASS, qsa_flash memcmp PASS`)
- `prompt_parity`: **ALL PASS** (`route_topk 300 tokens: id mismatches 0, worst weight rel err 2.28e-07`)
- `comb_fold_parity`: **ALL PASS** (`T = 1, 7, 64, 1024` bitwise PASS)
- `ab_bench` (`memcmp == 0` between baseline and patched kernels across all outputs `qcur`, `qidx`, `kq`, `vq`, `ksc`, `vsc`, `tail`, `dead`, `pooled`, `ids`, `w`): **100% BITWISE PASS (`fails = 0`)**.

---

## 3. Paired A/B Benchmark Results (Apple M5)

### `qsa_prep` (12-layer back-to-back command buffers, 9 alternating pairs, `static_tg_mem`: `1,088 B -> 320 B`)

| Config | `pos0` | `T` | `memcmp` | Baseline x12 Med (Min) ms | Candidate x12 Med (Min) ms | Per-Call Before -> After | Delta (Med) |
|---|---:|---:|:---:|---:|---:|---:|---:|
| `qf16=1, mrope=0` (flash prompt path) | `0` | `512` | `PASS` | `2.210 (2.205)` | `2.185 (2.180)` | `184.2 -> 182.1 us` | **-1.11%** |
| `qf16=1, mrope=0` (flash prompt path) | `0` | `1024` | `PASS` | `4.484 (4.477)` | `4.466 (4.462)` | `373.7 -> 372.1 us` | **-0.41%** |
| `qf16=1, mrope=0` (flash prompt path) | `2050` | `512` | `PASS` | `2.373 (2.367)` | `2.336 (2.327)` | `197.8 -> 194.6 us` | **-1.58%** |
| `qf16=1, mrope=0` (flash prompt path) | `2050` | `1024` | `PASS` | `4.834 (4.816)` | `4.785 (4.770)` | `402.8 -> 398.7 us` | **-1.01%** |
| `qf16=1, mrope=0` (flash prompt path) | `4096` | `512` | `PASS` | `2.372 (2.367)` | `2.328 (2.323)` | `197.7 -> 194.0 us` | **-1.86%** |
| `qf16=1, mrope=0` (flash prompt path) | `4096` | `1024` | `PASS` | `4.836 (4.821)` | `4.785 (4.757)` | `403.0 -> 398.7 us` | **-1.06%** |
| `qf16=1, mrope=0` (flash prompt path) | `4096` | `1984` | `PASS` | `9.484 (9.461)` | `9.502 (9.467)` | `790.4 -> 791.8 us` | +0.18% |
| `qf16=0, mrope=0` (fp32 path) | `2050` | `512` | `PASS` | `2.896 (2.862)` | `2.878 (2.852)` | `241.3 -> 239.8 us` | **-0.62%** |
| `qf16=0, mrope=0` (fp32 path) | `4096` | `512` | `PASS` | `2.898 (2.860)` | `2.880 (2.856)` | `241.5 -> 240.0 us` | **-0.60%** |
| `qf16=0, mrope=0` (fp32 path) | `4096` | `1024` | `PASS` | `5.976 (5.902)` | `5.921 (5.858)` | `498.0 -> 493.4 us` | **-0.92%** |

### `route_topk` (48-layer back-to-back command buffers, 9 alternating pairs, `n_exp = 512`)

| `n_exp` | `prompt_moce` | `T` | `ids` / `w` `memcmp` | Baseline x48 Med (Min) ms | Candidate x48 Med (Min) ms | Per-Call Before -> After | Delta (Med) |
|---:|---:|---:|:---:|---:|---:|---:|---:|
| `512` | `0.00` | `512` | `PASS / PASS` | `0.688 (0.687)` | `0.683 (0.675)` | `14.34 -> 14.22 us` | **-0.82%** |
| `512` | `0.00` | `1024` | `PASS / PASS` | `1.237 (1.229)` | `1.220 (1.217)` | `25.77 -> 25.42 us` | **-1.34%** |
| `512` | `0.00` | `1984` | `PASS / PASS` | `2.293 (2.283)` | `2.292 (2.289)` | `47.78 -> 47.74 us` | **-0.08%** |
| `512` | `320335.00` | `512` | `PASS / PASS` | `0.780 (0.779)` | `0.767 (0.760)` | `16.26 -> 15.98 us` | **-1.69%** |
| `512` | `320335.00` | `1024` | `PASS / PASS` | `1.406 (1.398)` | `1.375 (1.369)` | `29.29 -> 28.65 us` | **-2.19%** |
| `512` | `320335.00` | `1984` | `PASS / PASS` | `2.665 (2.662)` | `2.635 (2.628)` | `55.53 -> 54.89 us` | **-1.15%** |
| `512` | `-32.02` | `512` | `PASS / PASS` | `0.705 (0.702)` | `0.692 (0.681)` | `14.69 -> 14.42 us` | **-1.85%** |
| `512` | `-32.02` | `1024` | `PASS / PASS` | `1.254 (1.243)` | `1.233 (1.226)` | `26.13 -> 25.68 us` | **-1.72%** |
| `512` | `-32.02` | `1984` | `PASS / PASS` | `2.341 (2.328)` | `2.322 (2.315)` | `48.78 -> 48.38 us` | **-0.80%** |