Overview of three perf PRs for LLM inference (#3918, #3919, #3920)

[spokvulcan]

We run local LLM serving on MLX and profiled decode and prefill end to end on an M3 Max. Three small, independent PRs came out of it, and together they cover the three phases that dominate local serving. Filing one issue as a map so they are easy to evaluate as a set.

PR	Where	Change	Measured (M3 Max, Qwen3.6-35B-A3B)
#3918	gather_qmm_rhs (Metal)	Pick the tile size from rows per expert	+12 to 22% on the kernel, +6.2% MoE prefill at 32K
#3919	SDPA ops fallback (Metal)	Fuse the causal mask into the softmax	mask+softmax stage -45% at 32K, +2.8% prefill end to end
#3920	eval_impl and gather_mm (core)	Cheaper eval walk, cached identity indices	+2.0 to 2.8% decode, +3.6% prefill at 8K

#3918 replaces the "TODO: Tune the block sizes" in gather_qmm_rhs with a measured rule. Below 32 rows per expert the stock 16x32 tile stays, since wider tiles straddle expert runs and lose. From 32 up, a 32x64 tile wins outright and brings the kernel to 96% of an equivalent dense qmm. A few lines of dispatch logic plus the ahead-of-time instantiations for the second geometry.

#3919 removes pure waste from long-context causal prefill. The ops fallback builds a bool mask, applies it with a where, then immediately softmaxes the result, which at 32K materializes 512 MB of masked scores per layer chunk. The fused kernel applies the causal condition at the softmax load sites instead. The output is bit-identical to the chain it replaces, checked bitwise on 14 of 14 configurations, and the gate is deliberately narrow, so everything outside it keeps the stock chain.

#3920 is CPU-side overhead. The eval walk did several unordered_map operations per graph edge on every eval, and gather_mm rebuilt its identity row indices a few hundred times per decoded token. A small open-addressing table and a bounded per-shape cache remove both, with the tape order unchanged.

All three are verified token-identical against their baselines over interleaved A/B generation runs, on top of the per-PR checks above. They are independent and can land or be rejected separately. #3918 is the smallest review, #3919 the largest. Happy to split, rebase, or re-measure anything if it helps.