7860148eb2
- 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>
74 lines
2.7 KiB
Python
74 lines
2.7 KiB
Python
"""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
|