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