Accelerate Decider 0.8B, 2B, and 2B Vision on Apple Silicon

[simply-sunny]

Summary

Adds Apple Silicon MPS acceleration for Decider 0.8B, 2B, and 2B Vision through their shared dense Qwen3.5 implementation.

• Automatically selects CUDA, MPS, or CPU in Decider; MPS defaults to float16 and logs the resolved device/dtype once per construction.
• Adds optional MLX/Metal-assisted gated-delta attention and reuses Decider's fused causal convolution.
• Keeps the Transformers patch narrow: MPS tensors use the optimized path; CPU/CUDA tensors and unsupported chunk sizes retain the original functions.
• Requires Transformers >=5.17 with the expected gated-delta signature; incompatible versions warn once and retain the reference path.
• A failed MLX/DLPack operation disables the optional kernel for the process and falls back to the PyTorch inverse with a warning.
• Adds --device support to decider.evaluate.py so MPS evaluation uses the same evaluator and metrics as the published results.

Full-model measurements

Apple M1 Pro · 32 GB · macOS 27.2 · PyTorch 2.14.0 · Transformers 5.17.0 · FP16

The benchmark has three arms: pure Transformers reference, reference plus Decider's existing fused causal convolution, and the full MPS patch.

Model	Input	Reference	Conv-only	Optimized	Optimized speedup
0.8B	Billing	143.51 ms	149.08 ms	145.84 ms	0.98×
0.8B	Login	142.26 ms	161.34 ms	98.51 ms	1.44×
0.8B	Sales	161.91 ms	146.99 ms	99.56 ms	1.63×
2B	Billing	170.98 ms	172.89 ms	132.96 ms	1.29×
2B	Login	170.87 ms	173.48 ms	129.87 ms	1.32×
2B	Sales	171.94 ms	171.85 ms	135.48 ms	1.27×
2B Vision	Red image	386.55 ms	406.95 ms	277.35 ms	1.39×
2B Vision	Green image	385.65 ms	395.42 ms	276.65 ms	1.39×
2B Vision	Blue image	388.82 ms	393.24 ms	277.48 ms	1.40×

Each value is the median of five synchronized complete-request measurements after two warmups. Identical locally cached checkpoints and inputs were run in separate processes, sequentially; model loading is excluded. The committed JSON is direct output from the benchmark script and records checkpoint revision, patch states, dtype, versions, platform, timings, medians, and probabilities.

All nine top-ranked answers matched. Maximum absolute probability differences between reference and optimized were 0.0018024 (0.8B), 0.0000672 (2B), and 0.0000015 (2B Vision). Exact full-model numerical parity is not claimed. Vision inputs are synthetic 224×224 solid-color images.

Held-out calibration check

Ran decider.evaluate.py on all 1,500 examples in the held-out massive_scenario test set using the optimized 2B FP16 MPS path, temperature 1.30, batch size 8, and max context 1536:

Path	Accuracy	ECE
Published decider-2b v10 BF16 rebuilt-set row	0.756	0.041
This PR: decider-2b FP16 on MPS	0.7553	0.0438

The MPS run used snapshot b37f7e1ba3fbc9238004cf531fabbee2619973fd and took 210.7 seconds excluding model loading. This is one held-out dataset, not the aggregate 24-task score; long-context behavior was not evaluated. Full details and raw evaluator output are in docs/benchmarks/mps-heldout.*.

Reference-path behavior

With flash-linear-attention 0.5.2 installed, the unpatched dispatcher on the evidence Mac selected the pure-PyTorch path because importing fla.ops required Triton, unavailable on macOS. The unpatched MPS path ran successfully. The FLA/Triton implementation itself was not tested; the benchmark explicitly unwraps the Transformers dispatcher for its reference arm.

Validation

• 24 tests passed with FLA installed; the pure-PyTorch reference is selected explicitly in comparison tests.
• Metal/PyTorch inversion, runtime fallback, reference numerical comparison, device dispatch, compatibility signature guard, and repeated-patching tests.
• Full-model three-arm checks for all three models.
• Held-out evaluation through decider.evaluate.py on MPS.
• git diff --check and benchmark compilation passed.

