From 14d56d3d81b08a55ede213da023a51c1e85a7a8c Mon Sep 17 00:00:00 2001 From: oniisanlabs <1436535+oniisanlabs@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:16:37 +0100 Subject: [PATCH] perf(rec_mlx): optimize M5 Pro prefill, decode kernels, and GPU core scaling --- engine/src/rec_mlx/kernels/__init__.py | 9 ++-- engine/src/rec_mlx/kernels/device.py | 52 ++++++++++++++++++- engine/src/rec_mlx/kernels/m4_max.py | 15 ++++++ engine/src/rec_mlx/kernels/masked_sampling.py | 2 +- engine/src/rec_mlx/m4_max_prefill.py | 14 ++--- engine/tests/test_kernel_routing.py | 15 ++++++ 6 files changed, 95 insertions(+), 12 deletions(-) diff --git a/engine/src/rec_mlx/kernels/__init__.py b/engine/src/rec_mlx/kernels/__init__.py index 1812bb5..68a7325 100644 --- a/engine/src/rec_mlx/kernels/__init__.py +++ b/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"])) diff --git a/engine/src/rec_mlx/kernels/device.py b/engine/src/rec_mlx/kernels/device.py index 0cd8358..63cd5d6 100644 --- a/engine/src/rec_mlx/kernels/device.py +++ b/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"), ) diff --git a/engine/src/rec_mlx/kernels/m4_max.py b/engine/src/rec_mlx/kernels/m4_max.py index a333ca2..2d6d826 100644 --- a/engine/src/rec_mlx/kernels/m4_max.py +++ b/engine/src/rec_mlx/kernels/m4_max.py @@ -12,6 +12,21 @@ def transform(header): scalar + "\n vec load4(uint i) const {return *((const device vec*)(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]);}", diff --git a/engine/src/rec_mlx/kernels/masked_sampling.py b/engine/src/rec_mlx/kernels/masked_sampling.py index daa9370..f7f2ec8 100644 --- a/engine/src/rec_mlx/kernels/masked_sampling.py +++ b/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""" diff --git a/engine/src/rec_mlx/m4_max_prefill.py b/engine/src/rec_mlx/m4_max_prefill.py index ddfafb9..d7f35ba 100644 --- a/engine/src/rec_mlx/m4_max_prefill.py +++ b/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) diff --git a/engine/tests/test_kernel_routing.py b/engine/tests/test_kernel_routing.py index 884394e..eadb7cc 100644 --- a/engine/tests/test_kernel_routing.py +++ b/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()