gather_qmm_rhs_nax: fixed BM=64 discards most routed-expert arithmetic at short per-expert runs (MoE prefill) #3925

Summary

On the Neural-Accelerator path, gather_qmm_rhs_nax uses a fixed bm = 64 (quantized.cpp, marked // TODO: Tune the block sizes). The kernel walks the sorted indices inside a BM-row tile and re-runs its entire K loop once per distinct expert, storing only that expert's row slice — so useful arithmetic is roughly min(1, rows_per_expert / BM), and rows per expert is a caller-side quantity:

rows_per_expert = tokens_per_step * top_k / num_experts

For MoE prefill this is routinely far below 64. A 256-expert / top-8 model prefilling 512-token chunks (a realistic step size on memory-constrained devices) presents runs of 16, i.e. ~75% of the routed-expert arithmetic is computed and discarded.

Measurements

Base M5 (10 GPU cores), macOS 26, mlx built from 7c92ce12. Shape is a real model's MoE FFN: E=256, K=2048, N=512, 4-bit affine g64, sorted indices, chained timing (up-proj then down-proj so the chain is type-correct), % of the measured 15.4 TFLOPS fp16 NA ceiling on this machine.

Stock dispatch (BM=64), sweeping rows per expert:

min(1, R/BM) predicts the relative curve to within 0.02 everywhere.

Alignment control (isolates the mechanism): runs of exactly 64 starting on a 64-row tile boundary vs the same runs shifted by 32 — identical arithmetic, bytes, and run length, only tile phase differs:

Skewed run lengths (real routers are not balanced) are worse than uniform at the same total: 16.9% vs 20.8% at R≈16.

Fix, measured

Instantiating BM=32 (same WM=2/WN=2) and dispatching on M/E — same binary, choice forced either way via a local env hook:

Weight traffic is unchanged by BM (each expert's slab is read once either way), and at short runs the weight stream itself becomes the floor — which is why BM=32 lands near the arithmetic-law prediction rather than at long-run saturation.

BM=16 is not worth dispatching: the only valid 16-row config (WM=1/WN=2; WN=4 produces the TM=1/TN=1 tile from #3924 that silently no-ops) measured parity with BM=32 at R≤8 and slightly worse at R=16 — the 64-thread threadgroup halves load parallelism.

End-to-end: Qwen3.6-35B-A3B (4-bit), 32k-token prompt, prefill_step=512 (R=16), warm, same binary back-to-back: TTFT 62.8s → 58.9s (1.066x), decode tok/s unchanged (B=1 decode never takes this path — expected null).

Notes

• The non-NA gather_qmm_rhs uses bm=16 and is correspondingly more forgiving of short runs; only the NA path pins 64.
• All data is from one machine (base M5). The NA path only exists on this generation, so I can't offer a second-chip datapoint yet.
• Each extra BM instantiation adds ~108 kernels to the metallib across the type/group/bits fan-out — if that's a concern, BM=32-only (as measured above) is the minimal version.

I have a branch with the BM=32 instantiation and an M/E-keyed dispatch ready, plus the sweep as a benchmark script — happy to open a PR if the approach looks right, or adjust thresholds if you'd rather tune it differently.