sparse/config: probe flex_gemm on MPS, soft-fall-back to pytorch backend

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>

test_flex_gemm_integration.py

@@ -0,0 +1,85 @@
+"""End-to-end integration smoke test for SPARSE_CONV_BACKEND=flex_gemm on MPS.
+
+Mirrors what trellis2's sparse decoder does: builds a SparseTensor on MPS,
+runs a SparseConv3d through the flex_gemm backend, and feeds the result
+through a LayerNorm — the exact path that originally crashed with
+"Passed CPU tensor to MPS op" before mtlgemm's device-routing fix.
+
+Exercises both Algorithm.IMPLICIT_GEMM (default dense kernel) and
+Algorithm.MASKED_IMPLICIT_GEMM (the masked kernel that landed in mtlgemm
+round 2). Both should produce numerically equivalent output and pass the
+LayerNorm hand-off without crashing.
+
+Note: this script intentionally uses torch.nn.functional.layer_norm rather
+than trellis2's SparseLayerNorm wrapper, because that wrapper currently calls
+torch.zeros_like which lacks an MPS kernel in some PyTorch builds. That is a
+PyTorch issue, separate from the mtlgemm fix being verified here.
+"""
+
+import os
+os.environ.setdefault("SPARSE_CONV_BACKEND", "flex_gemm")
+os.environ.setdefault("SPARSE_ATTN_BACKEND", "sdpa")
+os.environ.setdefault("FLEX_GEMM_QUIET", "1")
+
+import torch
+
+assert torch.backends.mps.is_available(), "This test needs MPS"
+
+from trellis2.modules.sparse import SparseTensor, SparseConv3d
+from flex_gemm.ops.spconv import Algorithm, set_algorithm
+
+device = "mps"
+dtype = torch.float16
+
+# Build a small sparse voxel shell (mimics trellis2 decoder input scale)
+res = 16
+ch = 32
+coords = torch.stack(torch.meshgrid(
+    torch.arange(res), torch.arange(res), torch.arange(res), indexing="ij",
+), dim=-1).int().contiguous()
+dist = ((coords.float() - res / 2 + 0.5) ** 2).sum(dim=-1).sqrt()
+active = (dist <= res / 2) & (dist >= res / 2 - 1.25)
+coords = torch.nonzero(active).int()
+coords = torch.cat([torch.zeros(coords.shape[0], 1, dtype=torch.int32), coords], dim=-1)
+coords = coords.contiguous().to(device)
+feats = torch.randn(coords.shape[0], ch, dtype=dtype).to(device).contiguous()
+
+print(f"Sparse tensor: {feats.shape[0]} voxels, {ch} channels, dtype={dtype}, device={device}")
+
+x = SparseTensor(feats=feats, coords=coords, shape=torch.Size([1, ch]), spatial_shape=[res, res, res])
+print(f"Built SparseTensor: feats device={x.feats.device}, dtype={x.feats.dtype}")
+
+# Build the conv module on CPU then move to MPS to dodge LayerNorm-init MPS limits.
+conv = SparseConv3d(ch, ch, kernel_size=3, bias=True).to(dtype)
+conv.weight.data = conv.weight.data.to(device)
+if conv.bias is not None:
+    conv.bias.data = conv.bias.data.to(device)
+print(f"Conv weight device: {conv.weight.device}, dtype: {conv.weight.dtype}")
+
+def _run_with_algo(algo, label):
+    set_algorithm(algo)
+    out = conv(x)
+    assert out.feats.device.type == "mps", f"FAIL [{label}]: SparseConv3d on {out.feats.device}, expected mps"
+    assert out.feats.dtype == dtype, f"FAIL [{label}]: SparseConv3d dtype {out.feats.dtype}, expected {dtype}"
+    # LayerNorm hand-off — the original crash site
+    w = torch.ones(ch, dtype=dtype).to(device)
+    b = torch.zeros(ch, dtype=dtype).to(device)
+    z = torch.nn.functional.layer_norm(out.feats, (ch,), w, b)
+    assert z.device.type == "mps", f"FAIL [{label}]: LayerNorm on {z.device}"
+    total = z.sum().item()
+    print(f"  {label:30s} feats={tuple(out.feats.shape)} sum={total:.4f}")
+    return out.feats.detach().cpu().float()
+
+print()
+y_dense  = _run_with_algo(Algorithm.IMPLICIT_GEMM, "IMPLICIT_GEMM (dense)")
+y_masked = _run_with_algo(Algorithm.MASKED_IMPLICIT_GEMM, "MASKED_IMPLICIT_GEMM")
+
+# Numerical parity between the two algorithms — same inputs, equivalent output.
+diff = (y_dense - y_masked).abs().max().item()
+parity_tol = 2e-2  # fp16 — masked reduces in a different order
+assert diff <= parity_tol, f"FAIL: dense vs masked diff {diff:.4e} > tol {parity_tol:.4e}"
+print(f"  parity dense vs masked: max_diff={diff:.4e} (tol={parity_tol})")
+
+print()
+print("PASS — trellis2-apple SparseConv3d + LayerNorm runs end-to-end on MPS via flex_gemm")
+print("       Both IMPLICIT_GEMM and MASKED_IMPLICIT_GEMM produce equivalent output.")

trellis2/modules/sparse/config.py

@@ -11,15 +11,36 @@ def __detect_defaults():
     global CONV, ATTN
     if platform.system() == 'Darwin':
         ATTN = 'sdpa'
-        try:
-            import flex_gemm
+        if __flex_gemm_works_on_mps():
             CONV = 'flex_gemm'
-        except ImportError:
+        else:
             CONV = 'pytorch'
     elif not __has_cuda():
         CONV = 'pytorch'
         ATTN = 'sdpa'
 
+
+def __flex_gemm_works_on_mps():
+    """Probe flex_gemm with a tiny MPS conv. If the install pre-dates the
+    device-routing fix (or fails to build), it returns a CPU tensor — fall
+    back to the pure-PyTorch backend rather than crashing inside the model
+    on the first LayerNorm. Build tensors on CPU and move to MPS because some
+    PyTorch builds lack int/fp16 torch.zeros kernels on MPS."""
+    try:
+        import torch
+        if not torch.backends.mps.is_available():
+            return False
+        import flex_gemm
+        from flex_gemm.ops.spconv import sparse_submanifold_conv3d, Algorithm, set_algorithm
+        set_algorithm(Algorithm.IMPLICIT_GEMM)
+        coords = torch.tensor([[0, 0, 0, 0]], dtype=torch.int32).to('mps')
+        feats = torch.zeros((1, 4), dtype=torch.float16).to('mps')
+        weight = torch.zeros((4, 1, 1, 1, 4), dtype=torch.float16).to('mps')
+        out, _ = sparse_submanifold_conv3d(feats, coords, torch.Size([1, 4, 1, 1, 1]), weight)
+        return out.device.type == 'mps'
+    except Exception:
+        return False
+
 def __has_cuda():
     try:
         import torch