Reproduce the full-model measurements with locally cached checkpoints:

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

Reproduce the held-out check by following docs/benchmarks/mps-heldout.md.

[Mapika]

Thank you for this work. The limitations section is clear about what was and was not tested, and the patch dispatches per tensor, which is the correct design. I read every changed line and found no way for the change to alter CUDA or CPU behaviour, the CUDA-graph engine, the FP8 path or the HTTP server. I compared the replacement delta rule with the Transformers 5.17 reference line by line and re-derived both inversion routines independently. The maths is equivalent, and the Metal kernel's recurrence, indexing and threadgroup memory use are correct. I would like the following changed or answered before merging.

1. Evidence. Every comparison in the PR is MPS against MPS on nine short inputs. This model's output is a calibrated probability, so I would like one held-out set run through decider/evaluate.py on MPS, with accuracy and ECE reported next to the published bf16 numbers. If that is more than you want to take on, please add a sentence to the doc stating that the MPS path is validated for output agreement on nine short inputs and not for accuracy or calibration. I ran the same three prompts on the 0.8B on CPU. The largest activation anywhere in the model is 81, against an FP16 limit of 65504, and FP16 tracks FP32 more closely than the released bf16 does. So I am not concerned about FP16 at this input length. Long context is untested.

2. The numbers in docs/benchmarks/mps-full-model.json. I reproduced the three 0.8B text cases on CPU with the same checkpoint. Your optimized billing case matches our FP16 run with an explicit inverse to two float32 ULPs. Your optimized login case matches our FP16 run with the triangular solve to two ULPs. Your optimized sales case matches our bf16 run to four ULPs, while our FP16 result is 1.8e-3 away, which is the same 1.8e-3 you report as the largest deviation. Our CPU FP16 output is bit-stable across thread counts, so rounding does not explain this. My guess is that the committed JSON was assembled by hand from the six separate runs and one row carries the wrong label: the script prints one run per process, writes a local path instead of a revision for snapshot, and does not emit the hardware, macOS, dtype or warmups fields. Could you make decider/bench/mps.py record the dtype, the patch_mps() return value and the Transformers version for each run, and commit the script's own output?

3. Concrete changes.

• decider/mps_ops.py: mps_chunk_gated_delta_rule and l2norm are near-verbatim from Transformers 5.17 models/qwen3_5/modeling_qwen3_5.py. The licence is the same, so there is no conflict, but please add a header naming the source, the version and the Apache-2.0 notice.
• tests/test_mps_ops.py: the imported torch_chunk_gated_delta_rule is the fla dispatcher, not the PyTorch reference, because flash-linear-attention is a hard dependency of this project. On a CPU-only Linux machine the two comparison tests raise RuntimeError: 0 active drivers. Please unwrap to the pure-PyTorch function, or skip when the reference cannot run on the test device. Our CI installs without torch and skips the whole file, so it would not catch this.
• pyproject.toml: please give the extra a platform marker, mlx>=0.20; sys_platform == "darwin" and platform_machine == "arm64".
• decider/mps_ops.py:65-82: please assert L.dtype == torch.float32 before the kernel call, because Metal does not bounds-check a buffer of the wrong width. Please also wrap the MLX branch in a try/except that disables the kernel after the first failure, so that a DLPack problem falls through to the PyTorch inverse and does not fail the request.
• decider/infer.py: automatic device selection is welcome, but it changes "no GPU visible" from an exception into a slowdown of about 100 times with no message. Please log the resolved device and dtype once at construction, and update the Decider docstring, which still says bf16 on CUDA.
• decider/vision/model.py: the pixel_values cast appears to do nothing, because Qwen3_5Model.get_image_features already calls pixel_values.type(self.visual.dtype). If it fixed a real problem on MPS, please say which one. Otherwise please remove it. Also, patch_mps() runs there whenever MPS exists, before any device is chosen.
• docs/benchmarks: please add a third arm, reference plus patch_conv() only. patch_conv is device-agnostic and already in the tree, so this would show how much of the 1.4x comes from the new attention path and how much from the existing convolution patch.
• The patch is pinned to Transformers 5.17.x exactly. On 5.18 the path is disabled and a warning is printed at every construction. A version floor plus a check of the patched function's signature would be more durable.

