The mtlgemm Metal extension's metal_context.mm calls
at::mps::dispatch_sync_with_rethrow, which was added to that namespace
in PyTorch 2.11.0 (pytorch/pytorch#167445, merged 2025-11-11). Earlier
stable releases (2.6 – 2.10) only expose it under at::native::mps::,
so installing mtlgemm against them fails the C++ extension build.
torchvision pinned to >=0.26.0 (matches torch 2.11) to keep the
torch/torchvision wheel pair ABI-compatible on resolution.
The flex_gemm_sparse_attn backend is now flash-attention-v2 with
simdgroup_matrix_multiply_accumulate for Q@K^T and P@V plus simd-shuffle
softmax row reductions, and wins 5–15× over SDPA-padded-CPU-bounce at
every measured shape including max_seqlen=2048. Production decoder block
on M3 Max: 5.04× wall-clock vs all-SDPA-padded baseline (3.49 ms vs
17.57 ms), entirely from this change being the default.
The existing __flex_gemm_works_on_mps() probe covers the same package
the attention path lives in, so the gate is identical: if conv probe
passes, set both CONV='flex_gemm' and ATTN='flex_gemm_sparse_attn'.
SPARSE_ATTN_BACKEND= (or ATTN_BACKEND=) env override is unchanged.
Also fix benchmarks/e2e_decoder.py to inject the repo root into
sys.path so it runs as `python benchmarks/e2e_decoder.py` from any cwd.
End-to-end micro-bench for the production decoder block at res=32 ch=64
with seqlens=[256, 192, 128, 64]. Three stages reported (convs-only /
attn-only / combined block) at fp16, vs the all-SDPA-padded baseline.
Used to track the cumulative effect of the mtlgemm flash-attention-v2
fwd + bwd work on the actual decoder hot path.
The 256-cap inside the flex_gemm_sparse_attn branch was a hold-over from
the early naive Metal kernel (round 3). The current backend is flash-
attention-v2 with simdgroup matmul + simd-shuffle softmax row reductions
and wins at every measured shape including max_seqlen=2048. The CUDA
backends above (xformers / flash_attn / flash_attn_3) never fork on
max_seqlen — match that precedent.
Safety-valve preserved as FLEX_GEMM_ATTN_MAX_SEQLEN=N env var: when set,
falls back to SDPA-padded above the cap. Useful only on PyTorch builds
where the Accelerate-SDPA-CPU-bounce happens to win at a specific shape
(measured crossover sits beyond 768 on fp32, higher on fp16).
SparseGroupNorm / SparseLayerNorm used torch.zeros_like which fails with
"DispatchStub: missing kernel for mps" on PyTorch builds compiled with
both CUDA and MPS backends (the user's local build hit this). Added a
_zeros_like_safe helper that builds zeros on CPU and transfers when the
reference tensor is on MPS; Apple Silicon unified memory makes the
transfer metadata-only, so the overhead vs a working MPS zeros kernel
is negligible. On CPU, behaves identically to torch.zeros_like.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Wires the new mtlgemm fused sparse attention kernel through the ATTN
backend selector. Dispatches to the fused Metal kernel when
max(max_q, max_kv) <= 256 (where the naive per-thread-serial-KV kernel
beats SDPA-padded on M3 Max), and falls through to an inline SDPA-padded
path for larger max_seqlen. Opt in via ATTN_BACKEND=flex_gemm_sparse_attn
or SPARSE_ATTN_BACKEND=flex_gemm_sparse_attn. Default on Darwin stays
'sdpa' — the threshold-based fallback doesn't yet prove a universal
win across the pipeline's attention shape distribution.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The original probe only exercised Algorithm.IMPLICIT_GEMM. A stale install
where dense works but the masked cache/dispatch path is broken (e.g. the
pre-round-2 aliased-to-dense fallback) would select flex_gemm and crash at
the first MASKED_IMPLICIT_GEMM call inside the decoder. Probe both.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The default Darwin path used to be a try/except ImportError, which only
catches build failures. With the mtlgemm round 1 fixes shipped, the more
common failure mode for end users will be running an *older* mtlgemm that
still returns CPU tensors from MPS calls — that doesn't fail import, it
fails with a cryptic LayerNorm crash inside the model on the first conv.
The new probe runs a tiny SparseConv3d on MPS and checks the output
device. If anything breaks (import, build, dispatch, return-device), fall
back to the pure-PyTorch backend rather than crashing inside the model.
Tensors in the probe are built on CPU then moved to MPS because some
PyTorch builds lack int/fp16 torch.zeros kernels on MPS — that's a
PyTorch issue, separate from anything we control here.
Plus: end-to-end smoke test (test_flex_gemm_integration.py) that
exercises both Algorithm.IMPLICIT_GEMM and Algorithm.MASKED_IMPLICIT_GEMM
through a real SparseConv3d → F.layer_norm chain on MPS. Confirms both
algorithms return MPS tensors of the right dtype and produce equivalent
output.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
