Fix/speculative shape invariance

[aleroot]

Proposed changes

Related issue: ml-explore/mlx-swift-lm#542

Opt-in batch-invariant execution mode for selected Metal inference kernels. It is intended for workloads such as greedy speculative decoding, where evaluating a token independently and evaluating the same token inside a short verification block must produce identical results.

Problem

Autoregressive decoding normally evaluates one token at a time, while speculative decoding verifies several tokens in one target-model invocation.

Although these computations are mathematically equivalent, their shapes currently select different Metal execution paths:

• A single matrix row uses GEMV, while a short matrix can use gemv_wide or GEMM.
• Two-pass SDPA derives its split count from the query length, changing the softmax reduction tree between single-query decoding and multi-query verification.

Floating-point reductions are not associative. Consequently, the different kernels can produce slightly different BF16 results for the same logical token. Across many transformer layers and decoding rounds, these differences can eventually change the ordering of two close logits and cause greedy argmax decoding to diverge.

This is especially visible with quantized models and long contexts, but the underlying problem is shape-dependent floating-point reduction order rather than quantized matrix multiplication itself.

Changes

This PR introduces MLX_METAL_BATCH_INVARIANT_LIMIT, with matching runtime APIs:

mx.metal.set_batch_invariant_limit(limit)
mx.metal.get_batch_invariant_limit()

When the relevant leading/query dimension is within the configured limit:

• Short matrix rows are executed as a parallel batch of canonical single-row GEMVs. Rows remain parallel, but each row uses the same reduction kernel and order as token-by-token inference.
• Two-pass SDPA uses the single-query split configuration, preserving the same softmax reduction layout for a query evaluated alone or inside a short causal block.

The setting is process-wide and stored atomically. It defaults to 0, meaning disabled, so existing execution and performance are unchanged unless an application explicitly opts in. Long prefills and dimensions above the configured limit continue using the existing optimized paths.

The mode is intentionally opt-in because selecting canonical reduction kernels can reduce performance. Applications only need to cover their maximum speculative verification width.

Checklist

Put an x in the boxes that apply.

• I have read the [CONTRIBUTING](https://github.com/ml-explore/mlx/blob/main/CONTRIBUTING.md) document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[eyupcanakman]

Ran this on an M5 Pro (macOS 27.0) at dcedd1261, checking the half the tests here do not cover, whether the paths actually diverge with the limit off.

They do. Quantized matmul, over the grid your third test uses (K=4096, N=512, M=2..4, transposed):

limit=0 :  9 of 12 configs diverge, 9/9 rows each
           (all six f32, plus bf16 affine at gs 32/64/128)
limit=4 : 12 of 12 configs, 0/9 rows

SDPA, your first test's scenario with Lk in {1028, 8196, 16388, 65540}:

limit=0 : 32/32 rows mismatch
limit=8 :  0/32

Dense matmul at your M=4 shape fails f16, f32 and complex64 at limit=0, and passes all four dtypes at limit=4.

bf16 is already invariant at the narrower shapes on this chip. It passes with the feature off in the dense case up to M=9 and in mxfp4, mxfp8 and nvfp4, so a bf16-only check misses most of it.

Cost when enabled, median of 11 alternating A/B reps at K=N=4096 bf16, M=1 kept as a control since the gate needs M>1:

M=1  +2.6% / +1.7%   (control, noise floor)
M=2  +7.1% / +23.1%
M=4 +19.6% / +12.5%
M=8 +20.8% / +23.7%

Positive every run, but the magnitude moves around, so 10-25% at these shapes rather than a single number.

The matmul gate needs b_transposed, so a natural-layout B still diverges inside the limit (M=7 and M=8 here). Looks like that is on purpose, flagging in case not.

[aleroot]

@eyupcanakman Thanks for your tests. The natural-layout B exclusion wasn’t intentional. I’ve fixed the canonical GEMV dispatch to support both B orientations and added coverage test. It reproduces before the fix and passes afterward on my M2. I do not have an M5 unfortunately 😀 a rerun on your M5 would be appreciated.

[zcbenz]

It was brought up several times before about providing a way to get deterministic results with different batch size, but we do not plan to provide any of such support. It would complicate kernel dispatching a lot, and anyone needing that would definitely turn to other specialized inference engine in the end.