decider-35b-a3b and decider-2b on Apple Silicon: works, JevBench reproduced; two MPS patches make the 35B 8x faster

[nassersala]

Following up on #5, a report from a Mac (M5 Max, 128 GB, macOS 26.5.2, torch 2.14.0, transformers 5.17.0), library path
(Decider(path, device="mps", dtype=torch.bfloat16, use_graphs=False)), not the HTTP server.

Install. flash-linear-attention does not install here, so: pip install torch "transformers>=5.17" "numpy<2" huggingface_hub and then the package with --no-deps (we used git+https://github.com/Mapika/decider@c4daaac, before
1.1.2). transformers falls back to its reference kernels.

Fidelity. JevBench public items, easy / standard / hard:

• decider-2b v10: 48/48, 63/72, 51/111 (card: 48, 61, 51)
• decider-35b-a3b v1: 48/48, 70/72, 75/111 (card: 48, 70, 75, exact)

Memory. The 35B loads in bf16 in about 60 s and runs in 128 GB of unified memory, so the card's "one 80 GB GPU is the
minimum" has a Mac alternative.

Speed. Out of the box the 35B takes 2.8-4 s a decision; nearly all of it is two MPS ops in transformers' reference
code, not the model:

• torch.histc in integrations/moe.py (tokens per expert): ~45 ms per call at 256 experts, once per layer;
• torch.linalg.solve_triangular in the gated delta rule: ~17 ms per call, twice per linear-attention layer.

Replacing them in-process (a scatter_add_ count, and a block-doubling inverse of the unit lower-triangular system) gives
0.23-0.5 s a decision for typical inputs and 2.5 s for a 3,900-token one (was 11 s). The histc swap is bit-identical;
the solve swap moves probabilities by less than switching to the exact CPU solver does, no argmax changed, and the JevBench
counts above were taken with both patches on. The patch is mps_fix.py below (call mps_fix.install() before loading). Upstream: huggingface/transformers#49027 (the histc fix) and pytorch/pytorch#198304 (both ops); until then it may be worth a line in the card or an
opt-in in Decider(..., device="mps").

decider-2b is dense, so only the solve patch applies; I did not measure it with the patch.

mps_fix.py

"""Two MPS slow paths in transformers' Qwen3.5-MoE reference code, replaced in-process for decider-35b-a3b.

On an M-series Mac (torch 2.14, transformers 5.17) torch.histc takes 40 ms a call on MPS (once per MoE layer) and
torch.linalg.solve_triangular 17 ms (twice per linear-attention layer): 2.8 s of a 3 s decision. Checked on six items
(2026-09-23): histc swap bit-identical; block inverse within 1e-6 of the MPS solver on real systems, probabilities move
at most 0.033, less than switching to the exact CPU solver does (0.042); no argmax changed.
"""
import torch

_histc, _solve = torch.histc, torch.linalg.solve_triangular


def histc(x, bins=100, min=0, max=0, **kw):
    # expert ids 0..n-1 into n bins over [0, n-1]: bin k holds id k exactly, so this is a count (histc: 40 ms on MPS)
    if x.device.type == "mps" and x.dim() == 1 and min == 0 and bins == max - min + 1 and not kw:
        return torch.zeros(bins, device=x.device, dtype=x.dtype).scatter_add_(0, x.long(), torch.ones_like(x))
    return _histc(x, bins=bins, min=min, max=max, **kw)


def unit_lower_inverse(A):
    """Inverse of unit lower-triangular A by block doubling: [[A11,0],[A21,A22]]^-1 = [[X11,0],[-X22 A21 X11, X22]]."""
    n = A.shape[-1]
    inv = torch.ones(*A.shape[:-2], n, 1, 1, device=A.device, dtype=A.dtype)  # n diagonal 1x1 blocks
    b = 1
    while b < n:
        nb = n // (2 * b)
        blocks = A.reshape(*A.shape[:-2], nb, 2 * b, nb, 2 * b).diagonal(dim1=-4, dim2=-2).movedim(-1, -3)
        x11, x22, a21 = inv[..., 0::2, :, :], inv[..., 1::2, :, :], blocks[..., b:, :b]
        new = torch.zeros(*A.shape[:-2], nb, 2 * b, 2 * b, device=A.device, dtype=A.dtype)
        new[..., :b, :b], new[..., b:, b:], new[..., b:, :b] = x11, x22, -(x22 @ a21 @ x11)
        inv, b = new, 2 * b
    return inv[..., 0, :, :]


def solve_triangular(A, B, *, upper, left=True, unitriangular=False, out=None):
    # unit lower-triangular: invert by block doubling, all batched matmuls (the MPS solver: 17 ms per call)
    n = A.shape[-1]
    if A.device.type == "mps" and not upper and left and unitriangular and out is None and n & (n - 1) == 0:
        return unit_lower_inverse(A.tril(-1) + torch.eye(n, device=A.device, dtype=A.dtype)) @ B
    return _solve(A, B, upper=upper, left=left, unitriangular=unitriangular, out=out)


def install():
    torch.histc, torch.linalg.solve_triangular = histc, solve_triangular


def uninstall():
    torch.histc, torch.linalg.solve_triangular = _histc, _solve