# Round 23 Track C: `moe_mm.metal` Prompt Prefill `mm_gu3` & `mm_down3` Optimization (`argon-moe-mm-prompt-opt`) ## 1. Summary of Shipped Changes - **Deliverable patch:** `/tmp/argon_r23_moemm/argon_r23_moemm.patch` (`git -C vendor/strata apply --check` verified clean, exit `0`). - **Touched files:** `metal/kernels/moe_mm.metal` (lines 507–574, 688–765) and `metal/tests/moe_mm_parity.cpp` (lines 306–336). ### Key Optimizations in `metal/kernels/moe_mm.metal` 1. **4-bit Register Sign-Parity Helper `sign4_iq2xs(uint j, int hh)` (`W4` & `W8` for `T_IQ2_XS`, `T_IQ2_XXS`, `T_IQ3_XXS`)**: - `w4_signed(g, s, sc)` only reads the low 4 bits (`s & 1, s & 2, s & 4, s & 8`) of its second argument. - When `hh == 0` (50% of SIMD lanes), the low 4 bits of `ksigns_iq2xs[j]` are simply `j` itself (`v >> 9`) — bit 7 (`popcount(j) & 1u`) is not needed at all. - When `hh == 1`, `(ksigns_iq2xs[j] >> 4) & 15` is `(j >> 4) | ((popcount(j) & 1u) << 3)`. - `sign4_iq2xs(j, hh) = hh ? ((j >> 4) | ((popcount(j) & 1u) << 3)) : j` eliminates divergent threadgroup byte loads `ksg(tab, ..., j)` in `W4`, `W8`, `W4`, `W8`, `W4`, and `W8`, while avoiding `popcount` on `hh == 0` lanes. - Note on `moe_mm_parity --ab` sign convention: `moe_mm_parity.cpp:424` reports `diff = (mb - ma) / ma * 100.0` where `A` is the candidate and `B` is `--ab`, so positive `+X%` means `A` is faster than `B` (`ma < mb`). Direct head-to-head A/B comparison of `A = ksg` vs `B = sign4_iq2xs` (with all other changes identical) confirmed `B (sign4_iq2xs)` is `3.47%–10.63%` faster than `A (ksg)`. 2. **Constant-Offset Pointer Folding (`(k & 255) >> 3` and `(k & 255) >> 2`)**: - In `W8`, `4 * ib + l` (with `ib = (k & 255) >> 5`, `l = (k & 31) >> 3`, `l < 2`) is identical to `(k & 255) >> 3`. Computing `device const ushort* q = (device const ushort*)(bl + 2) + ((k & 255) >> 3)` turns `q[l], q[l + 2]` into constant-offset loads `q[0], q[2]` and eliminates `l`. - In `W8`, `8 * ib + 2 * l + hh` is `(k & 255) >> 2`, `4 * ib + l` is `(k & 255) >> 3`, and `2 * l + hh` is `(k & 31) >> 2`. Folding these into `qs = bl + 2 + ((k & 255) >> 2)` (`qs[0], qs[4]`) and `sp = bl + 74 + ((k & 255) >> 3)` (`sp[0], sp[2]`) eliminates `l` in the hot path and reduces VGPR pressure across the 4-way unrolled `mm_gu3` loop. - Applied the same constant-offset pointer folding to `W8`, `W8`, `W4`, `W4`, `W4`, and `W4`. 3. **Bitwise-Identical `fma` Scale Math under `MTLMathModeSafe`**: - Because `0.25f` ($2^{-2}$) and `0.5f` ($2^{-1}$) are exact powers of two and `rdh(bl)` is an fp16 value promoted to fp32 (exponent $\ge -24$), `d * fma(float(nib), 0.25f, 0.125f)` is 100% bitwise identical to `d * (0.5f + float(nib)) * 0.25f` across all 65,536 fp16 bit patterns and all 16 nibble values `0..15`, while replacing `FADD + FMUL + FMUL` (which `MTLMathModeSafe` cannot reassociate) with `FMA + FMUL` and removing `rdh(bl)` from the first stage of the dependency chain. 4. **Ruled-Out Candidates (Not Shipped)**: - Custom `W8` in `mm_down3`: tested both shared-`d` and split-`d` `W8` specializations; both regressed `mm_down3` by `0.7%–3.0%` at `T=1024, 2048` because CSEing the block pointer across the two `W4` halves allows LLVM to hoist all 16 loads across the 8-way unrolled loop, increasing VGPR pressure near the occupancy boundary. Kept the generic two-`W4` path untouched. - `W8` scalar byte packing: regressed `Q6_K` by `2.4%–4.5%` vs `uint4(as_type(...))`. Reverted so `W8` is untouched. --- ## 2. Paired A/B Benchmark Results on M5 (`moe_mm_parity --v3 --tile 16,32,32,4 --dtile 64,32,2 --gu-bench`) Note: In `moe_mm_parity --ab`, `A` is the candidate kernel, `B` is the baseline reference kernel, and `diff = (B - A) / A * 100%` (`+` = candidate `A` is faster than baseline `B`). | Format | `T` | Candidate `A` Med (Min) us | Baseline `B` Med (Min) us | Speedup `diff` Med (Min) | |---|---:|---:|---:|---:| | `IQ2_XS` | `512` | `7,682.3` (`7,669.9`) | `7,982.3` (`7,978.3`) | **`+3.91%` (`+4.02%`)** | | `IQ2_XS` | `1024` | `16,603.6` (`12,901.3`) | `17,684.0` (`13,439.3`) | **`+6.51%` (`+4.17%`)** | | `IQ2_XS` | `2048` | `25,535.6` (`25,486.3`) | `26,712.9` (`26,492.7`) | **`+4.61%` (`+3.95%`)** | | `IQ3_S` | `512` | `17,946.8` (`17,777.0`) | `19,823.0` (`19,711.9`) | **`+10.45%` (`+10.88%`)** | | `IQ3_S` | `1024` | `28,841.0` (`28,506.7`) | `31,010.2` (`30,957.8`) | **`+7.52%` (`+8.60%`)** | | `IQ3_S` | `2048` | `48,501.1` (`46,976.4`) | `52,371.5` (`51,207.6`) | **`+7.98%` (`+9.01%`)** | | `IQ2_S` | `512` | `35,342.2` (`34,052.6`) | `36,953.1` (`35,455.6`) | **`+4.56%` (`+4.12%`)** | | `IQ2_S` | `1024` | `44,448.1` (`42,382.8`) | `45,269.8` (`44,114.7`) | **`+1.85%` (`+4.09%`)** | | `IQ2_XXS` | `1024` | `32,583.5` (`30,793.0`) | `33,326.4` (`32,109.3`) | **`+2.28%` (`+4.27%`)** | | `IQ2_XXS` | `2048` | `56,247.3` (`54,307.5`) | `58,809.2` (`56,789.6`) | **`+4.55%` (`+4.57%`)** | | `IQ3_XXS` | `512` | `28,919.5` (`22,779.1`) | `33,020.1` (`24,348.8`) | **`+14.18%` (`+6.89%`)** | | `IQ3_XXS` | `1024` | `33,955.7` (`29,351.1`) | `35,193.0` (`32,111.7`) | **`+3.64%` (`+9.41%`)** | | `IQ3_XXS` | `2048` | `50,338.7` (`45,779.3`) | `53,857.8` (`50,270.8`) | **`+6.99%` (`+9.81%`)** | --- ## 3. Bitwise Parity Verification (`HS`, `YE`, `Y` `memcmp == 0`) Verified 100% bitwise identity (`HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS`) between candidate and unmodified baseline across all formats (`IQ2_XS`, `IQ3_S`, `Q6_K`, `Q2_0`, `IQ2_S`, `IQ2_XXS`, `IQ3_XXS`) at `T = 256, 512, 1024, 2048` on both `TK=32` (`W8`) and `TK=16` (`W4`): ``` [W8 TK=32] IQ2_XS T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS [W8 TK=32] IQ3_S T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS [W8 TK=32] Q6_K T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS [W8 TK=32] Q2_0 T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS [W8 TK=32] IQ2_S T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS [W8 TK=32] IQ2_XXS T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS [W8 TK=32] IQ3_XXS T=256..2048: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS [W4 TK=16] IQ2_XS..IQ3_XXS T=256: HS_memcmp=PASS YE_memcmp=PASS Y_memcmp=PASS ```