Round 4 Track C: Prefill MoE Plan (mm_clear + mm_hist + mm_scan + mm_place) & Epilogue / Gather Audit #

• Patch: /tmp/argon_prefill_plan/argon_prefill_plan.patch (verified clean with git -C vendor/strata apply --check)
• Target Device: Apple M5 (10-core GPU, 24 GB Unified Memory, Metal 4 matmul2d K4 right-input layout)

1. Executive Summary #

We audited and benchmarked the MoE prefill plan pipeline (mm_clear, mm_hist, mm_scan, mm_place, mm_gather) in metal/kernels/moe_mm.metal and all elementwise epilogue / conversion kernels (swiglu_h, act_f16, moe_out, gr_epi_down, gr_norm_h, mm_combine) in Window::prompt_chunk (metal/src/engine/window.cpp & prompt_chain.inc).

Key Findings & Optimizations Implemented #

1. 100% Elimination of mm_clear via Exact Padding-Row Initialization in mm_scan:

   • mm_hist initializes threadgroup atomic_uint h[1024] to 0 and unconditionally writes every element bcnt[tg.x * a.n_expert + e] (e < a.n_expert), so bcnt never needed external clearing.
   • mm_clear was dispatching (p.maxrows + 255) / 256 threadgroups (36,864 ints for C = 2048) on every layer solely to initialize p.stok to -1 for the padded tail rows [r0 + tot, r0 + rows) of each expert that mm_place does not touch.
   • In mm_scan (1 threadgroup of n_expert = 512 threads, e = tid), each thread e already computes r0, tot, and rows = nt * a.tm. Having thread e write srt_tok[r0 + r] = -1 for r = tot .. rows - 1 (0..15 ints per expert) completely eliminates the mm_clear dispatch across all 48 MoE layers.

2. Discovery & Fix of the Real Plan Bottleneck: $O(1024^2)$ Serial Threadgroup Scan in mm_place (5.3x–6.0x Plan Speedup):

   • Profiling individual plan kernels revealed that mm_hist (1.1–1.5 us) and mm_scan (2.9–4.8 us) were tiny, whereas mm_place alone took 50.4–81.5 us/layer (91% of total plan time)!
   • Root cause: for (uint j = 0; j < tid; ++j) rank += s_id[j] == e; in mm_place forced threads in the last warp (tid = 992..1023) to execute up to 1,023 serial threadgroup memory loads and comparisons per block.
   • Fix: Replaced the $O(1024)$ serial loop in mm_place with a two-level deterministic 32-lane sub-block nibble-packed counter (p_cnt[4096]) + short4 intra-warp scan (ltid = tid & 31, at most 7 short4 loads per thread). Because topk = 10 per token guarantees at most $\lceil 32/10 ceil = 4 \le 15$ occurrences of any expert within a 32-entry sub-block, 8 sub-block counts pack into a single 32-bit atomic_uint without overflow, and the resulting rank, pos[i], srt_tok[row], and srt_w[row] are 100% bitwise identical to the original loop.
   • Result: mm_place drops from 50.4–81.5 us to 3.9–6.2 us (12.9x faster), and the full MoE plan (mm_hist + mm_scan + mm_place) drops from 46.9–77.2 us/layer to 8.6–13.1 us/layer (5.4x–5.9x faster) in moe_mm_parity --bench.

3. Why Single-Threadgroup mm_hist + mm_scan Fusion Was Rejected:

   • We also prototyped and measured a single-threadgroup fused mm_hist_scan kernel (1 TG of 1024 threads looping over b = 0 .. n_blocks - 1). Because 1 TG cannot occupy all 10 GPU cores on M5 across n_blocks = 10..20, mm_hist_scan took 8.27–11.59 us vs 5.60 us for separate multi-TG mm_hist + 1-TG mm_scan_pad. Keeping mm_hist multi-TG and fusing padding-clear into mm_scan is strictly faster.

