monster/trellis-2-mrp-mlx
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

sparse/config: probe flex_gemm masked path too, not just dense

Browse Source 
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>
This commit is contained in:
Pedro Augusto 2026-04-21 14:54:27 +01:00 committed by Jourloy
parent 2780bc459a
commit 71f071fc76
1 changed files with 19 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
trellis2/modules/sparse/config.py
View File

@ -21,23 +21,34 @@ def __detect_defaults():
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."""
"""Probe flex_gemm with tiny MPS convs covering both IMPLICIT_GEMM and
MASKED_IMPLICIT_GEMM. If the install pre-dates the device-routing fix
(or pre-dates the real masked kernel in round 2), one of these 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)
# Exercise both algorithms — masked carries its own cache/dispatch path
# distinct from dense. A stale install may have one working and the
# other broken (e.g. the pre-round-2 aliased-to-dense fallback).
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'
shape = torch.Size([1, 4, 1, 1, 1])
for algo in (Algorithm.IMPLICIT_GEMM, Algorithm.MASKED_IMPLICIT_GEMM):
set_algorithm(algo)
out, _ = sparse_submanifold_conv3d(feats, coords, shape, weight)
if out.device.type != 'mps':
return False
return True
except Exception:
return False