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>