4. One question. With flash-linear-attention installed, which a plain pip install -e . gives you, transformers.models.qwen3_5.modeling_qwen3_5.torch_chunk_gated_delta_rule is bound to fla's Triton kernel at import and not to the PyTorch reference. What does the unpatched MPS path do on your machine: does it run, or does it fail? If it fails, this PR is what makes MPS work at all, not only a speedup, and the PR description and the doc should say so.

I cannot verify the following without a Mac: the DLPack hand-off between PyTorch MPS and MLX, the timings, and the 23 passing tests.

[simply-sunny]

Implemented the requested changes in f01b58d.

1. Evidence: added a real decider.evaluate.py MPS run over all 1,500 held-out MASSIVE Scenario examples at temperature 1.30. Result: accuracy 0.7553, ECE 0.0438, versus the published BF16 row 0.756 / 0.041. Raw output and methodology are in docs/benchmarks/mps-heldout.*. Long context and the aggregate 24-task score remain untested.

2. Benchmark provenance: decider/bench/mps.py now emits checkpoint revision, dtype, patch_mps, patch_conv, PyTorch/Transformers versions, platform, warmups, timings, medians, and probabilities. The committed JSON is generated from those run outputs. It includes reference, conv-only, and optimized arms.

3. Concrete code changes:

   • Added Transformers 5.17 source and Apache-2.0 attribution.
   • Tests use inspect.unwrap() to force the pure-PyTorch reference, even when FLA is installed; 24 tests pass with FLA installed.
   • Metal dependency is restricted to Darwin ARM64.
   • Metal inversion asserts float32 and disables itself after a runtime/DLPack failure, warning before falling back.
   • Decider logs its resolved device/dtype and documents automatic CUDA/MPS/CPU selection.
   • Removed the redundant vision pixel cast and moved MPS patch activation until slot_logits() knows the model device.
   • Compatibility now uses a Transformers >=5.17 floor plus required signature validation; incompatible-version warnings are emitted once.
   • Added a regression test for Metal failure fallback.

4. Unpatched MPS behavior: with flash-linear-attention 0.5.2 installed, the evidence Mac could not import fla.ops because Triton was unavailable. Transformers therefore selected the pure-PyTorch implementation, and the unpatched MPS path ran successfully. The FLA/Triton implementation itself was not tested. The benchmark explicitly unwraps the dispatcher for a deterministic pure-reference arm.

Local verification: 24 passed, git diff --check, and benchmark/module compilation passed. The branch is pushed at f01b58d; the working tree is clean.

[Mapika]

Thanks for the revision. I checked f01b58d against the earlier notes.

Done:

• The held-out run is the evidence I asked for. 0.7553 / 0.0438 on the 1,500 MASSIVE Scenario examples against the published 0.756 / 0.041 is inside run-to-run noise, and the raw output is committed.
• The compatibility check now has a version floor and resolves query and chunk_size positions from the signature instead of fixed indices.
• The Metal path disables itself after a failure, warns once, and has a regression test.
• The benchmark JSON carries provenance.
• On a CUDA machine without MPS, patch_mps() returns False, the modules import, and the CPU-runnable tests pass (5 passed, 1 skipped). The patch is inert off Apple hardware, which is what matters for everyone else using the repo.

Three small things remain. None blocks the merge; fix them here or in a follow-up as you prefer:

1. docs/benchmarks/mps-full-model.md: the 0.8B billing row still shows the optimized arm slower than the reference (0.98x) while the other two 0.8B rows show 1.4 to 1.6x. That row was measured before the kernels warmed. Rerun it or footnote it.
2. fast_invert_unitriangular_64: the float32 assert sits outside the try, so a non-float32 input raises instead of taking the fallback, and assert is removed under python -O. The caller always passes float32 today, so this is a robustness point, not a bug.
3. Decider.__init__ reports the resolved device and dtype through logger.info, which prints nothing unless the caller configures logging. A warnings.warn on the MPS path, or a print behind a verbose flag, would make the automatic device choice visible.

I cannot run the MPS path here, so the Apple-side numbers rest on your run and the committed output. Approving on that basis.