# opus R31 Track A: `mm_gu3` IQ2_S / IQ2_XXS / IQ2_XS / IQ3_S decode arithmetic

Patch: `/tmp/opus_r31_moemm/opus_r31_moemm.patch` (only `metal/kernels/moe_mm.metal`; `git -C vendor/strata apply --check` passes).
Scratch tree: `/tmp/opus_r31_moemm/strata`, build `/tmp/opus_r31_moemm/build`, baseline kernels `/tmp/opus_r31_moemm/base_kernels`.
Exhaustive checker: `/tmp/opus_r31_moemm/chk.{metal,cpp}` (binary `/tmp/opus_r31_moemm/chk`).

## What ships (all integer and pointer rewrites, no FP change)
1. **`W8<T_IQ2_S>::block`**: the lane's `hh = (kl >> 2) & 1` comes from `mm_gu3` (fixed for the whole K loop). The grid word is read as `(tab + hh)[2 * i]`, and the sign byte as `s >> (hh << 2)` instead of `hh ? s >> 4 : s`.
   Measured separately: the shift alone gives about +4.7%, and `tab + hh` adds about +3.5% on top.
2. **`W8<T_IQ2_XXS>`**: there's a new `block(bl, k, d, tab, hh)`, and `mm_gu3` runs IQ2_XXS with the 256-wide super-block loop (cached `dg/du`) the way IQ2_S does, with the lane `hh` and `tab + hh`. The `sign4_iq2xs` select stays. `run()` is a thin wrapper, so the other callers don't change.
3. **`W8<T_IQ2_XS>`**: same change as IQ2_XXS (`block` with cached d, lane hh, `tab + hh`).
4. **`W8<T_IQ3_S>::run`**: `s >> (4 * hh)` instead of the select.
5. `mm_gu3`: `constexpr bool sb = IQ2_S || IQ2_XXS || IQ2_XS` picks the 256 super-block loop. The QUANT_CTAB path is kept inline (`ctab ? grid[i] >> 32*hh : th[2*i]`).

## Results: `moe_mm_parity --v3 --tile 16,32,32,4 --dtile 16,64,32,2 --dn Q2_0 --ab base --gu-bench`
diff = (base - new) / new, so a positive number means new is faster. Medians are from paired, interleaved runs. The machine was shared with the live server, so I give the quiet-run ranges.

| gate/up | T=512 | T=1024 | T=2048 | bitwise |
|---|---|---|---|---|
| IQ2_S   | +6.6..+11.7% | +7.9..+8.9% | **+7.9..+8.8%** (25.1 ms vs 27.1 ms) | PASS |
| IQ2_XXS | +3.0..+5.9% | +4.4..+5.7% | **+5.8..+7.6%** | PASS |
| IQ3_S   | +4.6..+5.4% | +2.4..+2.6% | +2.0..+3.0% | PASS |
| IQ2_XS  | +3.5..+3.8% | +3.5..+5.4% | +6.6..+6.9% | PASS |

All of them hit `HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS`. I also checked 72 extra A/B parity runs: 4 formats × tiles {16,32,32,4 / 16,64,32,2 / 32,64,32,2 / 16,16,32,8 / 16,32,16,4 / 16,64,64,2} × {T=256, T=1000, T=512 with QUANT_CTAB=1}. All passed bitwise.

Estimated full-model kprof savings, extrapolated from the T=2048 standalone numbers (I didn't run a full-model check, per the rules): IQ2_S 763.8 ms × ~7.4% ≈ **-56 ms**, IQ2_XXS 220.1 × ~6% ≈ **-13 ms**, IQ3_S 46.7 × ~2.5% ≈ -1 ms. That's roughly **-70 ms total (about 2% of prefill GPU time)**.

## Exhaustive GPU checks (chk): 65,536 fp16 d × 16 nibbles × v in 1..63 × all 16 sign masks (1.06 G evaluations)
- `(d*0.125f) * float(2n+1)` vs `d * fma(n, 0.25f, 0.125f)`: bit-identical for every d. IQ3_XXS `(d*0.25f)*float(2n+1)` vs `d*fma(n,.5,.25)` is also identical. This is expected: scaling by a power of 2 is exact, the product `v*sc` (≤ 6+5+11 bits) is exact in fp32, so there's a single rounding to fp16 either way.
- These `w4_signed` rewrites are bit-identical to the original for every d, including NaN/Inf: the fp32 sign-XOR (`as_type<uint4>(float4(uchar4)) ^ (s << {31..28})`), and the fp16 0x6400 mantissa trick with sign XOR on the fp16 codes.
- Sign-XOR on the **fp16 output** is identical for every finite d but differs on NaN-d sign bits, so it is not 100% bit-exact.

## Tried and rejected (all bitwise PASS, all slower; measured on IQ2_S T=2048 unless noted)
| variant | result |
|---|---|
| `w4_signed` fp32 sign-XOR | -11.3% (IQ2_XXS -16%) |
| `w4_signed` fp16 0x6400 codes + sign-XOR | -12% (IQ2_XXS -15.5%) |
| `w4_signed` 0x6400 codes, keep select | -11.2% |
| `w4_signed` select on the scale (`s&1 ? -sc : sc`) | -3.4% (IQ2_XXS ±0) |
| scale as `fma(float(n), d/4, d/8)`, d/4 and d/8 cached | -9.5% (d/8 computed at use: -2.6%; with tab+hh: -13%) |
| pre-scaled `dg = d*0.125`, `sc = dg * float(((scb<<1)&30)|1)` | -1.5% vs base (about 9% behind the shipped one) |
| `tab + hh` with the select kept for the sign | -2.5% |
| index pre-doubled `(p0<<1) | ((qh<<(9-2l))&0x600)` | ±0 vs shipped |
| `l` taken from the lane as well | ±0 vs shipped (dropped) |
| IQ2_XXS `sign4_iq2xs` as a shift `(j|parity<<7) >> 4*hh` | -10% (also -8% with cached d + tab+hh) |
| IQ3_S cached d via the 256 loop + lane hh | +1.5% (the shift alone is better: +2.8%; lane hh without caching: same as shift alone) |

Takeaway: on the M5 the decode is ALU-issue bound (each extra op per W4 costs about 1%). `d * fma(float(nibble), 0.25f, 0.125f)` with constant operands, and `float(uchar)` plus the select-negate in `w4_signed`, are already the cheapest forms. The wins came from removing per-lookup integer work (`+ hh`, the `hh ?` select on the sign byte) and per-call scale loads. Selects are not uniformly bad, though: `sign4_iq2xs`'s select beats the shift form.

Not done: IQ3_XXS (not in this model's hot list). The non-coop_k4 (M4) path still calls `W4<F>::run`, which is unchanged; I couldn't run it on the M5.