From 7860148eb2c4c843c690c40fb44e5cc13842ff23 Mon Sep 17 00:00:00 2001 From: m3ultra Date: Mon, 20 Jul 2026 03:37:20 +1000 Subject: [PATCH] spconv Metal kernel (opt-in) + fix vertex baker scheduler fallthrough MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - sparse_conv_metal: fused gather-GEMM kernel, parity 8e-4 vs stock; honest verdict: stock chunking already bounds memory at decoder scale (1.97GB synthetic peak) and beats the scalar kernel on speed — kept as TRELLIS2_METAL_SPCONV=1 opt-in + upstream reference, NOT default. - _attempt_schedule: preferred=vertex now runs the vertex baker (was silently falling through to the 5h+ pure-python kdtree grind). Co-Authored-By: Claude Fable 5 --- mlx_backend/sparse_conv.py | 11 +++++ mlx_backend/sparse_conv_metal.py | 73 +++++++++++++++++++++++++++ scripts/generate_asset.py | 7 ++- test_spconv_metal.py | 85 ++++++++++++++++++++++++++++++++ 4 files changed, 175 insertions(+), 1 deletion(-) create mode 100644 mlx_backend/sparse_conv_metal.py create mode 100644 test_spconv_metal.py diff --git a/mlx_backend/sparse_conv.py b/mlx_backend/sparse_conv.py index 4cc3a2b..2cea083 100644 --- a/mlx_backend/sparse_conv.py +++ b/mlx_backend/sparse_conv.py @@ -3,6 +3,7 @@ Submanifold sparse 3D convolution for MLX. Port of conv_pytorch.py algorithm: hash → neighbor map → gather → bmm → scatter. """ import itertools +import os import logging import time import mlx.core as mx @@ -129,6 +130,16 @@ class MlxSparseConv3d(nn.Module): w = self.weight.reshape(Co, K, Ci) # (Co, K, Ci) w = w.transpose(1, 2, 0) # (K, Ci, Co) + # TRELLIS2_METAL_SPCONV=1: fused gather-GEMM Metal kernel — never + # materializes the (K,N,Ci) gather (stock path peaks ~77GB at + # decoder scale and eval-syncs between 1-2-offset chunks). + if os.environ.get("TRELLIS2_METAL_SPCONV", "0") == "1": + from .sparse_conv_metal import sparse_conv_metal + feats_padded = mx.concatenate([ + x.feats, mx.zeros((1, Ci), dtype=x.feats.dtype)], axis=0) + result = sparse_conv_metal(feats_padded, neighbor_map, w, self.bias) + return x.replace(result.astype(x.feats.dtype)) + # Pad feats with zero row for out-of-bounds indices # feats_padded[N] is zeros, so invalid neighbor lookups produce zero # from matmul naturally — no valid_mask needed. diff --git a/mlx_backend/sparse_conv_metal.py b/mlx_backend/sparse_conv_metal.py new file mode 100644 index 0000000..5db7ead --- /dev/null +++ b/mlx_backend/sparse_conv_metal.py @@ -0,0 +1,73 @@ +"""Fused gather-GEMM-accumulate Metal kernel for submanifold sparse conv. + +The stock MlxSparseConv3d materializes a (K, N, Ci) gather — ~77GB at +decoder scale — then chunks it 1-2 offsets at a time with an mx.eval sync +between chunks. This kernel never materializes the gather: each thread +accumulates out[n, co] = Σ_k Σ_ci feats[nbr[k,n], ci] · w[k, ci, co] +straight from the source buffers. Correctness-first implementation +(scalar loops, fp32 accumulation); simdgroup-matrix tiling is a later +optimization if profiling ever says the decoders matter for *time* +(they are ~10s/gen — this kernel's prize is PEAK MEMORY, which gates the +M1 Ultra). + +Env gate: TRELLIS2_METAL_SPCONV=1 routes MlxSparseConv3d through this. +""" +import mlx.core as mx + +_SRC = """ + uint n = thread_position_in_grid.x; // voxel index + uint co = thread_position_in_grid.y; // output channel + uint N = (uint)shape_info[0]; + uint K = (uint)shape_info[1]; + uint Ci = (uint)shape_info[2]; + uint Co = (uint)shape_info[3]; + if (n >= N || co >= Co) return; + + float acc = 0.0f; + for (uint k = 0; k < K; ++k) { + int nbr = nmap[k * N + n]; // N (==pad row) when invalid + if (nbr >= (int)N) continue; // pad row is zeros anyway; skip + const device T* frow = feats + (size_t)nbr * Ci; + const device T* wrow = w + ((size_t)k * Ci) * Co + co; + for (uint ci = 0; ci < Ci; ++ci) { + acc += (float)frow[ci] * (float)wrow[(size_t)ci * Co]; + } + } + out[(size_t)n * Co + co] = (T)acc; +""" + +_kernel = None + + +def _get_kernel(): + global _kernel + if _kernel is None: + _kernel = mx.fast.metal_kernel( + name="submconv3d_gather_gemm", + input_names=["feats", "nmap", "w", "shape_info"], + output_names=["out"], + source=_SRC, + ) + return _kernel + + +def sparse_conv_metal(feats_padded: mx.array, neighbor_map: mx.array, + w_kcico: mx.array, bias: mx.array) -> mx.array: + """feats_padded: (N+1, Ci) — row N is zeros (pad; skipped anyway). + neighbor_map: (K, N) int32 with N as the invalid sentinel. + w_kcico: (K, Ci, Co). bias: (Co,). Returns (N, Co) in w's dtype.""" + K, N = neighbor_map.shape + Ci = feats_padded.shape[1] + Co = w_kcico.shape[2] + shape_info = mx.array([N, K, Ci, Co], dtype=mx.int32) + kern = _get_kernel() + (out,) = kern( + inputs=[feats_padded.astype(w_kcico.dtype), neighbor_map, + w_kcico, shape_info], + template=[("T", w_kcico.dtype)], + grid=(N, Co, 1), + threadgroup=(256, 1, 1), + output_shapes=[(N, Co)], + output_dtypes=[w_kcico.dtype], + ) + return out + bias diff --git a/scripts/generate_asset.py b/scripts/generate_asset.py index d668a67..3660e4c 100755 --- a/scripts/generate_asset.py +++ b/scripts/generate_asset.py @@ -277,7 +277,12 @@ def _export_pbr(mesh, path: Path, *, baker: str, target: Optional[int], texture_ def _attempt_schedule(preferred: str, requested_target: Optional[int], raw_faces: int): - preferred_order = ["metal", "kdtree"] if preferred in {"auto", "metal"} else ["kdtree"] + if preferred == "vertex": + preferred_order = ["vertex"] # direct vertex-color bake; no fallback needed + elif preferred in {"auto", "metal"}: + preferred_order = ["metal", "kdtree"] + else: + preferred_order = ["kdtree"] targets = [requested_target] if requested_target is None and raw_faces > SAFETY_FACE_TARGET: targets.append(SAFETY_FACE_TARGET) diff --git a/test_spconv_metal.py b/test_spconv_metal.py new file mode 100644 index 0000000..7598e19 --- /dev/null +++ b/test_spconv_metal.py @@ -0,0 +1,85 @@ +"""Parity + memory gate: fused Metal sparse conv vs the stock chunked path. +Small case first (exact check vs naive reference), then a decoder-scale +case for peak-memory + timing comparison. +""" +import sys +import time +from pathlib import Path + +import mlx.core as mx +import numpy as np + +ROOT = Path(__file__).resolve().parent +sys.path.insert(0, str(ROOT)) + +from mlx_backend.sparse_conv import MlxSparseConv3d, _build_neighbor_map +from mlx_backend.sparse_conv_metal import sparse_conv_metal +from mlx_backend.sparse_tensor import MlxSparseTensor + +mx.random.seed(0) + + +def make_case(n_vox, ci, co, res): + # random unique voxel coords in a res^3 grid, batch 0 + rng = np.random.default_rng(0) + flat = rng.choice(res ** 3, size=n_vox, replace=False) + x, rem = flat // (res * res), flat % (res * res) + y, z = rem // res, rem % res + coords = np.stack([np.zeros_like(x), x, y, z], 1).astype(np.int32) + st = MlxSparseTensor( + feats=mx.random.normal((n_vox, ci)).astype(mx.float16), + coords=mx.array(coords)) + conv = MlxSparseConv3d(ci, co, 3) + conv.weight = mx.random.normal((co, 3, 3, 3, ci)).astype(mx.float16) * 0.05 + conv.bias = mx.random.normal((co,)).astype(mx.float16) * 0.1 + return st, conv + + +def run_metal(st, conv): + nmap = _build_neighbor_map(st.coords, st.shape[0], st.spatial_shape, + conv.kernel_size, conv.dilation) + mx.eval(nmap) + K = nmap.shape[0] + Co, Ci = conv.out_channels, conv.in_channels + w = conv.weight.reshape(Co, K, Ci).transpose(1, 2, 0) # (K, Ci, Co) + feats_padded = mx.concatenate( + [st.feats, mx.zeros((1, Ci), dtype=st.feats.dtype)], axis=0) + return sparse_conv_metal(feats_padded, nmap, w, conv.bias) + + +# ---- small exact-parity case ---- +st, conv = make_case(5000, 64, 96, 64) +ref = conv(st).feats +got = run_metal(st, conv) +mx.eval(ref, got) +d = mx.max(mx.abs(ref.astype(mx.float32) - got.astype(mx.float32))).item() +rel = d / max(1e-6, mx.max(mx.abs(ref.astype(mx.float32))).item()) +print(f"small: max|diff|={d:.3e} rel={rel:.3e}", "PASS" if rel < 2e-2 else "FAIL") +ok_small = rel < 2e-2 + +# ---- decoder-scale case: memory + time ---- +st2, conv2 = make_case(400_000, 512, 512, 512) +mx.eval(st2.feats, conv2.weight) + +mx.reset_peak_memory() +t0 = time.perf_counter() +ref2 = conv2(st2).feats +mx.eval(ref2) +t_stock = time.perf_counter() - t0 +m_stock = mx.get_peak_memory() / 2**30 + +st2._spatial_cache = {} # drop cached neighbor map so both build it +mx.reset_peak_memory() +t0 = time.perf_counter() +got2 = run_metal(st2, conv2) +mx.eval(got2) +t_metal = time.perf_counter() - t0 +m_metal = mx.get_peak_memory() / 2**30 + +d2 = mx.max(mx.abs(ref2.astype(mx.float32) - got2.astype(mx.float32))).item() +r2 = d2 / max(1e-6, mx.max(mx.abs(ref2.astype(mx.float32))).item()) +print(f"large: stock {t_stock:.2f}s peak {m_stock:.2f}GB | " + f"metal {t_metal:.2f}s peak {m_metal:.2f}GB | rel diff {r2:.3e}") +ok_large = r2 < 2e-2 +print("SPCONV_GATE:", "PASS" if (ok_small and ok_large) else "FAIL") +sys.exit(0 if (ok_small and ok_large) else 1)