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)