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>
