[oniisanlabs]

Summary

Optimizes rec_mlx prefill and resident decode on Apple M5 Pro (applegpu_g17s) while preserving 100% bit-exact (0 ULP / mx.array_equal) parity across prefill logits, all 50 decode layer traces, all 150,144 output logits, and both greedy (resident8) and filtered (top_p/top_k) sampling:

1. Dynamic M5 Pro Worker Scaling Across 16/18/20-Core GPU Bins (rec_mlx/kernels/device.py):

   • Queries AGXAccelerator → gpu-core-count via read-only ctypes IOKit (<0.05 ms, fallback to 18 cores) and sizes workers = head_workers = max(32, min(40, (cores * 2) // 4 * 4)) while keeping mlp_workers = 32:
     • 20-core M5 Pro (40 / 32 / 40): Reaches groups = 40 // 4 = 10, activating the static lock-free QKV path in kv_teams.py instead of the atomic ticket loop (shared_projection), and uses all 40 threadgroups for attention and the 150,144-row LM head.
     • 18-core M5 Pro (36 / 32 / 36): Matches 36-worker capacity and uses 36 head workers.
     • 16-core M5 Pro (32 / 32 / 32): Prevents sync_all deadlock on 16-core M5 Pro bins where 36 threadgroups exceed the 16 × 2 = 32 co-resident threadgroup limit.

2. Selective Prefill Fastpath for applegpu_g17s (rec_mlx/m4_max_prefill.py, rec_mlx/kernels/__init__.py):

   • Enables m4_max.FullModel (GlobalCache "causal" SDPA fastpath + compact power-of-two initial KVStorage block allocation) and PrefillAttention (512-token windowed sliding attention) on M5 Pro, while keeping dequantize_linear=False and fuse_switch_glu=False on applegpu_g17s where MLX's native Metal 4 quantized_matmul / gather_qmm is 35–40% faster than BF16 dequantization and flint_prefill_gate_up.
   • Delivers +40.8% faster prefill at L = 4,096 (1,841.8 → 2,594.1 tok/s) and 12× lower initial KV allocation (1.65 GiB → 138 MiB for 4K sessions).

3. Decode Shader & Fast Single-Bank KVView::address (rec_mlx/kernels/m4_max.py, rec_mlx/kernels/masked_sampling.py):

   • Enables kernels/m4_max.py (dot16_w197_bias23 IEEE-754 magic-float unpacking + vectorized 64-bit KVView::load4) and filtered_sampling.make_sampler on M5 Pro.
   • Adds a single-bank fastpath (if(!global || (i < 65536u*128u)) return p0 + offset + i;) and shift/mask (i >> 23, i & (65536u*128u - 1u)) to KVView::address in kernels/m4_max.py, eliminating integer division/modulo/ternary bank resolution per load4 on all sliding-window layers and global layers up to 64K context.

---

Benchmarks (Apple M5 Pro, applegpu_g17s, 20 GPU Cores, 48 GB Unified Memory)

Prefill Throughput (49 Layers, Bit-Exact mx.array_equal)

Configuration	L = 1,024	L = 2,048	L = 4,096	Bit-Exact Parity
main (Baseline M5 Pro)	2,518.6 tok/s	2,090.4 tok/s	1,841.8 tok/s	Baseline
Full m4_max_prefill (with PrefillLinear + PrefillSwitchGLU)	1,763.8 tok/s (0.70×)	1,819.3 tok/s (0.87×)	1,802.5 tok/s (0.98×)	Pass
This PR (GlobalCache + PrefillAttention, native applegpu_g17s QMM)	2,756.2 tok/s (+9.4%)	2,679.3 tok/s (+28.2%)	2,594.1 tok/s (+40.8%)	Pass (0 ULP)

Interleaved A/B Decode Throughput (A -> B -> C -> D / D -> C -> B -> A × 4 reps)

Context Length	Configuration	Greedy (resident8, temp=0)	Default (temp=1.0, top_p=0.95, top_k=20)	Bit-Exact Parity
CTX = 512	main (36/32/32)	136.5 tok/s (1.00×)	131.9 tok/s (1.00×)	Baseline
	This PR (40/32/40 + m4_max + fast KVView)	139.6–141.4 tok/s (+3.6%)	137.2 tok/s (+4.0%)	Pass (0 ULP)
CTX = 8,192	main (36/32/32)	116.9 tok/s (1.00×)	115.1 tok/s (1.00×)	Baseline
	This PR (40/32/40 + m4_max + fast KVView)	123.4 tok/s (+5.6%)	116.8 tok/s (+1.5%)	Pass (0 ULP)
CTX = 32,768	main (36/32/32)	80.2 tok/s (1.00×)	76.5 tok/s (1.00×)	Baseline
	This PR (40/32/40 + m4_max + fast KVView)	84.7 tok/s (+5.6%)	82.7 tok/s (+8.1%)	Pass (0 ULP)