round 12 track c: gemm_f16 2d l2/slc tile swizzling & gemm_hc_up_reg tile + epilogue optimization #

• patch: /tmp/argon_r12_gemm/argon_r12_gemm.patch (git -C vendor/strata apply --check exits 0; touches metal/kernels/prompt.metal and metal/src/engine/window.cpp only)
• scratch checkout & standalone benchmarks: /tmp/argon_r12_gemm/strata (build-metal), /tmp/argon_r12_gemm/bench_r12.cpp, /tmp/argon_r12_gemm/bench_swz2.cpp, /tmp/argon_r12_gemm/bench_swz_overhead.cpp, /tmp/argon_r12_gemm/bench_epi_ab.cpp, /tmp/argon_r12_gemm/window.o

1. executive summary & decision #

1. target 1 (gemm_f16 2D L2/SLC tile swizzling in metal/kernels/dense_mm.metal): REJECTED (0 changes to dense_mm.metal).

   • All tested 2D threadgroup swizzle mappings (GROUP_M = 2, 4, 8, 16, GROUP_N = 2, 4, 8, and pure grid transpose my = tg.x, nx = tg.y) are 100% bitwise identical (memcmp == 0) across all 12 production (N, K) shapes and all M = 256, 512, 1024, 1983, 2048.
   • However, on Apple M5 (64_128_4, 128_64_4, 32_128_2 selected in Round 11), the hardware threadgroup walker already schedules 2D (tg.x, tg.y) grids across the 10 GPU cores with strong L2/SLC locality. Any software swizzle either ties within < 0.2% noise (GN4, GN8, GM2 at M = 2048) or severely degrades M5's hardware L2/SLC prefetching on 64_128_4 at M = 1983, 2048 (e.g., pure transpose slows the 12-shape chunk from 978.2 ms to 1,966.1 ms, GM16 slows it to 1,302.6 ms, GM8 slows it to 1,013.5 ms).
   • Even when compiled as zero-argument compile-time template constants (bench_swz_overhead), no swizzle variant delivers a repeatable speedup across M = 256..2048. Following the falsifier rule ("Keep ONLY what delivers a clean, repeatable speedup across the production shapes"), dense_mm.metal is left untouched.

2. target 2 (gemm_hc_up_reg tile shape & epilogue optimization in metal/kernels/prompt.metal + metal/src/engine/window.cpp): SHIPPED (+6.9% to +7.8% faster, saving 7.3–7.9 ms/chunk at T = 1983..2048).

   • Epilogue Unroll & Pointer Hoisting (prompt.metal): Added #pragma clang loop unroll(full) to the outer cooperative-tensor loop (i += 4) and inner 4-stream reduction loop (c = 0..3), hoisted device const half* xnh_t = xnh + ulong(t) * 4 * n + d; and const ulong out_idx = ulong(t) * n + d;, and replaced s = s / 4.0f with s *= 0.25f. Isolated paired A/B testing (11 interleaved rounds) shows the epilogue optimization alone yields a consistent 1.008x–1.019x (~1.0–1.9%) speedup across every tile (64_64_4, 64_128_4, 32_128_4) and every T, while remaining 100% bitwise identical (memcmp == 0) on both mixh (fp16) and mixed (fp32).
   • Tile Shape Sweep (32_128_4 across all T on M5 k4): Sweeping 10 register tile instantiations (64_64_4, 64_128_4, 128_64_4, 64_32_4, 32_128_2, 32_128_4, 64_128_2, 32_256_4, 64_256_4, 32_64_2, 64_64_2, 16_128_4, 16_128_2, 32_64_4, 16_64_2) revealed that gemm_hc_up_reg_32_128_4 (TM = 32, TN = 128, SG = 4, 0 bytes threadgroup memory) is the fastest tile across every single T from 1 to 2048:
     • Why 32_128_4 wins: Halving TM from 64 to 32 while keeping TN = 128, SG = 4 cuts each thread's cooperative-tensor register footprint in half (acc.get_capacity() drops from 64 floats to 32 floats per lane = 8 outputs x 4 streams). Lower register pressure boosts M5 occupancy while keeping 4 simdgroups cooperating across the K = 320 reduction and 32 coalesced output columns per threadgroup. By contrast, SG = 2 variants (32_128_2, 64_128_2) increase per-lane accumulator capacity or halve K-reduction parallelism.
     • Furthermore, at T = 1983, TM = 64 pads to 31 * 64 = 1984 (1 wasted row), whereas at common non-multiple-of-64 lengths TM = 32 also reduces tail padding.
   • Combined (gemm_hc_up_reg_32_128_4 + unrolled/hoisted epilogue), the 96-call per-chunk time drops from 112.39 ms -> 105.10 ms (1.069x, -7.29 ms/chunk) at T = 1983 and 116.52 ms -> 108.76 ms (1.071x, -7.76 ms/chunk) at T = 2048.

