Accelerate large-model MLX prefill

[justinchuby]

Summary

• split large decoder fusion clusters into configurable layer groups to release prefill intermediates before they exceed unified memory
• keep initializer bytes in Plan-owned storage and drop constant fused-node runtime inputs
• materialize shape-keyed GQA output once before downstream projections
• add an opt-in stock-MLX fast Metal INT4 QMM path and document the profiling/decode lessons

ONNXRUNTIME_EP_MLX_LAYER_PARTITIONS=8 selects eight transformer layers per fused partition. ONNXRUNTIME_EP_MLX_BF16_QMM_FP16=1 enables the standalone stock-MLX fallback kernel. The best result below uses the private MLX kernel from ml-explore/mlx#4216 via MLX_BF16_QMM_FP16=1.

Muse-Glimmer-30B results

M1 Max 32 GB, 512-token prompt, asymmetric INT4 block 32:

Runtime	Prefill	200-token decode
llama.cpp UD-Q4_K_XL	137.84 tok/s	13.50 tok/s
MLX EP + this PR + MLX #4216	138.67 tok/s median	14.79 tok/s

The layer partitioning reduced MLX peak memory from 34.85 GiB to about 20 GiB, eliminating swap. Greedy token IDs were stable across measured runs.

Validation

• Rust library tests: 9 passed, 1 ignored benchmark
• focused GQA + asymmetric MatMulNBits tests: 47 passed
• same-binary repeated 512-token prefill runs
• 512-token prompt followed by 200 generated tokens

[github-actions]

🏎️ MLX EP op benchmark

Median session.run (MLX EP), base vs PR — 2 regression(s) 🔴 · 1 improvement(s) 🟢 (threshold ±10%).

Case	base ms	PR ms	Δ%	note
gqa_prefill_s256_h16kv4d64	1.708	2.026	+18.6%	🔴 regression
diffusion_resnet_block_320ch_64x64	12.868	15.204	+18.2%	🔴 regression
gelu_fp32_512x1024	0.846	0.883	+4.5%
add_broadcast_fp16_512x1024	0.652	0.670	+2.7%
diffusion_upblock_resize_concat_320ch_32x32	17.008	17.141	+0.8%
matmul_fp16_256x512x512	0.792	0.791	-0.0%
gqa_decode_s1_past256_h16kv4d64	1.293	1.291	-0.1%
matmulnbits_prefill_M256_K1024_N1024	1.351	1.335	-1.2%
layernorm_fp32_512x1024	0.873	0.857	-1.8%
simplified_ln_default_domain_512x1024	0.676	0.659	-2.5%
matmul_fp32_256x512x512	0.904	0.879	-2.8%
rmsnorm_512x1024	0.836	0.809	-3.1%
matmulnbits_decode_M1_K2048_N2048	0.661	0.638	-3.4%
softmax_fp32_1x16x256x256	1.288	1.226	-4.9%
matmul_fp16_gemv_1x2048x2048	1.677	1.579	-5.8%
conv2d_fp32_32to64_64x64	1.020	0.867	-15.0%	🟢 improvement

_{base base (main) · PR PR · 60 iters (median) · ORT 1.27.0 · lower ms is better. Timings on a shared runner are noisy; only |Δ%| > 10% is flagged.}

[justinchuby]

A self-contained preview wheel is now available: https://github.com/justinchuby/onnxruntime-mlx/releases/tag/v0.27.3-muse-prefill-preview.1

It bundles the pinned optimized MLX + mlx-c runtime, so users do not need to wait for MLX #4216 or replace their system MLX. Enable with MLX_BF16_QMM_FP16=1 and ONNXRUNTIME_EP_MLX_LAYER_PARTITIONS=8. The wheel was installed in a clean venv and successfully registered with ORT.