spconv Metal kernel (opt-in) + fix vertex baker scheduler fallthrough

- 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 <noreply@anthropic.com>

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.

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

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)

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)