perf(rec_mlx): optimize M5 Pro prefill (+40.8%), decode kernels, and GPU core scaling

engine/src/rec_mlx/kernels/__init__.py

@@ -29,7 +29,7 @@ def build_decoder(model, *, optimized):
 
         install_qat(megakernel, model, profile.mlp_workers, profile.specialize_mixed_experts)
     engine_type = megakernel.FullModel
-    if profile.chip == "M4 Max":
+    if profile.chip in ("M4 Max", "M5 Pro"):
         from rec_mlx.m4_max import FullModel as engine_type
 
     engine = engine_type(
@@ -43,10 +43,11 @@ def build_decoder(model, *, optimized):
     if not optimized:
         return engine, None
 
-    if profile.chip == "M4 Max":
+    if profile.chip in ("M4 Max", "M5 Pro"):
         from rec_mlx.m4_max_prefill import optimize_prefill
 
-        optimize_prefill(model)
+        is_m4_max = profile.chip == "M4 Max"
+        optimize_prefill(model, dequantize_linear=is_m4_max, fuse_switch_glu=is_m4_max)
 
     from rec_mlx.kernels.activation_optimization import guard
     from rec_mlx.kernels.activation_optimization import transform as fast_activation
@@ -57,7 +58,7 @@ def build_decoder(model, *, optimized):
     spec = captures["decode"]
     guard(megakernel, engine, spec)
     spec = dict(spec, header=fast_activation(spec["header"]))
-    if profile.chip == "M4 Max":
+    if profile.chip in ("M4 Max", "M5 Pro"):
         from rec_mlx.kernels.m4_max import transform as m4_max
 
         spec = dict(spec, header=m4_max(spec["header"]))

engine/src/rec_mlx/kernels/device.py

@@ -1,4 +1,7 @@
+import ctypes
+import ctypes.util
 from dataclasses import dataclass
+from functools import lru_cache
 
 
 @dataclass(frozen=True)
@@ -12,6 +15,42 @@ class DeviceProfile:
     specialize_mixed_experts: bool = True
 
 
+@lru_cache(maxsize=1)
+def _gpu_core_count(default=18):
+    try:
+        iokit = ctypes.cdll.LoadLibrary(ctypes.util.find_library("IOKit"))
+        cf = ctypes.cdll.LoadLibrary(ctypes.util.find_library("CoreFoundation"))
+        iokit.IOServiceMatching.restype = ctypes.c_void_p
+        iokit.IOServiceGetMatchingService.argtypes = [ctypes.c_uint, ctypes.c_void_p]
+        iokit.IORegistryEntryCreateCFProperty.argtypes = [
+            ctypes.c_uint,
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_uint32,
+        ]
+        iokit.IORegistryEntryCreateCFProperty.restype = ctypes.c_void_p
+        cf.CFStringCreateWithCString.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_uint32]
+        cf.CFStringCreateWithCString.restype = ctypes.c_void_p
+        cf.CFNumberGetValue.argtypes = [ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p]
+        cf.CFRelease.argtypes = [ctypes.c_void_p]
+
+        service = iokit.IOServiceGetMatchingService(0, iokit.IOServiceMatching(b"AGXAccelerator"))
+        if not service:
+            return default
+        key = cf.CFStringCreateWithCString(None, b"gpu-core-count", 0x08000100)
+        prop = iokit.IORegistryEntryCreateCFProperty(service, key, None, 0)
+        cf.CFRelease(key)
+        iokit.IOObjectRelease(service)
+        if not prop:
+            return default
+        value = ctypes.c_int64(0)
+        ok = cf.CFNumberGetValue(prop, 4, ctypes.byref(value))
+        cf.CFRelease(prop)
+        return int(value.value) if ok and value.value > 0 else default
+    except Exception:
+        return default
+
+
 def profile_for_device(info):
     chips = {
         ("applegpu_g16s", "Apple M4 Pro"): "M4 Pro",
@@ -32,10 +71,21 @@ def profile_for_device(info):
             head_workers=64,
             prefill_chunk_size=4096,
         )
