From f01b58d0da6da5e166b7bf7e90c3efd04e7399ba Mon Sep 17 00:00:00 2001 From: simply-sunny Date: Mon, 21 Sep 2026 14:45:26 -0400 Subject: [PATCH] Add Apple Silicon MPS acceleration for dense Qwen3.5 models Keep the Transformers patch device-gated, validate the pure reference path, and include reproducible full-model and held-out MPS evidence for the dense Decider models. --- README.md | 2 + decider/bench/mps.py | 110 ++++++ decider/engine.py | 9 +- decider/evaluate.py | 12 +- decider/infer.py | 21 +- decider/mps_ops.py | 264 +++++++++++++ decider/vision/model.py | 7 +- docs/benchmarks/mps-full-model.json | 580 ++++++++++++++++++++++++++++ docs/benchmarks/mps-full-model.md | 58 +++ docs/benchmarks/mps-heldout.json | 42 ++ docs/benchmarks/mps-heldout.md | 44 +++ pyproject.toml | 1 + tests/test_mps_ops.py | 123 ++++++ 13 files changed, 1264 insertions(+), 9 deletions(-) create mode 100644 decider/bench/mps.py create mode 100644 decider/mps_ops.py create mode 100644 docs/benchmarks/mps-full-model.json create mode 100644 docs/benchmarks/mps-full-model.md create mode 100644 docs/benchmarks/mps-heldout.json create mode 100644 docs/benchmarks/mps-heldout.md create mode 100644 tests/test_mps_ops.py diff --git a/README.md b/README.md index 3f2e8ed..420f73d 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,8 @@ The v8 weights stay available under the Hub tag `v8`. `docs/HISTORY.md` describe pip install git+https://github.com/Mapika/decider # or: git clone ... && pip install -e ".[serve]" ``` +On Apple Silicon, install the optional MLX/Metal kernel with `pip install -e ".[metal]"` from a clone. Without it, MPS inference uses the PyTorch implementation. + ```python from decider.infer import Decider d = Decider("Mapika/decider-2b") # one CUDA GPU, bf16, about 4 GB; downloads the weights on first use diff --git a/decider/bench/mps.py b/decider/bench/mps.py new file mode 100644 index 0000000..545bde2 --- /dev/null +++ b/decider/bench/mps.py @@ -0,0 +1,110 @@ +"""Offline full-model MPS check. + +Usage: python -m decider.bench.mps MODEL reference|conv|optimized + +Runs one mode per process without concurrent GPU work. JSON includes per-request +probabilities and synchronized timings; model loading is excluded. ``reference`` +is the pure Transformers PyTorch path, ``conv`` adds only Decider's fused +causal convolution, and ``optimized`` adds the MPS attention patch as well. +""" +import inspect +import json +import platform +import statistics +import sys +import time +from pathlib import Path + +import torch +import transformers +from huggingface_hub import snapshot_download + +from decider import mps_ops +from decider.engine import patch_conv +from decider.infer import Decider, Example, Q + + +def configure(mode): + import transformers.models.qwen3_5.modeling_qwen3_5 as mq + + if mode not in ("reference", "conv", "optimized"): + raise ValueError("mode must be reference, conv, or optimized") + # Avoid FLA/Triton in the reference arm: this is the pure Transformers path. + mq.torch_chunk_gated_delta_rule = inspect.unwrap(mq.torch_chunk_gated_delta_rule) + mq.causal_conv1d_fn = inspect.unwrap(mq.causal_conv1d_fn) + if mode == "reference": + mps_ops.patch_mps = lambda: False + return False, False + if mode == "conv": + mps_ops.patch_mps = lambda: False + patch_conv() + return False, True + return mps_ops.patch_mps(), True + + +def main(): + if len(sys.argv) != 3: + raise SystemExit(__doc__) + name, mode = sys.argv[1:] + assert torch.backends.mps.is_available(), "MPS required" + path = snapshot_download(name, local_files_only=True) + patch_result, conv_result = configure(mode) + if name.endswith("vision"): + from PIL import Image + from decider.vision.model import VisionDecisionModel + + model = VisionDecisionModel(path, grad_ckpt=False).to("mps").eval() + cases = [(Image.new("RGB", (224, 224), color), Example( + "Identify the dominant color in the image.", + [Q("What color is shown?", ["red", "green", "blue"])])) + for color in ("red", "green", "blue")] + + def run(case): + logits = model.slot_logits(model.prepare([case])) + return torch.softmax(logits, -1)[0, :3].cpu().tolist() + else: + model = Decider(path, device="mps", use_graphs=False) + cases = ["My card was charged twice for the same purchase.", + "I cannot log in after resetting my password.", + "I would like pricing for fifty licenses."] + + def run(case): + return model.decide(case, [{"question": "Which department should handle this?", + "options": ["billing", "technical", "sales"]}])[0]["probs_list"] + + results = [] + with torch.inference_mode(): + for index, case in enumerate(cases): + for _ in range(2): + run(case) + torch.mps.synchronize() + times = [] + for _ in range(5): + torch.mps.synchronize() + start = time.perf_counter() + probs = run(case) + torch.mps.synchronize() + times.append((time.perf_counter() - start) * 1000) + assert all(torch.isfinite(torch.tensor(probs))) + results.append(dict(case=index, probs=probs, times_ms=times, + median_ms=statistics.median(times))) + weights = model.lm if name.endswith("vision") else model.m.lm + print(json.dumps({ + "model": name, + "snapshot_revision": Path(path).name, + "mode": mode, + "patch_mps": patch_result, + "patch_conv": conv_result, + "torch": torch.__version__, + "transformers": transformers.__version__, + "dtype": str(next(weights.parameters()).dtype), + "hardware": platform.machine(), + "macOS": platform.mac_ver()[0], + "warmups": 2, + "measurements_per_case": 5, + "results": results, + }, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/decider/engine.py b/decider/engine.py index 1d13aec..c7dbdad 100644 --- a/decider/engine.py +++ b/decider/engine.py @@ -68,8 +68,13 @@ class Engine: def __init__(self, path, device="cuda", dtype=torch.bfloat16, use_graphs=True, max_ctx_tokens=1536, compile=True, fp8=False, conv_patch=True): if conv_patch: - patch_conv() + if str(device).startswith("mps"): + from decider.mps_ops import patch_mps + patch_mps() + else: + patch_conv() self.m = DecisionModel(path, dtype=dtype, grad_ckpt=False).to(device).eval() + use_graphs = use_graphs and torch.device(device).type == "cuda" self.tok = self.m.tok; self.dev = device; self.use_graphs = use_graphs; self.max_ctx = max_ctx_tokens self.core, self.W = self.m.lm.model, self.m.lm.lm_head.weight[self.m.letters].detach().clone() self.cfg = dict(compile=compile, fp8=fp8, conv_patch=conv_patch, graphs=use_graphs) @@ -82,7 +87,7 @@ def __init__(self, path, device="cuda", dtype=torch.bfloat16, use_graphs=True, m else: self._fwd_impl = self._fwd_eager self.graphs = {} # (B, T) -> (static_ids, static_out, graph) - self.pool = torch.cuda.graph_pool_handle() if use_graphs else None + self.pool = torch.cuda.graph_pool_handle() if (use_graphs and str(device).startswith("cuda")) else None self.stats = dict(graph_captures=0, forwards=0) def _fwd_eager(self, ids): diff --git a/decider/evaluate.py b/decider/evaluate.py index 322641c..f5e4333 100644 --- a/decider/evaluate.py +++ b/decider/evaluate.py @@ -70,6 +70,7 @@ def agg(keys): ap.add_argument("--temperature", type=float, default=1.0) ap.add_argument("--limit", type=int, default=0) ap.add_argument("--engine", default="eager", help="eager | graph | compile | fp8") + ap.add_argument("--device", default=None, help="cuda | mps | cpu (auto-detected when omitted)") ap.add_argument("--max_options", type=int, default=0, help="0 = sub-sample large label sets to 10 (the original protocol); 255 = offer the full label set") ap.add_argument("--max_ctx", type=int, default=1536) ap.add_argument("--layout", default="state_first", help="state_first | schema_first") @@ -80,16 +81,21 @@ def agg(keys): evals = {k: v for k, v in evals.items() if k in a.tasks.split(",")} if a.limit: evals = {k: v[:a.limit] for k, v in evals.items()} + device = a.device or ("cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu")) + dtype = torch.float16 if device == "mps" else torch.bfloat16 eng = None if a.engine: from decider.engine import Engine - eng = Engine(a.model, compile=a.engine in ("compile", "fp8"), fp8=a.engine == "fp8", conv_patch=a.engine in ("compile", "fp8")) + eng = Engine(a.model, device=device, dtype=dtype, compile=a.engine in ("compile", "fp8"), fp8=a.engine == "fp8", conv_patch=a.engine in ("compile", "fp8")) m = eng.m else: - m = DecisionModel(a.model, grad_ckpt=False).cuda() + if device == "mps": + from decider.mps_ops import patch_mps + patch_mps() + m = DecisionModel(a.model, dtype=dtype, grad_ckpt=False).to(device).eval() os.makedirs(a.out, exist_ok=True) res, dump = run_eval(m, evals, bs=a.bs, temperature=a.temperature, engine=eng, max_options=a.max_options or None, max_ctx=a.max_ctx, layout=a.layout) agg = aggregate(res) print("[agg]", json.dumps(agg, indent=1)) - json.dump(dict(results=res, agg=agg, model=a.model, engine=a.engine), open(f"{a.out}/eval.json", "w"), indent=1) + json.dump(dict(results=res, agg=agg, model=a.model, engine=a.engine, device=device, dtype=str(dtype)), open(f"{a.out}/eval.json", "w"), indent=1) pickle.dump(dump, open(f"{a.out}/preds.pkl", "wb")) diff --git a/decider/infer.py b/decider/infer.py index 379d479..e1c854a 100644 --- a/decider/infer.py +++ b/decider/infer.py @@ -7,6 +7,8 @@ {"question": "How urgent is this?", "options": ["low", "medium", "high"]}]) # -> [{'choice': 'billing', 'confidence': 0.97, 'probs': {...}}, {...}] """ +import logging + import torch from decider.model import DecisionModel, collate from decider.prompt import build, MAX_OPTIONS @@ -23,6 +25,7 @@ class Example: context: str; qs: list; task: str = "infer"; image: bytes = None +logger = logging.getLogger(__name__) NEUTRAL_NONE = "not listed here" @@ -53,9 +56,18 @@ def __call__(self, state, max_state_tokens=32768): class Decider: - """use_graphs=True (default on CUDA) routes scoring through decider.engine.Engine: shape-bucketed - CUDA graphs, ~7x lower single-request latency than eager. Set False for CPU or debugging.""" - def __init__(self, path, device="cuda", dtype=torch.bfloat16, temperature=None, abstain_below=0.0, use_graphs=None): + """One-pass decisions with automatic CUDA, MPS, or CPU device selection. + + CUDA uses shape-bucketed graphs by default. MPS defaults to float16 and uses + the optional MPS patch; CPU defaults to bfloat16. Set ``use_graphs=False`` + for eager execution or debugging. + """ + def __init__(self, path, device=None, dtype=None, temperature=None, abstain_below=0.0, use_graphs=None): + if device is None: + device = "cuda" if torch.cuda.is_available() else ("mps" if torch.backends.mps.is_available() else "cpu") + if dtype is None: + dtype = torch.float16 if str(device).startswith("mps") else torch.bfloat16 + logger.info("Decider device=%s dtype=%s", device, dtype) import json, os cfg = {} try: # model folder may carry decider_config.json (temperature, flags) @@ -73,6 +85,9 @@ def __init__(self, path, device="cuda", dtype=torch.bfloat16, temperature=None, from decider.engine import Engine self.eng = Engine(path, device=device, dtype=dtype); self.m = self.eng.m else: + if str(device).startswith("mps"): + from decider.mps_ops import patch_mps + patch_mps() self.eng = None; self.m = DecisionModel(path, dtype=dtype, grad_ckpt=False).to(device).eval() self.dev = device; self.T = temperature; self.abstain_below = abstain_below self.name = "decider-" + str(cfg.get("version", "dev")) diff --git a/decider/mps_ops.py b/decider/mps_ops.py new file mode 100644 index 0000000..43d4af8 --- /dev/null +++ b/decider/mps_ops.py @@ -0,0 +1,264 @@ +"""MPS kernels for Qwen3.5 gated-delta attention. + +The optimization is measured against Transformers' PyTorch reference fallback for +this kernel; it is not an end-to-end model speedup claim. + +The gated-delta and L2-normalization code follows Transformers 5.17's +``modeling_qwen3_5.py`` under the Apache License 2.0. The MPS inversion and +backend dispatch are Decider additions. +""" +import functools +import inspect +import logging +import warnings + +import torch, torch.nn.functional as F +from decider.engine import fused_causal_conv1d_fn + + +# Adapted from transformers 5.17.0's modeling_qwen3_5.py (Apache-2.0). +# The original license applies to the gated-delta and normalization portions; +# the MPS inversion and dispatch below are Decider additions. +# Native Metal Shading Language (MSL) JIT kernel via MLX/Metal +_metal_invert_kernel = None +_metal_failure_reported = False +_compat_warning_reported = False +try: + import mlx.core as mx + import mlx.core.fast as fast + + _METAL_INVERT_SRC = """ + uint col = thread_position_in_threadgroup.x; + uint mat_idx = threadgroup_position_in_grid.x; + threadgroup float s_inv[4096]; + + for (uint r = 0; r < 64; ++r) { + s_inv[r * 64 + col] = (r == col) ? 1.0f : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + uint mat_offset = mat_idx * 4096; + for (uint r = 1; r < 64; ++r) { + float sum = 0.0f; + if (col < r) { + for (uint k = col; k < r; ++k) { + sum += L[mat_offset + r * 64 + k] * s_inv[k * 64 + col]; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (col < r) { + s_inv[r * 64 + col] = -sum; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint r = 0; r < 64; ++r) { + out_inv[mat_offset + r * 64 + col] = s_inv[r * 64 + col]; + } + """ + _metal_invert_kernel = fast.metal_kernel( + name="invert_unitriangular_64", + input_names=["L"], + output_names=["out_inv"], + source=_METAL_INVERT_SRC, + ) +except (ImportError, OSError, RuntimeError, AttributeError): + logging.getLogger(__name__).debug("MLX Metal kernel unavailable; using PyTorch MPS fallback", exc_info=True) + _metal_invert_kernel = None + + +def l2norm(x: torch.Tensor, dim: int = -1, eps: float = 1e-6) -> torch.Tensor: + return x * torch.rsqrt((x * x).sum(dim=dim, keepdim=True) + eps) + + +def fast_invert_unitriangular_64(L: torch.Tensor) -> torch.Tensor: + """Invert batches of 64x64 unit lower-triangular matrices.""" + global _metal_invert_kernel, _metal_failure_reported + if _metal_invert_kernel is not None and L.is_mps and not L.requires_grad: + assert L.dtype == torch.float32, "Metal inversion requires float32 input" + try: + orig_shape = L.shape + L_flat = L.reshape(-1, 64, 64).contiguous() + num_m = L_flat.shape[0] + torch.mps.synchronize() + L_mx = mx.from_dlpack(torch.to_dlpack(L_flat)) + mx.eval(L_mx) + out_mx = _metal_invert_kernel( + inputs=[L_mx], + grid=(num_m * 64, 1, 1), + threadgroup=(64, 1, 1), + output_shapes=[(num_m, 64, 64)], + output_dtypes=[mx.float32], + )[0] + mx.eval(out_mx) + torch.mps.synchronize() + return torch.from_dlpack(out_mx).to(L.device).reshape(orig_shape) + except Exception: + # Optional acceleration must not turn a recoverable backend issue into + # a failed request; disable the failing kernel for this process. + _metal_invert_kernel = None + if not _metal_failure_reported: + logging.getLogger(__name__).warning("MLX Metal inversion failed; using the PyTorch MPS fallback", exc_info=True) + _metal_failure_reported = True + + N = 64 + inv = torch.eye(N, device=L.device, dtype=L.dtype).expand_as(L).clone() + + # b = 1: 32 blocks of 2x2. For [[1, 0], [l, 1]], inverse is [[1, 0], [-l, 1]] + idx_row = torch.arange(1, 64, 2, device=L.device) + idx_col = torch.arange(0, 64, 2, device=L.device) + inv[..., idx_row, idx_col] = -L[..., idx_row, idx_col] + + # b = 2: 16 blocks of 4x4 + for j in range(16): + r1, r2, r3 = j * 4, j * 4 + 2, j * 4 + 4 + inv[..., r2:r3, r1:r2] = -inv[..., r2:r3, r2:r3] @ L[..., r2:r3, r1:r2] @ inv[..., r1:r2, r1:r2] + + # b = 4: 8 blocks of 8x8 + for j in range(8): + r1, r2, r3 = j * 8, j * 8 + 4, j * 8 + 8 + inv[..., r2:r3, r1:r2] = -inv[..., r2:r3, r2:r3] @ L[..., r2:r3, r1:r2] @ inv[..., r1:r2, r1:r2] + + # b = 8: 4 blocks of 16x16 + for j in range(4): + r1, r2, r3 = j * 16, j * 16 + 8, j * 16 + 16 + inv[..., r2:r3, r1:r2] = -inv[..., r2:r3, r2:r3] @ L[..., r2:r3, r1:r2] @ inv[..., r1:r2, r1:r2] + + # b = 16: 2 blocks of 32x32 + for j in range(2): + r1, r2, r3 = j * 32, j * 32 + 16, j * 32 + 32 + inv[..., r2:r3, r1:r2] = -inv[..., r2:r3, r2:r3] @ L[..., r2:r3, r1:r2] @ inv[..., r1:r2, r1:r2] + + # b = 32: 1 block of 64x64 + inv[..., 32:64, 0:32] = -inv[..., 32:64, 32:64] @ L[..., 32:64, 0:32] @ inv[..., 0:32, 0:32] + return inv + + +def mps_chunk_gated_delta_rule( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + chunk_size: int = 64, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + **kwargs, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Optimized chunk-gated delta rule for Apple Silicon (MPS).""" + assert chunk_size == 64, f"MPS optimized delta rule requires chunk_size=64, got {chunk_size}" + initial_dtype = query.dtype + batch_size, sequence_length, _, k_head_dim = key.shape + num_v_heads, v_head_dim = value.shape[-2:] + recurrent_state_shape = (batch_size, num_v_heads, k_head_dim, v_head_dim) + padded_output_shape = (batch_size, num_v_heads, -1, v_head_dim) + decay = g + + query, key, value, beta, decay = [ + x.transpose(1, 2).to(torch.float32, memory_format=torch.contiguous_format) + for x in (query, key, value, beta, decay) + ] + if use_qk_l2norm_in_kernel: + query = l2norm(query, dim=-1, eps=1e-6) + key = l2norm(key, dim=-1, eps=1e-6) + scaling = query.shape[-1] ** -0.5 + query = query * scaling + + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + query = F.pad(query, (0, 0, 0, pad_size)) + key = F.pad(key, (0, 0, 0, pad_size)) + value = F.pad(value, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + decay = F.pad(decay, (0, pad_size)) + + v_beta = value * beta.unsqueeze(-1) + k_beta = key * beta.unsqueeze(-1) + + query, key, k_beta, v_beta = [ + x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) + for x in (query, key, k_beta, v_beta) + ] + decay = decay.reshape(decay.shape[0], decay.shape[1], -1, chunk_size) + + strictly_upper_mask = torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device).triu(1) + cum_decay = decay.cumsum(dim=3) + pairwise_decay = (cum_decay.unsqueeze(4) - cum_decay.unsqueeze(3)).masked_fill(strictly_upper_mask, float("-inf")).exp() + + ut_system = (k_beta @ key.transpose(-1, -2)) * pairwise_decay + intra_chunk_attn = (query @ key.transpose(-1, -2)) * pairwise_decay + decayed_k_beta = k_beta * cum_decay.exp().unsqueeze(-1) + + # Fast block divide-and-conquer unitriangular inverse on MPS + L = ut_system.tril(-1) + inv = fast_invert_unitriangular_64(L) + new_values = inv @ v_beta + k_cumdecay = inv @ decayed_k_beta + + if initial_state is None: + last_recurrent_state = torch.zeros(recurrent_state_shape, dtype=new_values.dtype, device=new_values.device) + else: + last_recurrent_state = initial_state.to(new_values) + core_attn_out = torch.zeros_like(new_values) + + query = query * cum_decay.exp().unsqueeze(-1) + key = key * (cum_decay[..., -1:] - cum_decay).exp().unsqueeze(-1) + chunk_decay = cum_decay[..., -1].exp()[..., None, None] + + num_chunks = query.shape[2] + qk = torch.cat([query, k_cumdecay], dim=3) + kt = key.transpose(-1, -2) + + for i in range(num_chunks): + qk_state = qk[:, :, i] @ last_recurrent_state + inter_chunk_attn = qk_state[:, :, :chunk_size] + v_new = new_values[:, :, i] - qk_state[:, :, chunk_size:] + core_attn_out[:, :, i] = inter_chunk_attn + intra_chunk_attn[:, :, i] @ v_new + last_recurrent_state = last_recurrent_state * chunk_decay[:, :, i] + kt[:, :, i] @ v_new + + last_recurrent_state = None if not output_final_state else last_recurrent_state + core_attn_out = core_attn_out.reshape(padded_output_shape)[:, :, :sequence_length] + core_attn_out = core_attn_out.transpose(1, 2).to(initial_dtype, memory_format=torch.contiguous_format) + return core_attn_out, last_recurrent_state + + +def patch_mps(): + """Patch Qwen3.5 operations on MPS tensors while preserving other backends.""" + global _compat_warning_reported + if not torch.backends.mps.is_available(): + return False + try: + import transformers + import transformers.models.qwen3_5.modeling_qwen3_5 as mq + version = tuple(int(part) for part in transformers.__version__.split(".")[:2]) + if not all(callable(getattr(mq, name, None)) for name in ("torch_chunk_gated_delta_rule", "causal_conv1d_fn")): + raise AttributeError("unsupported Transformers Qwen3.5 implementation") + params = inspect.signature(mq.torch_chunk_gated_delta_rule).parameters + required = {"query", "key", "value", "g", "beta", "chunk_size"} + if version < (5, 17) or not required.issubset(params): + if not _compat_warning_reported: + warnings.warn(f"MPS patch requires Transformers >=5.17 with the Qwen3.5 gated-delta signature; got {transformers.__version__}", RuntimeWarning, stacklevel=2) + _compat_warning_reported = True + return False + if getattr(mq.torch_chunk_gated_delta_rule, "_decider_mps_patch", False): + return True + + original_delta, original_conv = mq.torch_chunk_gated_delta_rule, mq.causal_conv1d_fn + delta_params = tuple(inspect.signature(original_delta).parameters) + query_index, chunk_index = delta_params.index("query"), delta_params.index("chunk_size") + + @functools.wraps(original_delta) + def delta(*args, **kwargs): + query = args[query_index] if len(args) > query_index else kwargs["query"] + chunk_size = args[chunk_index] if len(args) > chunk_index else kwargs.get("chunk_size", 64) + return mps_chunk_gated_delta_rule(*args, **kwargs) if query.is_mps and chunk_size == 64 else original_delta(*args, **kwargs) + + @functools.wraps(original_conv) + def conv(hidden_states, *args, **kwargs): + return fused_causal_conv1d_fn(hidden_states, *args, **kwargs) if hidden_states.is_mps else original_conv(hidden_states, *args, **kwargs) + + delta._decider_mps_patch = True + mq.torch_chunk_gated_delta_rule, mq.causal_conv1d_fn = delta, conv + return True + except (ImportError, AttributeError, TypeError, ValueError): + logging.getLogger(__name__).debug("MPS patch unavailable; using Transformers fallback", exc_info=True) + return False diff --git a/decider/vision/model.py b/decider/vision/model.py index 6d55711..78c6710 100644 --- a/decider/vision/model.py +++ b/decider/vision/model.py @@ -14,8 +14,10 @@ def to_pil(x): class VisionDecisionModel(nn.Module): - def __init__(self, name, dtype=torch.bfloat16, grad_ckpt=True): + def __init__(self, name, dtype=None, grad_ckpt=True): super().__init__() + if dtype is None: + dtype = torch.float16 if torch.backends.mps.is_available() else torch.bfloat16 self.proc = AutoProcessor.from_pretrained(name); self.tok = self.proc.tokenizer self.lm = AutoModelForImageTextToText.from_pretrained(name, dtype=dtype) if grad_ckpt: self.lm.gradient_checkpointing_enable() @@ -49,6 +51,9 @@ def prepare(self, examples, max_ctx_tokens=1536): def slot_logits(self, inp): dev = self.letters.device + if dev.type == "mps": + from decider.mps_ops import patch_mps + patch_mps() kw = {k: v.to(dev) for k, v in inp.items() if k in ("input_ids", "attention_mask", "pixel_values", "image_grid_thw", "mm_token_type_ids")} h = self.lm.model(**kw, use_cache=False).last_hidden_state # [B, T, H] hs = h[inp["slot_batch"].to(dev), inp["slot_idx"].to(dev)] # [N, H] diff --git a/docs/benchmarks/mps-full-model.json b/docs/benchmarks/mps-full-model.json new file mode 100644 index 0000000..6a22d81 --- /dev/null +++ b/docs/benchmarks/mps-full-model.json @@ -0,0 +1,580 @@ +{ + "runs": [ + { + "model": "Mapika/decider-0.8b", + "snapshot_revision": "1ea54127d3bd52f6d753d9257b32a6380b873907", + "mode": "conv", + "patch_mps": false, + "patch_conv": true, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9906930327415466, + 0.008457512594759464, + 0.0008494335343129933 + ], + "times_ms": [ + 147.1938329996192, + 149.08208300039405, + 151.9858749998093, + 153.05950000038138, + 146.1539170013566 + ], + "median_ms": 149.08208300039405 + }, + { + "case": 1, + "probs": [ + 0.005709604825824499, + 0.9921925663948059, + 0.002097897697240114 + ], + "times_ms": [ + 147.69070799957262, + 165.61800000090443, + 156.68779199950222, + 172.6192499991157, + 161.34441699978197 + ], + "median_ms": 161.34441699978197 + }, + { + "case": 2, + "probs": [ + 0.3767502009868622, + 0.011073778383433819, + 0.6121760010719299 + ], + "times_ms": [ + 144.24420800060034, + 146.9018749994575, + 146.99420899887627, + 152.0920839993778, + 153.09554200030107 + ], + "median_ms": 146.99420899887627 + } + ] + }, + { + "model": "Mapika/decider-0.8b", + "snapshot_revision": "1ea54127d3bd52f6d753d9257b32a6380b873907", + "mode": "optimized", + "patch_mps": true, + "patch_conv": true, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9906930327415466, + 0.008457512594759464, + 0.0008494335343129933 + ], + "times_ms": [ + 173.9600420005445, + 170.48820800118847, + 145.8357089995843, + 128.68499999967753, + 99.56808300012199 + ], + "median_ms": 145.8357089995843 + }, + { + "case": 1, + "probs": [ + 0.005752823781222105, + 0.9921493530273438, + 0.002097806427627802 + ], + "times_ms": [ + 98.5057500001858, + 96.95287500107952, + 101.59995800131583, + 99.90291600115597, + 98.24325000045064 + ], + "median_ms": 98.5057500001858 + }, + { + "case": 2, + "probs": [ + 0.3767186999320984, + 0.01115715503692627, + 0.6121242046356201 + ], + "times_ms": [ + 97.34708300129569, + 99.5587079996767, + 98.98433399939677, + 101.38904199993704, + 135.6377079991944 + ], + "median_ms": 99.5587079996767 + } + ] + }, + { + "model": "Mapika/decider-0.8b", + "snapshot_revision": "1ea54127d3bd52f6d753d9257b32a6380b873907", + "mode": "reference", + "patch_mps": false, + "patch_conv": false, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9908191561698914, + 0.008331239223480225, + 0.0008495416841469705 + ], + "times_ms": [ + 143.51175000047078, + 140.9479169997212, + 140.79929100080335, + 150.36620799946832, + 156.48395900097967 + ], + "median_ms": 143.51175000047078 + }, + { + "case": 1, + "probs": [ + 0.0057962816208601, + 0.992090106010437, + 0.002113653812557459 + ], + "times_ms": [ + 142.38641699921573, + 142.81708399903437, + 142.24654200006626, + 141.35062499917694, + 142.25687500038475 + ], + "median_ms": 142.25687500038475 + }, + { + "case": 2, + "probs": [ + 0.3784691095352173, + 0.011208995245397091, + 0.6103218197822571 + ], + "times_ms": [ + 161.90645799906633, + 159.2158749990631, + 155.37495900025533, + 171.4579590006906, + 178.3607920006034 + ], + "median_ms": 161.90645799906633 + } + ] + }, + { + "model": "Mapika/decider-2b", + "snapshot_revision": "b37f7e1ba3fbc9238004cf531fabbee2619973fd", + "mode": "conv", + "patch_mps": false, + "patch_conv": true, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9854162335395813, + 0.00705093564465642, + 0.00753279123455286 + ], + "times_ms": [ + 171.3152499996795, + 172.89420800079824, + 182.06000000100175, + 170.31337499975052, + 175.06854199928057 + ], + "median_ms": 172.89420800079824 + }, + { + "case": 1, + "probs": [ + 0.003030716907233, + 0.9943298101425171, + 0.0026394694577902555 + ], + "times_ms": [ + 172.5064160000329, + 173.27641699921514, + 173.96304099929694, + 173.4836669984361, + 174.41150000013295 + ], + "median_ms": 173.4836669984361 + }, + { + "case": 2, + "probs": [ + 0.36520224809646606, + 0.011323649436235428, + 0.62347412109375 + ], + "times_ms": [ + 171.85408399927837, + 174.10445799941954, + 178.6571249995177, + 170.36533300051815, + 171.12066700065043 + ], + "median_ms": 171.85408399927837 + } + ] + }, + { + "model": "Mapika/decider-2b", + "snapshot_revision": "b37f7e1ba3fbc9238004cf531fabbee2619973fd", + "mode": "optimized", + "patch_mps": true, + "patch_conv": true, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9854606986045837, + 0.007051253691315651, + 0.007487996015697718 + ], + "times_ms": [ + 130.30574999902456, + 132.46291700124857, + 136.03441700070107, + 132.96358300067368, + 137.251334001121 + ], + "median_ms": 132.96358300067368 + }, + { + "case": 1, + "probs": [ + 0.003030716907233, + 0.9943298101425171, + 0.0026394694577902555 + ], + "times_ms": [ + 136.81341699884797, + 132.82266600072035, + 127.78708299993013, + 129.8745419990155, + 129.06720899991342 + ], + "median_ms": 129.8745419990155 + }, + { + "case": 2, + "probs": [ + 0.36383503675460815, + 0.011281256563961506, + 0.6248837113380432 + ], + "times_ms": [ + 130.87658299991745, + 131.76137499976903, + 138.73420799973246, + 135.48220899974694, + 135.69687499875727 + ], + "median_ms": 135.48220899974694 + } + ] + }, + { + "model": "Mapika/decider-2b", + "snapshot_revision": "b37f7e1ba3fbc9238004cf531fabbee2619973fd", + "mode": "reference", + "patch_mps": false, + "patch_conv": false, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9854187965393066, + 0.007093454711139202, + 0.007487677503377199 + ], + "times_ms": [ + 172.34524999912537, + 170.9845419991325, + 173.80195800069487, + 170.2168329993583, + 170.33079099928727 + ], + "median_ms": 170.9845419991325 + }, + { + "case": 1, + "probs": [ + 0.003030668944120407, + 0.9943140745162964, + 0.002655337331816554 + ], + "times_ms": [ + 169.0252499993221, + 170.86904200004938, + 178.09154200040211, + 167.71116599920788, + 182.68116699982784 + ], + "median_ms": 170.86904200004938 + }, + { + "case": 2, + "probs": [ + 0.36381006240844727, + 0.01134848315268755, + 0.6248414516448975 + ], + "times_ms": [ + 169.43104099846096, + 169.8419999993348, + 171.94458399899304, + 172.01070899864135, + 176.3894169998821 + ], + "median_ms": 171.94458399899304 + } + ] + }, + { + "model": "Mapika/decider-2b-vision", + "snapshot_revision": "446c8c5e334e53ae3526a1c5384f93d5caee68cd", + "mode": "conv", + "patch_mps": false, + "patch_conv": true, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9997536540031433, + 0.00013710305211134255, + 0.00010930807911790907 + ], + "times_ms": [ + 398.3357500001148, + 397.08587499990244, + 407.70579199852364, + 492.72212500000023, + 406.95045800021035 + ], + "median_ms": 406.95045800021035 + }, + { + "case": 1, + "probs": [ + 0.00019713742949534208, + 0.999640703201294, + 0.00016216083895415068 + ], + "times_ms": [ + 398.63974999934726, + 400.94854200106056, + 394.2771250003716, + 395.4182080015016, + 393.89945799848647 + ], + "median_ms": 395.4182080015016 + }, + { + "case": 2, + "probs": [ + 0.0002715098671615124, + 0.0003028911305591464, + 0.9994256496429443 + ], + "times_ms": [ + 391.624416999548, + 404.1594589998567, + 396.4951250000013, + 393.24149999993097, + 393.2371659993805 + ], + "median_ms": 393.24149999993097 + } + ] + }, + { + "model": "Mapika/decider-2b-vision", + "snapshot_revision": "446c8c5e334e53ae3526a1c5384f93d5caee68cd", + "mode": "optimized", + "patch_mps": true, + "patch_conv": true, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9997536540031433, + 0.00013710305211134255, + 0.00010930807911790907 + ], + "times_ms": [ + 275.2892079988669, + 280.462875000012, + 277.7905829989322, + 277.3510830011219, + 274.18125000076543 + ], + "median_ms": 277.3510830011219 + }, + { + "case": 1, + "probs": [ + 0.00019560383225325495, + 0.9996434450149536, + 0.00016089931887108833 + ], + "times_ms": [ + 276.64883299985377, + 277.90812499915774, + 279.31687499949476, + 274.3287089997466, + 275.2075000007608 + ], + "median_ms": 276.64883299985377 + }, + { + "case": 2, + "probs": [ + 0.0002715098671615124, + 0.0003028911305591464, + 0.9994256496429443 + ], + "times_ms": [ + 275.93125000021246, + 277.46612500050105, + 279.31904199977, + 277.4788329988951, + 281.12366699861013 + ], + "median_ms": 277.4788329988951 + } + ] + }, + { + "model": "Mapika/decider-2b-vision", + "snapshot_revision": "446c8c5e334e53ae3526a1c5384f93d5caee68cd", + "mode": "reference", + "patch_mps": false, + "patch_conv": false, + "torch": "2.14.0", + "transformers": "5.17.0", + "dtype": "torch.float16", + "hardware": "arm64", + "macOS": "27.2", + "warmups": 2, + "measurements_per_case": 5, + "results": [ + { + "case": 0, + "probs": [ + 0.9997536540031433, + 0.00013710305211134255, + 0.00010930807911790907 + ], + "times_ms": [ + 381.8799999990006, + 388.8329169985809, + 380.74745799895027, + 387.30608399964694, + 386.5536670000438 + ], + "median_ms": 386.5536670000438 + }, + { + "case": 1, + "probs": [ + 0.00019713766232598573, + 0.9996418952941895, + 0.00016089907148852944 + ], + "times_ms": [ + 385.64641700031643, + 384.3574160000571, + 380.93866600138426, + 388.39999999981956, + 388.53820800068206 + ], + "median_ms": 385.64641700031643 + }, + { + "case": 2, + "probs": [ + 0.0002715098671615124, + 0.0003028911305591464, + 0.9994256496429443 + ], + "times_ms": [ + 388.81791600033466, + 390.2366250003979, + 391.11066599980404, + 385.1849580005364, + 384.0816249994532 + ], + "median_ms": 388.81791600033466 + } + ] + } + ] +} diff --git a/docs/benchmarks/mps-full-model.md b/docs/benchmarks/mps-full-model.md new file mode 100644 index 0000000..390dfca --- /dev/null +++ b/docs/benchmarks/mps-full-model.md @@ -0,0 +1,58 @@ +# Full-model MPS measurements + +Apple M1 Pro, 32 GB; macOS 27.2; PyTorch 2.14.0; Transformers 5.17.0. +All arms use float16 weights on MPS. MLX/Metal was available for the optimized path. + +Each row is the median of five synchronized complete request measurements after two +warmups. Model loading is excluded. Text timing includes `Decider.decide()`; vision +timing includes image preparation, `slot_logits()`, softmax, and transfer of +probabilities to CPU. Each arm runs in a separate process, sequentially, with +identical locally cached checkpoints and inputs. These are nine specific smoke-test +workloads, not a representative accuracy suite. + +The **reference** arm uses Transformers' pure-PyTorch gated-delta and convolution +implementations. In the evidence environment, `flash-linear-attention` 0.5.2 was +installed, but its `fla.ops` import required unavailable Triton support on macOS, so +Transformers selected its pure-PyTorch fallback and the unpatched MPS path ran. The +FLA/Triton implementation itself was not tested. The **conv-only** arm adds Decider's existing fused causal +convolution. The **optimized** arm adds the MPS gated-delta patch as well. This +separates the attention contribution from the already-existing convolution patch. + +| Model | Input | Reference ms | Conv-only ms | Optimized ms | Optimized speedup | Maximum probability difference | +|---|---|---:|---:|---:|---:|---:| +| 0.8B | Billing request | 143.51 | 149.08 | 145.84 | 0.98× | 0.0001263 | +| 0.8B | Login request | 142.26 | 161.34 | 98.51 | 1.44× | 0.0000592 | +| 0.8B | Sales request | 161.91 | 146.99 | 99.56 | 1.63× | 0.0018024 | +| 2B | Billing request | 170.98 | 172.89 | 132.96 | 1.29× | 0.0000422 | +| 2B | Login request | 170.87 | 173.48 | 129.87 | 1.32× | 0.0000159 | +| 2B | Sales request | 171.94 | 171.85 | 135.48 | 1.27× | 0.0000672 | +| 2B Vision | Red image | 386.55 | 406.95 | 277.35 | 1.39× | 0 | +| 2B Vision | Green image | 385.65 | 395.42 | 276.65 | 1.39× | 0.0000015 | +| 2B Vision | Blue image | 388.82 | 393.24 | 277.48 | 1.40× | 0 | + +All nine top-ranked answers matched. The largest probability change was 0.0018024 +(0.18024 percentage points) on the 0.8B sales request. This exceeds a provisional +0.001 absolute probability-difference check; exact full-model numerical parity is +not claimed. Close decision boundaries may be sensitive to numerical differences. +Vision inputs are synthetic 224×224 solid-color images, not real-world vision data. +The separate [held-out evaluation](mps-heldout.md) covers 1,500 MASSIVE Scenario +examples; no aggregate held-out score or long-context evaluation was run here. + +## Reproduce + +Install the project and optional `metal` dependencies, and cache the three model +checkpoints. Run without other GPU work: + +```sh +for model in decider-0.8b decider-2b decider-2b-vision; do + for mode in reference conv optimized; do + HF_HUB_OFFLINE=1 python -m decider.bench.mps "Mapika/$model" "$mode" > "$model-$mode.json" + done +done +``` + +See [raw measurements](mps-full-model.json) for all timings, probabilities, +checkpoint revisions, dtype, patch return values, and software versions. The +runnable input definitions are in `decider/bench/mps.py`. +The optional Metal path accelerates part of the PyTorch computation; this is not +an all-MLX model. Full-model speedups include both attention and convolution changes. diff --git a/docs/benchmarks/mps-heldout.json b/docs/benchmarks/mps-heldout.json new file mode 100644 index 0000000..6d78e9c --- /dev/null +++ b/docs/benchmarks/mps-heldout.json @@ -0,0 +1,42 @@ +{ + "results": { + "massive_scenario": { + "n": 1500, + "acc": 0.7553333333333333, + "nll": 0.6979675889015198, + "brier": 0.3348042070865631, + "ece": 0.043757498661677025, + "aurc": 0.08389327558743383, + "acc_at_80": 0.845, + "acc_at_50": 0.936, + "chance": 0.1, + "mean_conf": 0.7968708276748657, + "sec": 210.7, + "heldout": true + } + }, + "agg": { + "in_task": {}, + "heldout": { + "acc": 0.7553333333333333, + "nll": 0.6979675889015198, + "brier": 0.3348042070865631, + "ece": 0.043757498661677025, + "aurc": 0.08389327558743383, + "acc_at_80": 0.845, + "chance": 0.1 + }, + "n_in": 0, + "n_heldout": 1 + }, + "model": "Mapika/decider-2b", + "engine": "", + "device": "mps", + "dtype": "torch.float16", + "snapshot_revision": "b37f7e1ba3fbc9238004cf531fabbee2619973fd", + "published_bf16_massive_scenario": { + "accuracy": 0.756, + "ece": 0.041, + "source": "MODEL_CARD.md per-task held-out table" + } +} diff --git a/docs/benchmarks/mps-heldout.md b/docs/benchmarks/mps-heldout.md new file mode 100644 index 0000000..d00473b --- /dev/null +++ b/docs/benchmarks/mps-heldout.md @@ -0,0 +1,44 @@ +# Held-out MPS evaluation + +This is the requested `decider.evaluate.py` check, not a smoke-test prompt. +It evaluates all 1,500 examples in the held-out `massive_scenario` test set from +the repository's data loader. + +| Model/path | Examples | Temperature | Accuracy | ECE | NLL | +|---|---:|---:|---:|---:|---:| +| Published decider-2b v10 BF16, rebuilt-set reference | 1,500-task row | 1.30 | 0.756 | 0.041 | — | +| This PR, decider-2b FP16 on MPS | 1,500 | 1.30 | 0.7553 | 0.0438 | 0.6980 | + +The published per-task table reports accuracy/ECE, not per-task NLL. Its MASSIVE +Scenario row is `0.756 / 0.041`; the aggregate held-out NLL elsewhere in the model +card is not comparable to this single-task NLL. + +The MPS result was produced by `python -m decider.evaluate` with `--device mps`, +`--temperature 1.30`, batch size 8, and max context 1536. It used the cached +`Mapika/decider-2b` snapshot `b37f7e1ba3fbc9238004cf531fabbee2619973fd`, PyTorch +2.14.0, Transformers 5.17.0, and float16 weights. The run took 210.7 seconds; +model loading is excluded from the reported task timer. + +This is one held-out dataset, not the aggregate 24-task published score. It shows +that the MPS path's accuracy and calibration are close to the published BF16 row +on this set; it does not establish long-context behavior or the aggregate held-out +score. The raw `evaluate.py` output is recorded in +[`mps-heldout.json`](mps-heldout.json). + +## Reproduce + +The repository does not commit generated dataset caches. Build the same one-task +cache and run the evaluator on an Apple Silicon MPS machine: + +```sh +python -m decider.data.core massive_scenario --out /tmp/massive_scenario.pkl +python -m decider.evaluate \ + --model Mapika/decider-2b \ + --data /tmp/massive_scenario.pkl \ + --out /tmp/mps-eval \ + --tasks massive_scenario \ + --device mps \ + --temperature 1.30 \ + --bs 8 \ + --max_ctx 1536 +``` diff --git a/pyproject.toml b/pyproject.toml index 2f31554..1122123 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,6 +12,7 @@ dependencies = ["torch", "transformers>=5", "flash-linear-attention", "numpy<2", serve = ["fastapi", "uvicorn[standard]", "httpx"] train = ["datasets", "accelerate", "scikit-learn", "pillow"] games = ["gymnasium", "gym-super-mario-bros", "nes-py", "minigrid", "ale-py", "imageio"] +metal = ["mlx>=0.20; sys_platform == 'darwin' and platform_machine == 'arm64'"] [project.urls] Repository = "https://github.com/Mapika/decider" diff --git a/tests/test_mps_ops.py b/tests/test_mps_ops.py new file mode 100644 index 0000000..27fe605 --- /dev/null +++ b/tests/test_mps_ops.py @@ -0,0 +1,123 @@ +import inspect + +import pytest +torch = pytest.importorskip("torch") +F = pytest.importorskip("torch.nn.functional") +from transformers.models.qwen3_5.modeling_qwen3_5 import torch_chunk_gated_delta_rule as _dispatch_delta_rule + +# Transformers decorates this function to prefer FLA when installed. Unwrap it +# so these tests always compare with the pure-PyTorch reference implementation. +torch_chunk_gated_delta_rule = inspect.unwrap(_dispatch_delta_rule) +from decider.mps_ops import mps_chunk_gated_delta_rule, fast_invert_unitriangular_64, patch_mps + + +@pytest.mark.parametrize("disable_metal", [False, True]) +def test_unitriangular_inverse(monkeypatch, disable_metal): + """Check both optional Metal and PyTorch fallback inversion.""" + if disable_metal: + monkeypatch.setattr("decider.mps_ops._metal_invert_kernel", None) + device = "mps" if torch.backends.mps.is_available() else "cpu" + torch.manual_seed(42) + L = torch.randn(16, 64, 64, device=device).tril(-1) * 0.1 + inv = fast_invert_unitriangular_64(L) + I = torch.eye(64, device=device).expand_as(L) + A = I + L + prod = A @ inv + err = (prod - I).abs().max().item() + assert err < 1e-5, f"Inversion error too high: {err}" + print(f"[PASS] test_unitriangular_inverse max err: {err:.2e}") + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS required") +def test_metal_failure_falls_back(monkeypatch): + import decider.mps_ops as ops + + class BrokenKernel: + def __call__(self, **kwargs): + raise RuntimeError("test Metal failure") + + monkeypatch.setattr(ops, "_metal_invert_kernel", BrokenKernel()) + monkeypatch.setattr(ops, "_metal_failure_reported", False) + L = torch.randn(1, 64, 64, device="mps", dtype=torch.float32).tril(-1) * 0.1 + inv = ops.fast_invert_unitriangular_64(L) + I = torch.eye(64, device="mps").expand_as(L) + assert torch.allclose((I + L) @ inv, I, atol=1e-5, rtol=1e-5) + assert ops._metal_invert_kernel is None + + +def test_delta_rule_shapes(): + """Verify numerical correctness against reference for B in (1, 2), S in (64, 512, 1024, 1596).""" + device = "mps" if torch.backends.mps.is_available() else "cpu" + shapes = [64, 512, 1024, 1596] + H, D = 16, 128 + + for B in [1, 2]: + for S in shapes: + torch.manual_seed(S + B * 1000) + q = torch.randn(B, S, H, D, device=device, dtype=torch.float32) + k = torch.randn(B, S, H, D, device=device, dtype=torch.float32) / (D ** 0.5) + v = torch.randn(B, S, H, D, device=device, dtype=torch.float32) + g = -torch.rand(B, S, H, device=device, dtype=torch.float32) + beta = torch.sigmoid(torch.randn(B, S, H, device=device, dtype=torch.float32)) + + ref_out, ref_state = torch_chunk_gated_delta_rule(q, k, v, g, beta, output_final_state=True) + opt_out, opt_state = mps_chunk_gated_delta_rule(q, k, v, g, beta, output_final_state=True) + + max_diff = (ref_out - opt_out).abs().max().item() + cos_sim = F.cosine_similarity(ref_out.flatten(), opt_out.flatten(), dim=0).item() + state_diff = (ref_state - opt_state).abs().max().item() + + assert torch.allclose(ref_out, opt_out, atol=1e-3, rtol=1e-3), f"Failed for B={B}, S={S}, max_diff={max_diff}" + assert torch.allclose(ref_state, opt_state, atol=1e-3, rtol=1e-3), f"State diff failed for B={B}, S={S}: {state_diff}" + print(f"[PASS] B={B} S={S:4d}: max_diff={max_diff:.2e}, cos_sim={cos_sim:.8f}, state_diff={state_diff:.2e}") + + +def test_patch_dispatch_and_guards(monkeypatch): + """Exercise dispatch without requiring MPS hardware or leaving global patches behind.""" + from types import SimpleNamespace + import transformers + import transformers.models.qwen3_5.modeling_qwen3_5 as mq + import decider.mps_ops as ops + monkeypatch.setattr(transformers, "__version__", "5.17.0") + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: False) + assert not patch_mps() + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + def reference_delta(query, key=None, value=None, g=None, beta=None, chunk_size=64, **kwargs): + return "reference delta" + + def reference_conv(hidden_states, *args, **kwargs): + return "reference conv" + + monkeypatch.setattr(mq, "torch_chunk_gated_delta_rule", reference_delta) + monkeypatch.setattr(mq, "causal_conv1d_fn", reference_conv) + monkeypatch.setattr(ops, "mps_chunk_gated_delta_rule", lambda *a, **kw: "MPS delta") + monkeypatch.setattr(ops, "fused_causal_conv1d_fn", lambda *a, **kw: "MPS conv") + assert patch_mps() + patched = mq.torch_chunk_gated_delta_rule + assert patch_mps() and mq.torch_chunk_gated_delta_rule is patched + for is_mps in (False, True): + tensor = SimpleNamespace(is_mps=is_mps) + prefix = "MPS" if is_mps else "reference" + assert patched(query=tensor) == prefix + " delta" + assert mq.causal_conv1d_fn(tensor) == prefix + " conv" + assert patched(tensor, chunk_size=32) == "reference delta" + monkeypatch.setattr(transformers, "__version__", "5.16.0") + with pytest.warns(RuntimeWarning, match="requires Transformers"): + assert not patch_mps() + + +def test_l2norm_option(): + """Verify use_qk_l2norm_in_kernel branch.""" + device = "mps" if torch.backends.mps.is_available() else "cpu" + B, S, H, D = 1, 128, 16, 128 + torch.manual_seed(123) + q = torch.randn(B, S, H, D, device=device, dtype=torch.float32) + k = torch.randn(B, S, H, D, device=device, dtype=torch.float32) + v = torch.randn(B, S, H, D, device=device, dtype=torch.float32) + g = -torch.rand(B, S, H, device=device, dtype=torch.float32) + beta = torch.sigmoid(torch.randn(B, S, H, device=device, dtype=torch.float32)) + + ref_out, _ = torch_chunk_gated_delta_rule(q, k, v, g, beta, use_qk_l2norm_in_kernel=True) + opt_out, _ = mps_chunk_gated_delta_rule(q, k, v, g, beta, use_qk_l2norm_in_kernel=True) + assert torch.allclose(ref_out, opt_out, atol=1e-3, rtol=1e-3) + print("[PASS] test_l2norm_option")