2. target 1 evidence: gemm_f16 2d l2/slc tile swizzling #

2.1 tested swizzle mappings (bench_swz_overhead.cpp & bench_swz2.cpp) #

All variants were compiled as zero-overhead template specializations (orig, GM4, GM8, GM16, GN4, GN8, Transp) and verified against orig (memcmp == 0 on 0xFF-poisoned output buffer Y across all 12 production shapes at M = 256, 512, 1024, 1983, 2048, fails = 0).

2.2 paired 12-shape chunk totals on quiet M5 (7 interleaved rounds per shape) #

2.3 per-shape breakdown at M = 1983 (us/call, median of 7 interleaved rounds) #

Analysis: On M5, gemm_f16_64_128_4 reads 64 rows of A (320 KB at K = 2560) and 128 rows of B (640 KB at K = 2560) per threadgroup. When Metal dispatches (tg.x, tg.y) in its native order, simultaneously active threadgroups across the 10 GPU cores share the 320 KB mA slice in L2 while streaming consecutive mB slices sequentially. Remapping tg into M-major strips (GM >= 8 or Transp) forces concurrent threadgroups to read disjoint mA rows (31 * 320 KB = 9.9 MB) while hammering the same mB slice, thrashing L2 and causing up to a 2.4x slowdown on Wqkv/Wout/Wq.

3. target 2 evidence: gemm_hc_up_reg tile sweep & epilogue optimization #

3.1 full tile + swizzle sweep (bench_r12.cpp & bench_hcup2.cpp, 96 calls/chunk, unbound mixed) #

All candidates verified memcmp == 0 on both mixh (fp16) and mixed (fp32, when bound) across T = 1, 4, 7, 8, 16, 24, 32, 48, 64, 128, 256, 384, 512, 768, 1024, 1536, 1983, 2048 (0 failures).

3.2 isolated paired A/B: baseline epilogue vs unrolled/hoisted epilogue (bench_epi_ab.cpp, 11 interleaved rounds, 96 calls/chunk) #

4. verification & parity results #

1. dense_mm_parity:
   • /tmp/argon_r12_gemm/strata/build-metal/dense_mm_parity: ALL PASS (3.05s, all 13 quant formats x 6 shapes, memcmp(wh)=0, memcmp(xh)=0).
   • /tmp/argon_r12_gemm/strata/build-metal/dense_mm_parity --tiles M (M = 256, 512, 1024) and --ab M (M = 1983, 2048): ALL BITWISE (< 3s per invocation).
2. comb_fold_parity:
   • /tmp/argon_r12_gemm/strata/build-metal/comb_fold_parity 64: ALL PASS (4.67s, T = 1, 7, 64).
   • /tmp/argon_r12_gemm/strata/build-metal/comb_fold_parity T hcup confirm (T = 64, 512, 1024, 1983, 2048): ALL PASS (< 2.1s per invocation, tgmem = 0 on all reg kernels).
3. Engine compilation & patch check:
   • window.cpp compiled to /tmp/argon_r12_gemm/window.o using strata_metal_kernels release flags: exit 0 (0 warnings/errors).
   • git -C vendor/strata apply --check /tmp/argon_r12_gemm/argon_r12_gemm.patch: exit 0.