+    if chips[identity] == "M5 Pro":
+        cores = info.get("gpu_cores") or _gpu_core_count()
+        workers = max(32, min(40, (cores * 2) // 4 * 4))
+        return DeviceProfile(
+            chip=chips[identity],
+            architecture=architecture,
+            workers=workers,
+            mlp_workers=32,
+            head_workers=workers,
+            prefill_chunk_size=4096,
+        )
     return DeviceProfile(
         chip=chips[identity],
         architecture=architecture,
-        prefill_chunk_size=4096 if chips[identity] == "M5 Pro" else 2048,
+        prefill_chunk_size=2048,
         # The specialized M4 resident shader fails scalar/packet parity.
         specialize_mixed_experts=chips[identity] not in ("M4 Pro", "M4 Max"),
     )

engine/src/rec_mlx/kernels/m4_max.py

@@ -12,6 +12,21 @@ def transform(header):
         scalar
         + "\n    vec<bfloat16_t,4> load4(uint i) const {return *((const device vec<bfloat16_t,4>*)(address(i)));}",
     )
+    old_addr = """    const device bfloat16_t* address(uint i) const {
+        uint block=global?(i/128)>>16:0;
+        uint local=global?i%(65536u*128u):i;
+        const device bfloat16_t* p=block==0?p0:block==1?p1:block==2?p2:p3;
+        return p+offset+local;
+    }"""
+    new_addr = """    const device bfloat16_t* address(uint i) const {
+        if(!global || (i < 65536u*128u)) return p0 + offset + i;
+        uint block = i >> 23;
+        uint local = i & (65536u*128u - 1u);
+        const device bfloat16_t* p = block==1 ? p1 : block==2 ? p2 : p3;
+        return p + offset + local;
+    }"""
+    assert header.count(old_addr) == 1
+    header = header.replace(old_addr, new_addr)
     original = function(header, "attention_partial_path_w53")
     replacement = original.replace(
         "float key[4],value[4];for(uint j=0;j<4;++j){key[j]=float(key_base[at+j]);value[j]=float(value_base[at+j]);}",

engine/src/rec_mlx/kernels/masked_sampling.py

@@ -8,7 +8,7 @@ class FilteredResidentSampler:
     def __init__(self, module, engine, specification):
         self.module, self.engine = module, engine
         make_sampler = native_sampler
-        if engine.device_profile.chip == "M4 Max":
+        if engine.device_profile.chip in ("M4 Max", "M5 Pro"):
             from .filtered_sampling import make_sampler
         # After EOS, remaining graph nodes must not write KV or consume positions.
         inactive = f"""

engine/src/rec_mlx/m4_max_prefill.py

@@ -104,16 +104,18 @@ def window_attention(q, k, v, *, scale, mask, sinks):
     return mx.concatenate(outputs, axis=-2)
 
 
-def optimize_prefill(model):
+def optimize_prefill(model, *, dequantize_linear=True, fuse_switch_glu=True):
     dense_layers = set(model.config["mlp_only_layers"])
     for index, layer in enumerate(model.native.layers):
-        for name in ("q_proj", "o_proj"):
-            setattr(layer.self_attn, name, PrefillLinear(getattr(layer.self_attn, name)))
+        if dequantize_linear:
+            for name in ("q_proj", "o_proj"):
+                setattr(layer.self_attn, name, PrefillLinear(getattr(layer.self_attn, name)))
         if model.config["layer_types"][index] == "sliding_attention":
             layer.self_attn = PrefillAttention(layer.self_attn)
         dense = index in dense_layers
         mlp = layer.mlp if dense else layer.mlp.shared_expert
-        for name in ("gate_proj", "up_proj", "down_proj") if dense else ("down_proj",):
-            setattr(mlp, name, PrefillLinear(getattr(mlp, name)))
-        if not dense:
+        if dequantize_linear:
+            for name in ("gate_proj", "up_proj", "down_proj") if dense else ("down_proj",):
+                setattr(mlp, name, PrefillLinear(getattr(mlp, name)))
+        if fuse_switch_glu and not dense:
             layer.mlp.switch_mlp = PrefillSwitchGLU(layer.mlp.switch_mlp)

engine/tests/test_kernel_routing.py

@@ -171,6 +171,21 @@ def check_reduction(self, workers, bits, specialized, ranked):
             )
             self.assertTrue(mx.array_equal(trace[2048:], expected).item(), seed)
 
+    def test_m5_pro_profile_scales_with_gpu_cores(self):
+        for cores, expected_workers in ((16, 32), (18, 36), (20, 40)):
+            profile = profile_for_device(
+                {
+                    "architecture": "applegpu_g17s",
+                    "device_name": "Apple M5 Pro",
+                    "gpu_cores": cores,
+                }
+            )
+            self.assertEqual(profile.chip, "M5 Pro")
+            self.assertEqual(profile.workers, expected_workers)
+            self.assertEqual(profile.mlp_workers, 32)
+            self.assertEqual(profile.head_workers, expected_workers)
+            self.assertEqual(profile.prefill_chunk_size, 4096)
+
 
 if __name__ == "__main__":
     unittest.main()