Round 10 Track C Audit: mm_gather_h 128-Bit uint4 Vectorization + mm_scan/mm_place Plan Optimization #

• Agent: argon-mm-gather-and-plan (opencode-go/gemini-4-argon-contributor)
• Scratch workspace: /tmp/argon_r10_gather/strata
• Unified patch: /tmp/argon_r10_gather/argon_r10_gather.patch (verified clean with git -C vendor/strata apply --check)
• Files touched in patch: metal/kernels/moe_mm.metal, metal/tests/moe_mm_parity.cpp (metal/src/engine/window.cpp line 3540 g.tg[0] = 256 kept unchanged; compatible with prompt_chain.inc and prompt_chain.cpp too).

1. mm_gather_h Investigation (half4 vs uint4, Threadgroup Sizing, Unrolling & Multi-Row-per-TG) #

1.1 Roofline & Geometry Analysis on Apple M5 (24 GB Unified Memory, 153.6 GB/s Peak DRAM Bandwidth) #

For N_EMBD = 2560 (5,120 bytes per row = 640 half4s = 320 uint4s):

• At T = 512 (512 experts, TM = 16): tot_rows = 8,432 -> XS write footprint is 43.17 MB per layer (2.07 GB across 48 layers) plus gathered reads from xh (2.62 MB).
• At T = 1024: tot_rows = 14,352 -> XS write footprint is 73.48 MB per layer (3.53 GB across 48 layers) plus gathered reads from xh (5.24 MB).
• At T = 2048: tot_rows = 24,576 -> XS write footprint is 125.83 MB per layer (6.04 GB across 48 layers) plus scattered reads from xh (10.49 MB). Because 125.83 MB exceeds the M5 System Level Cache (SLC), every layer's XS write and subsequent read stream through DRAM at ~130–153 GB/s (85–99% of hardware DRAM peak).

1.2 Paired 48-Layer Sweep Results (mm_gather_h Variants) #

Measured via paired alternating 48-dispatch command buffers (7–9 pairs after GPU warmup) on Apple M5:

1.3 Key Findings for mm_gather_h #

1. 128-bit uint4 copies beat 64-bit half4 copies across all threadgroup sizes (32, 64, 128, 256) at T=512 (-0.16 ms / -1.2%) and T=1024 (-0.50 ms / -2.1%), and tie/edge out baseline at T=2048 (55.508 ms vs 55.543 ms), because 128-bit loads/stores halve the number of loop iterations (320 uint4s vs 640 half4s per row).
2. Why tg[0] = 256 with uint4 beats tg[0] = 64 at T = 2048 (55.51 ms vs 56.06 ms):
   • With uint4 and tg[0] = 256, iteration 0 issues 256 * 16B = 4,096B (32 cache lines of 128B) of the 5,120B row concurrently across 8 simdgroups, and iteration 1 issues the remaining 64 * 16B = 1,024B (8 cache lines) on the first 2 simdgroups (tid < 64). Because 64 is an exact multiple of the 32-wide SIMD width, simdgroups 2..7 exit cleanly at the SIMD level without intra-warp divergence, while launching 8 simdgroups per row maximizes memory-level parallelism (MLP) when xh (10.49 MB at T=2048) spills out of L1/L2 into SLC/DRAM.
   • Making mm_gather_h take ushort nth [[threads_per_threadgroup]] (with scalar uint tg [[threadgroup_position_in_grid]] per MSL scalar/vector attribute rules) supports any tg[0] = 32, 64, 128, 256 while keeping g.tg[0] = 256 in window.cpp, prompt_chain.inc, and prompt_chain.cpp.

2. Target 2 Investigation: mm_scan and mm_place Plan Kernels #

We also inspected and benchmarked mm_hist, mm_scan, and mm_place in metal/kernels/moe_mm.metal:

1. mm_scan barrier + dead-slot elimination:
   • Previously, mm_scan zeroed s_rows[0..31] and s_tiles[0..31], executed an extra threadgroup_barrier(mem_flags::mem_threadgroup), and then looped for (int i = 0; i < 32; ++i) after the prefix sum barrier because 512 threads only populate 16 of the 32 simdgroup slots.
   • By taking ushort nsg [[simdgroups_per_threadgroup]] and looping for (uint i = 0; i < nsg; ++i), we eliminate both the initial zero-fill and the first threadgroup_barrier (cutting mm_scan's barriers in half from 2 to 1).
2. mm_place threadgroup-memory & barrier reduction (s_id[1024] -> register simd_broadcast):
   • Previously, mm_place allocated threadgroup short s_id[1024] (2 KiB), wrote s_id[tid] = short(e) after zeroing p_cnt, and after the second barrier did a threadgroup-memory short4 scan over the caller's 32-lane sub-block.
   • Because each 32-lane sub-block (sub = tid >> 5) is an exact 32-lane Metal simdgroup, each lane can compute its intra-simdgroup rank directly in registers via simd_broadcast(e, l) while waiting for the p_cnt atomic updates, eliminating threadgroup short s_id[1024] completely.

2.1 Paired 48-Layer Plan Timing Results (mm_hist + mm_scan + mm_place) #

3. Verification (memcmp == 0 & moe_mm_parity) #

• Added fast (1.40s total runtime) --gather-bench mode to metal/tests/moe_mm_parity.cpp (using #ifdef MM_GATHER_ONLY to skip compiling the heavy GEMM template matrix when testing/benchmarking plan + gather):
  • Verifies 100% bitwise identity (memcmp == 0) on XS (across all tg[0] in {32, 64, 128, 256} against mm_gather) and exact plan invariants across T = 1, 7, 64, 256, 512, 1024, 2048.
  • Benchmarks 48-layer command buffers for both plan and mm_gather_h across T = 512, 1024, 2048.
• Full end-to-end moe_mm_parity --v3 --tile 16,32,32,4 --dtile 64,32,2 verified PASS (rel L2 1.11e-05 at T=1, 3.20e-06 at T=7, 2.65e-06 at T=64).