From 2780bc459abbef594c812ff9471ae70288145b03 Mon Sep 17 00:00:00 2001 From: Pedro Augusto <16762743+pedronaugusto@users.noreply.github.com> Date: Tue, 21 Apr 2026 03:04:22 +0100 Subject: [PATCH] sparse/config: probe flex_gemm on MPS, soft-fall-back to pytorch backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- test_flex_gemm_integration.py | 85 +++++++++++++++++++++++++++++++ trellis2/modules/sparse/config.py | 27 ++++++++-- 2 files changed, 109 insertions(+), 3 deletions(-) create mode 100644 test_flex_gemm_integration.py diff --git a/test_flex_gemm_integration.py b/test_flex_gemm_integration.py new file mode 100644 index 0000000..281a1cc --- /dev/null +++ b/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.") diff --git a/trellis2/modules/sparse/config.py b/trellis2/modules/sparse/config.py index f6d1070..f298f89 100644 --- a/trellis2/modules/sparse/config.py +++ b/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