4. Halving mm_gather Memory Traffic via mm_gather_h (p.mixh FP16 Source + half4 Loads/Stores):

   • In Window::prompt_chunk (line 3332), gemm_hc_up_64_64_4 already outputs both FP32 p.mixed and FP16 p.mixh (mixh[i] = half(mixed[i])).
   • Previously, mm_gather read FP32 p.mixed (4 bytes/elem) and converted it to half again (half(x[...])) for every routed token row (24,576 x 2560 = 251.6 MB read + 125.8 MB write per layer at T = 2048).
   • Added mm_gather_h, which reads directly from FP16 p.mixh using vectorized 64-bit half4 loads and stores (125.8 MB read + 125.8 MB write, bit-for-bit identical to mm_gather since mixh is already half(mixed)).
   • Result: mm_gather time in moe_mm_parity --bench at T = 1024 drops from 762.2 us to 495.2 us (-267.0 us/layer, 1.54x faster) and at T = 2048 drops from 2028.3 us to 1190.5 us (-837.8 us/layer, 1.70x faster, saving 40.2 ms across 48 layers).

5. Audit of swiglu_h, act_f16, and Other Elementwise Kernels + Barrier Placement:

   • Audited all barriers in Window::prompt_chunk:
     • Removed the redundant mm_clear dispatch and fixed the barrier flag before mm_hist: previously mm_clear had nb_ = true and mm_hist ALSO had nb_ = true, which only worked when stream_ && pred_on && l + 1 < 48 inserted a barrier inside the read-ahead loop. Now nb_ = stream_ && pred_on && l + 1 < 48; before mm_hist, guaranteeing a barrier after route_topk(p.ids) when read-ahead is off or on layer 47, while overlapping mm_hist with the read-ahead route_topk(predicted_ids) when read-ahead is active.
     • Verified swiglu_h (line 3567): the three shared-expert GEMMs (Wshg -> p.gsh, Wshu -> p.ush, Wsg -> p.sgv) run concurrently in one barrier group (nb_ = true on lines 3562, 3564); swiglu_h must wait for p.gsh and p.ush, and gemm(p.hh, p.Wshd, p.ys) must wait for p.hh.
     • Benchmarked swiglu_h (9.6–68.8 us) vs float4-vectorized swiglu_h_v4 (17.2–75.0 us) vs fused gemm_swiglu_f16 (113.6–1294.6 us), and act_f16 vs act_f16_v4. On M5, scalar coalesced 256-thread dispatches saturate the memory controller and hardware coalescer slightly better than float4-strided threads (64.8 us vs 75.0 us for swiglu_h at T=2048; 546.1 us vs 558.5 us for act_f16), and fusing Wshg + Wshu into one double-GEMM kernel hurts occupancy (1042 us -> 1295 us at T=2048). Thus swiglu_h and act_f16 are already optimal as-is.

2. Benchmark Results on Apple M5 (moe_mm_parity --bench --v3 --gu IQ2_XS --dn IQ4_NL) #

2.1 Microbenchmark: Isolated Plan Kernels (48-layer run_graph median per layer) #

2.2 End-to-End moe_mm_parity --bench (512 Experts, IQ2_XS Gate/Up 16,32,32,4, IQ4_NL Down 16,64,32,2) #

3. Bitwise & Numerical Parity Verification #

• Bitwise Plan & Gather Parity (plan_epi_bench):
  • Tested across T = 128, 512, 1024, 2048 with stok and bcnt pre-polluted with 0xDEADBEEF / 777777 and with skipped experts (ids = -1 from MoCE tail skipping):
  • std::memcmp across meta[0..10], boff[0..nblk*NE-1], eoff[0..NE-1], tiles[0..nt-1], stok[0..nr-1], sw[0..nr-1], pos[0..nent-1], and XS[0..nr*N-1]: 100% bitwise identical (0 byte differences).
• moe_mm_parity (--v3 --gu IQ2_XS --dn IQ4_NL --tile 16,32,32,4 --dtile 64,32,2):
  • T = 128: plan ok; y vs fp16-operand reference rel L2 2.02e-05, vs exact float reference 4.93e-04 PASS
  • T = 256: plan ok; y vs fp16-operand reference rel L2 1.80e-05, vs exact float reference 4.90e-04 PASS (identical 1.80e-05 / 4.90e-04 to baseline)
• prompt_parity: ALL PASS (route_topk, HC read as GEMMs, swiglu_h, moe_out)
• dense_mm_parity: ALL PASS across all 13 quantization formats (F32, BF16, Q8_0, Q4_0, Q5_0, IQ4_NL, IQ4_XS, Q2_0, IQ3_S, Q3_K, Q4_K, Q5_K, Q6_K).