Add Metal SDPA support for D96/V64

[wyanzhao]

Summary

Add Metal scaled dot-product attention support for Dqk=96, Dv=64, the
asymmetric head dimensions used by MiniCPM3. This adds one-pass and two-pass
vector kernels, lets NAX use an independent value tile width, and registers the
new kernels in JIT and no-JIT builds. Tests cover one- and two-pass MHA/GQA,
masks, sinks, non-contiguous inputs, interleaved value tensors, and the existing
D=96, V=96 path.

The existing short-query rules select the vector path. Full attention requires
NAX, and float32 also requires TF32; other cases keep the existing fallback.

MiniCPM3 performance

Measured on an M5 Max with
mlx-community/MiniCPM3-4B-4bit
snapshot 7a5f80ff3161be21f0a90e9be7137495633b2b99. The baseline kept both
new D96/V64 kernels off. The candidate enabled NAX for prefill, Vector for
decode, or both for generation. Ratio is baseline time divided by candidate
time; each row uses nine alternating adjacent pairs.

Workload	Prompt	Comparison	Median baseline	Median candidate	Paired ratio (95% CI)	Latency reduction
NAX prefill TTFT	128	Vector off, NAX off → Vector off, NAX on	57.548 ms	56.299 ms	1.0279 [1.0084, 1.0478]	2.71%
NAX prefill TTFT	512	Vector off, NAX off → Vector off, NAX on	152.703 ms	139.609 ms	1.0937 [1.0894, 1.0981]	8.57%
NAX prefill TTFT	1536	Vector off, NAX off → Vector off, NAX on	485.327 ms	415.369 ms	1.1672 [1.1633, 1.1710]	14.32%
Vector decode	128	Vector off, NAX off → Vector on, NAX off	12.603 ms/step	11.757 ms/step	1.0660 [1.0511, 1.0811]	6.19%
Vector decode	512	Vector off, NAX off → Vector on, NAX off	11.760 ms/step	11.120 ms/step	1.0590 [1.0474, 1.0707]	5.57%
Vector decode	1536	Vector off, NAX off → Vector on, NAX off	16.874 ms/step	15.925 ms/step	1.05374 [1.04287, 1.06472]	5.10%
16-token generation	128	Vector off, NAX off → Vector on, NAX on	239.641 ms	224.417 ms	1.06981 [1.05403, 1.08582]	6.53%
16-token generation	512	Vector off, NAX off → Vector on, NAX on	349.347 ms	327.835 ms	1.0653 [1.0572, 1.0734]	6.13%
16-token generation	1536	Vector off, NAX off → Vector on, NAX on	852.363 ms	752.896 ms	1.1263 [1.1178, 1.1349]	11.21%

NAX prefill measures materialized first-token logits with 512-token chunking.
Vector decode starts both configurations from the same baseline KV cache and
consumes the same eight tokens. Generation uses a fresh cache and includes
generation and detokenization, but not input tokenization.

Validation

Route tracing confirmed the NAX kernel during prefill and the vector kernel
during decode; configurations with those kernels off used the existing path.
First-token and decode logits stayed within the preset tolerance, top-1
predictions matched, and both configurations generated the same 16 tokens.

• Focused JIT and no-JIT SDPA tests passed. With TF32 disabled, the no-JIT
  test_fast_sdpa module passed 28 tests and skipped 2 by platform guards.
• The kernel-off configuration matched the pre-change baseline byte-for-byte
  on all six exact operator cells.
• pre-commit run --all-files passed.

Quick reproduction

Build main and this PR in separate clean checkouts and virtual environments,
then install MLX and mlx-lm in each:

python -m pip install -e .
python -m pip install mlx-lm

Save the script below as bench_d96_v64.py. Run the same command with the
Python executable from each checkout; main is the baseline and this PR is the
candidate.

import argparse
import json
import statistics
import time

import mlx.core as mx
from mlx_lm import load, stream_generate
from mlx_lm.generate import generation_stream
from mlx_lm.models.cache import make_prompt_cache

MODEL = "mlx-community/MiniCPM3-4B-4bit"
REVISION = "7a5f80ff3161be21f0a90e9be7137495633b2b99"
MOTIF = [
    53400, 1980, 1348, 6035, 11989, 2434, 72, 7591, 4104, 8381, 72,
    32435, 59415, 28262, 38531, 4408, 7627, 3120, 10150, 2427, 66,
]

p = argparse.ArgumentParser()
p.add_argument("kind", choices=["prefill", "decode", "generation"])
p.add_argument("length", type=int, choices=[128, 512, 1536])
a = p.parse_args()

mx.set_default_device(mx.gpu)
mx.random.seed(0)
model, tokenizer = load(MODEL, revision=REVISION)
tokens = ([1] + MOTIF * (a.length // len(MOTIF) + 1))[:a.length]
prompt = mx.array(tokens)
continuation = [MOTIF[(a.length - 1 + i) % len(MOTIF)] for i in range(8)]
baseline_cache = None

if a.kind == "decode":
    cache = make_prompt_cache(model)
    for cursor in range(0, a.length, 512):
        logits = model(prompt[None, cursor:cursor + 512], cache=cache)
        mx.eval(logits, [c.state for c in cache])
    baseline_cache = (
        [c.state for c in cache],
        [c.meta_state for c in cache],
        [type(c) for c in cache],
    )
    mx.eval(baseline_cache[0])


def run():
    start = time.perf_counter_ns()
    if a.kind == "prefill":
        cache = make_prompt_cache(model)
        cursor = 0
        with mx.stream(generation_stream):
            while cursor < a.length - 1:
                count = min(512, a.length - 1 - cursor)
                logits = model(prompt[None, cursor:cursor + count], cache=cache)
                mx.eval([c.state for c in cache])
                cursor += count
            logits = model(prompt[None, -1:], cache=cache)[:, -1, :]
        mx.eval(logits)
        elapsed = (time.perf_counter_ns() - start) / 1e6
    elif a.kind == "decode":
        states, metas, classes = baseline_cache
        cache = [cls.from_state(s, m) for cls, s, m in zip(classes, states, metas)]
        with mx.stream(generation_stream):
            for token in continuation:
                logits = model(mx.array([[token]]), cache=cache)[:, -1, :]
                mx.async_eval(logits)
        mx.synchronize()
        elapsed = (time.perf_counter_ns() - start) / 1e6 / len(continuation)
    else:
        rows = list(stream_generate(
            model, tokenizer, prompt, max_tokens=16, prefill_step_size=512,
            sampler=lambda x: mx.argmax(x, axis=-1),
        ))
        mx.synchronize()
        elapsed = (time.perf_counter_ns() - start) / 1e6
        assert len(rows) == 16
    mx.clear_cache()
    return elapsed


for _ in range(5):
    run()
samples = [run() for _ in range(9)]
unit = "ms/step" if a.kind == "decode" else "ms"
print(json.dumps({"median": statistics.median(samples), "unit": unit,
                  "samples": samples}))

Run all nine cells in both environments with no other GPU work active:

for kind in prefill decode generation; do
  for length in 128 512 1536; do
    python bench_d96_v64.py "$kind" "$length"
  done
done

This is a quick workload comparison; results will vary with the machine and
build.

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: AI-assisted coding and testing were used; I reviewed
  the diff and verified the results above.