diff --git a/README.md b/README.md index aa824b4..5690f6f 100644 --- a/README.md +++ b/README.md @@ -26,9 +26,13 @@ defaults `stride=1, padding=None` — which upstream dispatches to `SubMConv3d`. So the whole blocker is one operation: **submanifold 3×3×3 convolution**. +The sparse core now lives in its own package, **`trellis_sparse_mlx`**, because Pixal3D +and the rest of the TRELLIS.2 family need exactly the same thing. This repo is the LATO.2 +model on top of it. + | Upstream | Here | |---|---| -| `spconv.SubMConv3d` | `lato_mlx/sparse/conv.py` — gather/scatter over a sorted-key indice map | +| `spconv.SubMConv3d` | `trellis_sparse_mlx` — gather/scatter over a sorted-key indice map | | `SparseInverseConv3d` | never instantiated upstream; not needed | | strided sparse conv | never used; upsampling is `SparseSubdivide` (coord expansion ×8) | | `nn.Conv3d` (voxel encoder) | dense, maps to `mlx.nn.Conv3d` | diff --git a/bench/run_fleet.sh b/bench/run_fleet.sh index d77882a..2302c15 100755 --- a/bench/run_fleet.sh +++ b/bench/run_fleet.sh @@ -11,7 +11,7 @@ SSH_STR="ssh -o ConnectTimeout=10 -o BatchMode=yes -o StrictHostKeyChecking=acce HOSTS=( "m1max@100.92.78.24:m1max" "m2max@100.120.83.110:m2max" - "100.69.21.128:m4pro" + "m4pro@100.69.21.128:m4pro" "johnking@100.91.239.7:m1ultra" ) diff --git a/lato_mlx/models/vvae.py b/lato_mlx/models/vvae.py index fbc055a..5065786 100644 --- a/lato_mlx/models/vvae.py +++ b/lato_mlx/models/vvae.py @@ -28,13 +28,14 @@ from typing import Optional import mlx.core as mx import mlx.nn as nn -from ..sparse.ops import ( +from trellis_sparse_mlx import ( LayerNorm32, SparseLinear, SparseResBlock, + SparseTensor, SparseTransformerBlock, + downsample, ) -from ..sparse.tensor import SparseTensor, downsample class DownResBlock(nn.Module): diff --git a/lato_mlx/sparse/__init__.py b/lato_mlx/sparse/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/lato_mlx/sparse/conv.py b/lato_mlx/sparse/conv.py deleted file mode 100644 index 224404a..0000000 --- a/lato_mlx/sparse/conv.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Submanifold sparse 3D convolution in pure MLX. - -This is the ONLY genuinely CUDA-locked operation in LATO.2's inference path. Upstream -routes it to `spconv.SubMConv3d` (or torchsparse); neither has a Metal build, which is -what has kept LATO.2 — and TRELLIS before it — off Apple Silicon. - -Every SparseConv3d in the LATO.2 model code is constructed with the defaults -`stride=1, padding=None`, which upstream dispatches to SubMConv3d. So only the -submanifold case is needed, at kernel sizes 3 and 1. - -Submanifold semantics: the output occupies EXACTLY the input coordinates (no dilation of -the occupied set). For each output voxel c: - - out[c] = bias + sum over kernel offsets d of W[d] @ feats[c + d] (c+d occupied) - -Absent neighbours contribute nothing. Implemented by gathering into a feature matrix with -one appended zero row, so "missing" is index N and needs no masking in the hot loop. - -Building the indice map is the expensive part and depends only on the coordinate set, so -it is cached on the SparseTensor under `indice_key` — the same trick spconv uses, and the -reason upstream threads `indice_key=f"res_{resolution}"` through the ResBlocks. -""" - -from __future__ import annotations - -from typing import Optional - -import mlx.core as mx -import mlx.nn as nn -import numpy as np - -from .tensor import SparseTensor - -_MISSING = -1 - - -def _kernel_offsets(k: int) -> np.ndarray: - """Kernel offsets in C order over (dz, dy, dx), centred — matches spconv's ordering.""" - r = np.arange(k) - (k // 2) - return np.stack(np.meshgrid(r, r, r, indexing="ij"), axis=-1).reshape(-1, 3) - - -def build_indice_map(coords: mx.array, kernel_size: int) -> np.ndarray: - """[K^3, N] int32 — for each offset, the row of the neighbour, or -1 if unoccupied. - - Uses a sorted-key binary search rather than a Python dict: at k=3 this is 27 lookups - per voxel, and a per-voxel dict lookup would dominate runtime for any real mesh. - """ - c = np.asarray(coords, dtype=np.int64) - n = c.shape[0] - if n == 0: - return np.full((kernel_size**3, 0), _MISSING, dtype=np.int32) - - offsets = _kernel_offsets(kernel_size) - pad = kernel_size // 2 - - # Encode (batch,z,y,x) into one int64. Shift by `pad` so that neighbour coordinates - # of -1 stay non-negative and cannot alias onto a real cell at the opposite edge. - lo = c.min(axis=0) - pad - ext = (c.max(axis=0) + pad) - lo + 1 - strides = np.array( - [ext[1] * ext[2] * ext[3], ext[2] * ext[3], ext[3], 1], dtype=np.int64 - ) - - def encode(arr: np.ndarray) -> np.ndarray: - return ((arr - lo) * strides).sum(axis=1) - - keys = encode(c) - order = np.argsort(keys, kind="stable") - sorted_keys = keys[order] - - imap = np.empty((offsets.shape[0], n), dtype=np.int32) - for i, d in enumerate(offsets): - probe = c.copy() - probe[:, 1:] += d # batch index (column 0) never shifts - pk = encode(probe) - pos = np.searchsorted(sorted_keys, pk) - pos_clipped = np.clip(pos, 0, n - 1) - hit = sorted_keys[pos_clipped] == pk - imap[i] = np.where(hit, order[pos_clipped], _MISSING).astype(np.int32) - return imap - - -class SubMConv3d(nn.Module): - """Submanifold sparse conv. Weight layout [K^3, in_channels, out_channels].""" - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int = 3, - bias: bool = True, - indice_key: Optional[str] = None, - ): - super().__init__() - if kernel_size % 2 != 1: - raise ValueError(f"kernel_size must be odd, got {kernel_size}") - self.in_channels = in_channels - self.out_channels = out_channels - self.kernel_size = kernel_size - self.indice_key = indice_key - - scale = (in_channels * kernel_size**3) ** -0.5 - self.weight = mx.random.uniform( - -scale, scale, (kernel_size**3, in_channels, out_channels) - ) - if bias: - self.bias = mx.zeros((out_channels,)) - - def _gather_index(self, x: SparseTensor, n: int) -> mx.array: - """Cached [N, K^3] gather index, already on-device. - - Caching only the raw indice map is not enough: rebuilding the "missing -> N" - substitution and re-uploading the index cost more than the convolution itself. - At 128^3/128ch that overhead was ~26ms against ~13ms of actual GPU work, i.e. - two thirds of the measured time was CPU-side bookkeeping repeated every layer. - The transposed, sentinel-substituted device array depends only on the - coordinate set, so it is cached whole. - """ - key = f"gidx_k{self.kernel_size}_{self.indice_key}" - cached = x.cache_get(key) - if cached is not None: - return cached - imap = build_indice_map(x.coords, self.kernel_size) - idx_t = mx.array(np.where(imap == _MISSING, n, imap).T) # [N, K^3] - x.cache_put(key, idx_t) - return idx_t - - def __call__(self, x: SparseTensor) -> SparseTensor: - n = x.feats.shape[0] - - # k=1 touches only the centre voxel, so it is exactly a per-voxel linear — - # skip the indice map entirely. - if self.kernel_size == 1: - out = x.feats @ self.weight[0] - if hasattr(self, "bias"): - out = out + self.bias - return x.replace(out) - - # One appended zero row: absent neighbours index it and contribute nothing, - # which avoids a per-offset boolean mask. - feats_pad = mx.concatenate( - [x.feats, mx.zeros((1, self.in_channels), dtype=x.feats.dtype)], axis=0 - ) - idx_t = self._gather_index(x, n) # [N, K^3], cached on device - k3 = self.kernel_size**3 - - # Fuse the K^3 taps into ONE gather + ONE matmul. - # - # The obvious implementation loops over the K^3 offsets accumulating - # `gather(i) @ W[i]`, but that issues 2*K^3 tiny GPU dispatches with Python - # between them, and each is far too small to fill the machine. Fleet - # benchmarking made this unmistakable: the 80-core M3 Ultra came in SLOWER - # than a 38-core M2 Max (81.6ms vs 58.3ms at 128^3/128ch), i.e. throughput was - # anti-correlated with core count — the signature of launch-latency binding - # rather than compute binding. - # - # Concatenating neighbours along the channel axis turns the whole thing into a - # single [N, K^3*Cin] x [K^3*Cin, Cout] matmul, which is one dispatch big - # enough to actually occupy the GPU. - # - # That buffer is N*K^3*Cin floats, so it is chunked over rows to keep peak - # memory bounded (~256MB/chunk) — at 128^3 and 128ch the unchunked form alone - # would be ~2.9GB, which is fine on a Studio and not fine on an 8GB mini. - w_flat = self.weight.reshape(k3 * self.in_channels, self.out_channels).astype( - x.feats.dtype - ) - bytes_per_row = k3 * self.in_channels * 4 - chunk = max(1, min(n, (256 << 20) // max(bytes_per_row, 1))) - - outs = [] - for start in range(0, n, chunk): - stop = min(start + chunk, n) - g = mx.take(feats_pad, idx_t[start:stop].reshape(-1), axis=0) - g = g.reshape(stop - start, k3 * self.in_channels) - outs.append(g @ w_flat) - out = outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0) - - if hasattr(self, "bias"): - out = out + self.bias.astype(x.feats.dtype) - return x.replace(out) diff --git a/lato_mlx/sparse/ops.py b/lato_mlx/sparse/ops.py deleted file mode 100644 index 44fc92f..0000000 --- a/lato_mlx/sparse/ops.py +++ /dev/null @@ -1,256 +0,0 @@ -"""The rest of the sparse layer set, in MLX. - -Nothing here is CUDA-locked upstream — these are ordinary linear/norm/attention layers -that merely take a SparseTensor instead of a dense one. They are reimplemented rather -than adapted because upstream's versions inherit from torch modules. - -Two normalisation shapes are easy to conflate, and upstream uses both: - - LayerNorm32 applied to `x.feats` directly -> per-voxel over channels. - SparseGroupNorm32 reshapes [N_b, C] -> [1, C, N_b] per batch item, so statistics - are over (channels-in-group x voxels) WITHIN one batch item. - Getting this wrong is silent: shapes match either way. - -Attention runs in `attn_mode="full"`, which upstream defines as full attention *within* -each batch item (never across). Batch rows are contiguous, so each item is one slice. -""" - -from __future__ import annotations - -from typing import Optional - -import mlx.core as mx -import mlx.nn as nn - -from .conv import SubMConv3d -from .tensor import SparseTensor - -# ---------------------------------------------------------------- primitives - - -class SparseLinear(nn.Module): - def __init__(self, in_features: int, out_features: int, bias: bool = True): - super().__init__() - self.linear = nn.Linear(in_features, out_features, bias=bias) - - def __call__(self, x: SparseTensor) -> SparseTensor: - return x.replace(self.linear(x.feats)) - - -class LayerNorm32(nn.Module): - """Per-voxel LayerNorm over channels; computed in fp32 as upstream does.""" - - def __init__(self, dim: int, affine: bool = False, eps: float = 1e-6): - super().__init__() - self.eps = eps - self.affine = affine - if affine: - self.weight = mx.ones((dim,)) - self.bias = mx.zeros((dim,)) - - def __call__(self, feats: mx.array) -> mx.array: - dt = feats.dtype - f = feats.astype(mx.float32) - mu = mx.mean(f, axis=-1, keepdims=True) - var = mx.var(f, axis=-1, keepdims=True) - f = (f - mu) * mx.rsqrt(var + self.eps) - if self.affine: - f = f * self.weight + self.bias - return f.astype(dt) - - -class SparseGroupNorm32(nn.Module): - """GroupNorm over (channels-in-group x voxels), per batch item. fp32 internally.""" - - def __init__(self, num_groups: int, num_channels: int, eps: float = 1e-5): - super().__init__() - if num_channels % num_groups != 0: - raise ValueError(f"{num_channels} channels not divisible by {num_groups}") - self.num_groups = num_groups - self.num_channels = num_channels - self.eps = eps - self.weight = mx.ones((num_channels,)) - self.bias = mx.zeros((num_channels,)) - - def __call__(self, x: SparseTensor) -> SparseTensor: - dt = x.feats.dtype - g, c = self.num_groups, self.num_channels - parts = [] - for sl in x.layout: - f = x.feats[sl].astype(mx.float32) # [n_b, C] - n_b = f.shape[0] - if n_b == 0: - parts.append(f) - continue - # -> [G, (C/G)*n_b] so mean/var cover channels *and* voxels in the group - grouped = f.T.reshape(g, (c // g) * n_b) - mu = mx.mean(grouped, axis=1, keepdims=True) - var = mx.var(grouped, axis=1, keepdims=True) - grouped = (grouped - mu) * mx.rsqrt(var + self.eps) - f = grouped.reshape(c, n_b).T - parts.append(f * self.weight + self.bias) - out = parts[0] if len(parts) == 1 else mx.concatenate(parts, axis=0) - return x.replace(out.astype(dt)) - - -class SparseSiLU(nn.Module): - def __call__(self, x: SparseTensor) -> SparseTensor: - return x.replace(nn.silu(x.feats)) - - -class SparseGELU(nn.Module): - def __call__(self, x: SparseTensor) -> SparseTensor: - return x.replace(nn.gelu(x.feats)) - - -# ---------------------------------------------------------------- blocks - - -class SparseResBlock(nn.Module): - """norm1(affine) -> silu -> conv1 -> norm2(no affine) -> silu -> conv2 + skip.""" - - def __init__(self, channels: int, out_channels: Optional[int] = None): - super().__init__() - self.channels = channels - self.out_channels = out_channels or channels - self.norm1 = LayerNorm32(channels, affine=True, eps=1e-6) - self.norm2 = LayerNorm32(self.out_channels, affine=False, eps=1e-6) - self.conv1 = SubMConv3d(channels, self.out_channels, 3) - self.conv2 = SubMConv3d(self.out_channels, self.out_channels, 3) - self.skip_connection = ( - SparseLinear(channels, self.out_channels) - if channels != self.out_channels - else None - ) - - def __call__(self, x: SparseTensor) -> SparseTensor: - h = x.replace(self.norm1(x.feats)) - h = h.replace(nn.silu(h.feats)) - h = self.conv1(h) - h = h.replace(self.norm2(h.feats)) - h = h.replace(nn.silu(h.feats)) - h = self.conv2(h) - skip = self.skip_connection(x).feats if self.skip_connection else x.feats - return h.replace(h.feats + skip) - - -class SparseFeedForwardNet(nn.Module): - """Upstream is nn.Sequential(Linear, GELU, Linear), so its checkpoint keys are - `mlp.mlp.0` and `mlp.mlp.2` — index 1 is the activation and carries no weights. - Named `mlp_0`/`mlp_2` here because a Python list with a None hole does not survive - MLX's parameter tree; the loader remaps the dotted indices onto these.""" - - def __init__(self, channels: int, mlp_ratio: float = 4.0): - super().__init__() - hidden = int(channels * mlp_ratio) - self.mlp_0 = nn.Linear(channels, hidden) - self.mlp_2 = nn.Linear(hidden, channels) - - def __call__(self, x: SparseTensor) -> SparseTensor: - return x.replace(self.mlp_2(nn.gelu_approx(self.mlp_0(x.feats)))) - - -def _sdpa_per_batch( - q: mx.array, k: mx.array, v: mx.array, layout_q, layout_kv, heads: int, scale: float -) -> mx.array: - """Full attention inside each batch item. q/k/v are [N, H, D] flattened over batch.""" - outs = [] - for sq, skv in zip(layout_q, layout_kv): - qi = q[sq].transpose(1, 0, 2)[None] # [1, H, n, D] - ki = k[skv].transpose(1, 0, 2)[None] - vi = v[skv].transpose(1, 0, 2)[None] - o = mx.fast.scaled_dot_product_attention(qi, ki, vi, scale=scale) - outs.append(o[0].transpose(1, 0, 2)) # [n, H, D] - return outs[0] if len(outs) == 1 else mx.concatenate(outs, axis=0) - - -class SparseMultiHeadAttention(nn.Module): - """attn_mode='full' only — the sole mode LATO.2's model code instantiates.""" - - def __init__( - self, - channels: int, - num_heads: int, - ctx_channels: Optional[int] = None, - attn_type: str = "self", - qkv_bias: bool = True, - ): - super().__init__() - if channels % num_heads != 0: - raise ValueError(f"{channels} channels not divisible by {num_heads} heads") - self.channels = channels - self.num_heads = num_heads - self.head_dim = channels // num_heads - self.scale = self.head_dim**-0.5 - self._type = attn_type - self.ctx_channels = ctx_channels if ctx_channels is not None else channels - if attn_type == "self": - self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias) - else: - self.to_q = nn.Linear(channels, channels, bias=qkv_bias) - self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias) - self.to_out = nn.Linear(channels, channels) - - def __call__( - self, x: SparseTensor, context: Optional[SparseTensor] = None - ) -> SparseTensor: - n = x.feats.shape[0] - h, d = self.num_heads, self.head_dim - if self._type == "self": - qkv = self.to_qkv(x.feats).reshape(n, 3, h, d) - q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] - lq = lkv = x.layout - else: - if context is None: - raise ValueError("cross-attention needs a context") - q = self.to_q(x.feats).reshape(n, h, d) - m = context.feats.shape[0] - kv = self.to_kv(context.feats).reshape(m, 2, h, d) - k, v = kv[:, 0], kv[:, 1] - lq, lkv = x.layout, context.layout - o = _sdpa_per_batch(q, k, v, lq, lkv, h, self.scale) - return x.replace(self.to_out(o.reshape(n, self.channels))) - - -class SparseTransformerBlock(nn.Module): - """Pre-norm self-attention + FFN. Norms are non-affine (ln_affine=False upstream).""" - - def __init__(self, channels: int, num_heads: int, mlp_ratio: float = 4.0): - super().__init__() - self.norm1 = LayerNorm32(channels, affine=False, eps=1e-6) - self.norm2 = LayerNorm32(channels, affine=False, eps=1e-6) - self.attn = SparseMultiHeadAttention(channels, num_heads) - self.mlp = SparseFeedForwardNet(channels, mlp_ratio) - - def __call__(self, x: SparseTensor) -> SparseTensor: - h = self.attn(x.replace(self.norm1(x.feats))) - x = x.replace(x.feats + h.feats) - h = self.mlp(x.replace(self.norm2(x.feats))) - return x.replace(x.feats + h.feats) - - -class SparseTransformerCrossBlock(nn.Module): - """Pre-norm self-attn -> cross-attn -> FFN.""" - - def __init__( - self, channels: int, ctx_channels: int, num_heads: int, mlp_ratio: float = 4.0 - ): - super().__init__() - self.norm1 = LayerNorm32(channels, affine=False, eps=1e-6) - self.norm2 = LayerNorm32(channels, affine=False, eps=1e-6) - self.norm3 = LayerNorm32(channels, affine=False, eps=1e-6) - self.context_norm = LayerNorm32(ctx_channels, affine=False, eps=1e-6) - self.self_attn = SparseMultiHeadAttention(channels, num_heads) - self.cross_attn = SparseMultiHeadAttention( - channels, num_heads, ctx_channels=ctx_channels, attn_type="cross" - ) - self.mlp = SparseFeedForwardNet(channels, mlp_ratio) - - def __call__(self, x: SparseTensor, context: SparseTensor) -> SparseTensor: - h = self.self_attn(x.replace(self.norm1(x.feats))) - x = x.replace(x.feats + h.feats) - ctx = context.replace(self.context_norm(context.feats)) - h = self.cross_attn(x.replace(self.norm2(x.feats)), ctx) - x = x.replace(x.feats + h.feats) - h = self.mlp(x.replace(self.norm3(x.feats))) - return x.replace(x.feats + h.feats) diff --git a/lato_mlx/sparse/tensor.py b/lato_mlx/sparse/tensor.py deleted file mode 100644 index 6261a01..0000000 --- a/lato_mlx/sparse/tensor.py +++ /dev/null @@ -1,159 +0,0 @@ -"""SparseTensor for MLX — a coords/feats pair, mirroring LATO.2's upstream container. - -Upstream wraps either `spconv.SparseConvTensor` or `torchsparse.SparseTensor`; both are -CUDA-only, which is what blocks LATO.2 on Apple Silicon. Nothing about the *data* needs -CUDA — it is just coordinates plus features — so this is a plain MLX reimplementation. - -Layout matches upstream exactly so weight conversion stays a straight mapping: - coords : int32 [N, 4] -> (batch, z, y, x) - feats : float [N, C] -and rows belonging to one batch item are contiguous (upstream asserts this). -""" - -from __future__ import annotations - -from typing import List, Optional, Tuple - -import mlx.core as mx -import numpy as np - - -class SparseTensor: - """N non-empty voxels, each with a coordinate and a feature vector.""" - - __slots__ = ("feats", "coords", "_scale", "_spatial_cache", "_layout") - - def __init__( - self, - feats: mx.array, - coords: mx.array, - scale: Tuple[int, int, int] = (1, 1, 1), - spatial_cache: Optional[dict] = None, - layout: Optional[List[slice]] = None, - ): - if feats.shape[0] != coords.shape[0]: - raise ValueError( - f"feats/coords length mismatch: {feats.shape[0]} vs {coords.shape[0]}" - ) - if coords.ndim != 2 or coords.shape[1] != 4: - raise ValueError(f"coords must be [N, 4] (batch,z,y,x), got {coords.shape}") - self.feats = feats - self.coords = coords - self._scale = tuple(scale) - # Indice maps are expensive to build and identical for every conv that shares a - # coordinate set — upstream exploits this via spconv's `indice_key`. Same idea. - self._spatial_cache = spatial_cache if spatial_cache is not None else {} - self._layout = layout - - # -- basics --------------------------------------------------------------- - @property - def shape(self) -> Tuple[int, int]: - return (self.batch_size, self.feats.shape[1]) - - @property - def batch_size(self) -> int: - if self.coords.shape[0] == 0: - return 0 - return int(mx.max(self.coords[:, 0]).item()) + 1 - - @property - def layout(self) -> List[slice]: - """One slice per batch item. Relies on batch-contiguity, as upstream does.""" - if self._layout is None: - b = np.asarray(self.coords[:, 0], dtype=np.int64) - counts = np.bincount(b, minlength=self.batch_size) - offs = np.cumsum(counts) - self._layout = [ - slice(int(offs[i] - counts[i]), int(offs[i])) for i in range(len(counts)) - ] - return self._layout - - def replace(self, feats: mx.array) -> "SparseTensor": - """New tensor, same coordinates — so the indice-map cache stays valid.""" - return SparseTensor( - feats, - self.coords, - scale=self._scale, - spatial_cache=self._spatial_cache, - layout=self._layout, - ) - - # -- cache ---------------------------------------------------------------- - def cache_get(self, key: str): - return self._spatial_cache.get(key) - - def cache_put(self, key: str, value) -> None: - self._spatial_cache[key] = value - - def __repr__(self) -> str: - return ( - f"SparseTensor(N={self.coords.shape[0]}, C={self.feats.shape[1]}, " - f"batch={self.batch_size}, scale={self._scale})" - ) - - -def subdivide(x: SparseTensor) -> SparseTensor: - """Upsample ×2 by splitting each voxel into its 8 children (nearest-neighbour). - - Mirrors upstream `SparseSubdivide`: coords are doubled then offset by the unit - cube, features are replicated. Child order is the C-order of nonzero(ones(2,2,2)), - matching upstream's `torch.nonzero`, so replicated features line up identically. - """ - n = x.coords.shape[0] - offsets = np.stack(np.meshgrid(*[np.arange(2)] * 3, indexing="ij"), -1).reshape(-1, 3) - offsets = np.concatenate([np.zeros((8, 1), dtype=offsets.dtype), offsets], axis=1) - - base = np.asarray(x.coords, dtype=np.int32).copy() - base[:, 1:] *= 2 - new_coords = (base[:, None, :] + offsets[None, :, :]).reshape(n * 8, 4) - - new_feats = mx.repeat(x.feats, 8, axis=0) - return SparseTensor( - new_feats, - mx.array(new_coords, dtype=mx.int32), - scale=tuple(s * 2 for s in x._scale), - ) - - -def downsample(x: "SparseTensor", factor: int = 2) -> "SparseTensor": - """Downsample by `factor`, reducing colliding voxels with MAX. - - Upstream's docstring says "average pooling" but the implementation passes - reduce="amax" (the `reduce='mean'` line is commented out). Following the code, - not the docstring — mean vs max here is numerically silent in shape and would - quietly change every downsampled feature. - - Output coordinates come out sorted by the same packed code upstream sorts on, so - rows stay batch-contiguous as SparseTensor requires. - """ - import numpy as _np - - c = _np.asarray(x.coords, dtype=_np.int64).copy() - c[:, 1:] //= factor - - maxs = c[:, 1:].max(axis=0) + 1 - # OFFSET = reversed cumprod, matching upstream's packing - off = _np.array( - [maxs[0] * maxs[1] * maxs[2], maxs[1] * maxs[2], maxs[2], 1], dtype=_np.int64 - ) - code = (c * off).sum(axis=1) - - uniq, inv = _np.unique(code, return_inverse=True) - feats = _np.asarray(x.feats) - out = _np.full((uniq.shape[0], feats.shape[1]), -_np.inf, dtype=feats.dtype) - _np.maximum.at(out, inv, feats) - - new_coords = _np.stack( - [ - uniq // off[0], - (uniq // off[1]) % maxs[0], - (uniq // off[2]) % maxs[1], - uniq % maxs[2], - ], - axis=-1, - ).astype(_np.int32) - return SparseTensor( - mx.array(out), - mx.array(new_coords), - scale=tuple(s * factor for s in x._scale), - ) diff --git a/tests/test_ops.py b/tests/test_ops.py deleted file mode 100644 index 5406fde..0000000 --- a/tests/test_ops.py +++ /dev/null @@ -1,171 +0,0 @@ -"""Sparse layer tests against torch. - -Unlike the submanifold conv — where spconv is uninstallable and the oracle had to be -hand-written — every layer here has a real torch counterpart, so these compare against -upstream's actual semantics rather than a paraphrase of them. The upstream forward -bodies are reproduced verbatim (see modules/sparse/{norm,linear,nonlinearity}.py), -including the [N_b,C] -> [1,C,N_b] GroupNorm reshape, which is the one that would fail -silently if guessed. -""" - -import sys -from pathlib import Path - -import mlx.core as mx -import numpy as np -import torch -import torch.nn.functional as F - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from lato_mlx.sparse.ops import ( # noqa: E402 - LayerNorm32, - SparseGroupNorm32, - SparseMultiHeadAttention, - SparseTransformerBlock, -) -from lato_mlx.sparse.tensor import SparseTensor # noqa: E402 - - -def make_batched(n_per_batch, channels, seed=0): - """Batch-contiguous coords, as upstream requires.""" - rng = np.random.default_rng(seed) - coords, feats = [], [] - for b, nb in enumerate(n_per_batch): - for i in range(nb): - coords.append((b, i // 16, (i // 4) % 4, i % 4)) - feats.append(rng.standard_normal((nb, channels)).astype(np.float32)) - return np.array(coords, dtype=np.int32), np.concatenate(feats, 0) - - -def test_group_norm_matches_torch(groups=8, channels=32): - n_per_batch = [37, 51] - coords, feats = make_batched(n_per_batch, channels, seed=1) - rng = np.random.default_rng(2) - w = rng.standard_normal(channels).astype(np.float32) - b = rng.standard_normal(channels).astype(np.float32) - - gn = SparseGroupNorm32(groups, channels) - gn.weight, gn.bias = mx.array(w), mx.array(b) - got = np.asarray(gn(SparseTensor(mx.array(feats), mx.array(coords))).feats) - - # upstream: per batch item, [N_b,C] -> permute -> [1,C,N_b] -> nn.GroupNorm - tg = torch.nn.GroupNorm(groups, channels, eps=1e-5, affine=True) - tg.weight.data = torch.tensor(w) - tg.bias.data = torch.tensor(b) - want = np.zeros_like(feats) - off = 0 - for nb in n_per_batch: - bf = torch.tensor(feats[off : off + nb]) - bf = bf.permute(1, 0).reshape(1, channels, -1) - bf = tg(bf) - want[off : off + nb] = bf.reshape(channels, -1).permute(1, 0).detach().numpy() - off += nb - - err = np.abs(got - want).max() - assert err < 2e-4, f"group norm err {err:.3g}" - return err - - -def test_group_norm_is_not_per_voxel(groups=8, channels=32): - """Guard the easy-to-miss distinction: GroupNorm here is NOT a per-voxel norm.""" - coords, feats = make_batched([40], channels, seed=5) - gn = SparseGroupNorm32(groups, channels) - got = np.asarray(gn(SparseTensor(mx.array(feats), mx.array(coords))).feats) - per_voxel = torch.nn.functional.group_norm( - torch.tensor(feats).reshape(40, channels), groups - ).numpy() - assert np.abs(got - per_voxel).max() > 1e-3, "matched per-voxel norm — reshape lost" - return 0.0 - - -def test_layer_norm_matches_torch(channels=64): - rng = np.random.default_rng(3) - feats = rng.standard_normal((50, channels)).astype(np.float32) - ln = LayerNorm32(channels, affine=False, eps=1e-6) - got = np.asarray(ln(mx.array(feats))) - want = F.layer_norm(torch.tensor(feats), (channels,), eps=1e-6).numpy() - err = np.abs(got - want).max() - assert err < 1e-5, f"layer norm err {err:.3g}" - return err - - -def test_self_attention_matches_torch(channels=64, heads=8): - n_per_batch = [23, 31] - coords, feats = make_batched(n_per_batch, channels, seed=4) - rng = np.random.default_rng(6) - wq = rng.standard_normal((channels * 3, channels)).astype(np.float32) * 0.05 - bq = rng.standard_normal((channels * 3,)).astype(np.float32) * 0.05 - wo = rng.standard_normal((channels, channels)).astype(np.float32) * 0.05 - bo = rng.standard_normal((channels,)).astype(np.float32) * 0.05 - - attn = SparseMultiHeadAttention(channels, heads) - attn.to_qkv.weight, attn.to_qkv.bias = mx.array(wq), mx.array(bq) - attn.to_out.weight, attn.to_out.bias = mx.array(wo), mx.array(bo) - got = np.asarray(attn(SparseTensor(mx.array(feats), mx.array(coords))).feats) - - # reference: attention strictly within each batch item - d = channels // heads - want = np.zeros_like(feats) - off = 0 - for nb in n_per_batch: - f = torch.tensor(feats[off : off + nb]) - qkv = F.linear(f, torch.tensor(wq), torch.tensor(bq)).reshape(nb, 3, heads, d) - q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2] - o = F.scaled_dot_product_attention( - q.permute(1, 0, 2)[None], k.permute(1, 0, 2)[None], v.permute(1, 0, 2)[None] - ) - o = o[0].permute(1, 0, 2).reshape(nb, channels) - want[off : off + nb] = ( - F.linear(o, torch.tensor(wo), torch.tensor(bo)).detach().numpy() - ) - off += nb - - err = np.abs(got - want).max() - assert err < 2e-4, f"attention err {err:.3g}" - return err - - -def test_attention_does_not_cross_batches(channels=32, heads=4): - """Perturbing batch 1 must never change batch 0's output.""" - coords, feats = make_batched([12, 12], channels, seed=7) - attn = SparseMultiHeadAttention(channels, heads) - a = np.asarray(attn(SparseTensor(mx.array(feats), mx.array(coords))).feats) - f2 = feats.copy() - f2[12:] += 10.0 - b = np.asarray(attn(SparseTensor(mx.array(f2), mx.array(coords))).feats) - assert np.abs(a[:12] - b[:12]).max() < 1e-5, "batch 0 changed — attention leaked" - return 0.0 - - -def test_transformer_block_shape(channels=64, heads=8): - coords, feats = make_batched([20, 20], channels, seed=8) - blk = SparseTransformerBlock(channels, heads) - out = blk(SparseTensor(mx.array(feats), mx.array(coords))) - assert out.feats.shape == (40, channels) - assert np.isfinite(np.asarray(out.feats)).all(), "non-finite output" - return 0.0 - - -if __name__ == "__main__": - tests = [ - ("group norm vs torch", test_group_norm_matches_torch), - ("group norm != per-voxel", test_group_norm_is_not_per_voxel), - ("layer norm vs torch", test_layer_norm_matches_torch), - ("self-attn vs torch", test_self_attention_matches_torch), - ("attn batch isolation", test_attention_does_not_cross_batches), - ("transformer block", test_transformer_block_shape), - ] - failed = 0 - for name, fn in tests: - try: - err = fn() - print(f" PASS {name:26s} (max err {err:.2e})") - except AssertionError as e: - print(f" FAIL {name:26s} {e}") - failed += 1 - except Exception as e: # noqa: BLE001 - print(f" ERROR {name:26s} {type(e).__name__}: {e}") - failed += 1 - print(f"\n{len(tests)-failed}/{len(tests)} passed") - sys.exit(1 if failed else 0) diff --git a/tests/test_sparse.py b/tests/test_sparse.py deleted file mode 100644 index 942f5c7..0000000 --- a/tests/test_sparse.py +++ /dev/null @@ -1,162 +0,0 @@ -"""Correctness tests for the MLX sparse core. - -spconv cannot be installed on this machine — that is the entire reason this port exists — -so there is no way to diff against upstream numerically here. Instead the vectorised MLX -implementation is checked against a deliberately naive, obviously-correct reference -written straight from the definition of submanifold convolution (dict lookup, per-voxel -Python loop). The two share no indexing code, so an off-by-one in the fast path cannot -hide in both. -""" - -import sys -from pathlib import Path - -import mlx.core as mx -import numpy as np - -sys.path.insert(0, str(Path(__file__).resolve().parents[1])) - -from lato_mlx.sparse.conv import SubMConv3d, build_indice_map, _kernel_offsets -from lato_mlx.sparse.tensor import SparseTensor, subdivide - - -def reference_subm_conv(coords, feats, weight, bias, k): - """Definition of submanifold conv, written for obviousness, not speed.""" - occupied = {tuple(c): i for i, c in enumerate(coords.tolist())} - offsets = _kernel_offsets(k) - n, out_c = coords.shape[0], weight.shape[2] - out = np.zeros((n, out_c), dtype=np.float64) - for i, c in enumerate(coords.tolist()): - for oi, d in enumerate(offsets): - nb = (c[0], c[1] + d[0], c[2] + d[1], c[3] + d[2]) - j = occupied.get(nb) - if j is not None: - out[i] += feats[j].astype(np.float64) @ weight[oi].astype(np.float64) - if bias is not None: - out += bias.astype(np.float64) - return out - - -def random_sparse(n_vox, channels, batch=2, res=8, seed=0): - rng = np.random.default_rng(seed) - seen, coords = set(), [] - while len(coords) < n_vox: - b = int(rng.integers(0, batch)) - z, y, x = (int(v) for v in rng.integers(0, res, 3)) - if (b, z, y, x) in seen: - continue - seen.add((b, z, y, x)) - coords.append((b, z, y, x)) - coords.sort() # upstream requires batch-contiguous rows - coords = np.array(coords, dtype=np.int32) - feats = rng.standard_normal((n_vox, channels)).astype(np.float32) - return coords, feats - - -def test_subm_conv_matches_reference(k=3, n=180, cin=12, cout=7): - coords, feats = random_sparse(n, cin, seed=1) - rng = np.random.default_rng(2) - w = rng.standard_normal((k**3, cin, cout)).astype(np.float32) * 0.1 - b = rng.standard_normal((cout,)).astype(np.float32) - - conv = SubMConv3d(cin, cout, k, bias=True, indice_key="t") - conv.weight = mx.array(w) - conv.bias = mx.array(b) - - got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) - want = reference_subm_conv(coords, feats, w, b, k) - err = np.abs(got - want).max() - assert err < 2e-4, f"k={k} max abs err {err:.3g}" - return err - - -def test_kernel1_is_pointwise(n=64, cin=8, cout=5): - coords, feats = random_sparse(n, cin, seed=3) - rng = np.random.default_rng(4) - w = rng.standard_normal((1, cin, cout)).astype(np.float32) - conv = SubMConv3d(cin, cout, 1, bias=False, indice_key="p") - conv.weight = mx.array(w) - got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) - err = np.abs(got - feats @ w[0]).max() - assert err < 1e-4, f"k=1 err {err:.3g}" - return err - - -def test_isolated_voxel_sees_only_itself(): - """A voxel with no occupied neighbours must reduce to the centre tap alone.""" - coords = np.array([[0, 0, 0, 0], [0, 50, 50, 50]], dtype=np.int32) - feats = np.ones((2, 3), dtype=np.float32) - w = np.zeros((27, 3, 3), dtype=np.float32) - centre = 13 # index of (0,0,0) in centred C-order offsets - assert tuple(_kernel_offsets(3)[centre]) == (0, 0, 0) - w[centre] = np.eye(3) - conv = SubMConv3d(3, 3, 3, bias=False, indice_key="iso") - conv.weight = mx.array(w) - got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) - assert np.abs(got - feats).max() < 1e-6, got - return 0.0 - - -def test_batches_do_not_leak(): - """Same spatial cell in two batch items must not become neighbours.""" - coords = np.array([[0, 1, 1, 1], [1, 1, 1, 1]], dtype=np.int32) - feats = np.array([[1.0], [100.0]], dtype=np.float32) - w = np.ones((27, 1, 1), dtype=np.float32) # sum every occupied neighbour - conv = SubMConv3d(1, 1, 3, bias=False, indice_key="b") - conv.weight = mx.array(w) - got = np.asarray(conv(SparseTensor(mx.array(feats), mx.array(coords))).feats) - assert np.allclose(got, [[1.0], [100.0]]), f"batch leak: {got}" - return 0.0 - - -def test_indice_map_hit_rate(): - """A fully dense block: interior voxels must find all 27 neighbours.""" - coords = np.array( - [[0, z, y, x] for z in range(4) for y in range(4) for x in range(4)], - dtype=np.int32, - ) - imap = build_indice_map(mx.array(coords), 3) - lin = {tuple(c): i for i, c in enumerate(coords.tolist())} - interior = [lin[(0, z, y, x)] for z in (1, 2) for y in (1, 2) for x in (1, 2)] - assert (imap[:, interior] != -1).all(), "interior voxel missing a neighbour" - corner = lin[(0, 0, 0, 0)] - assert (imap[:, corner] != -1).sum() == 8, "corner should see exactly 8 of 27" - return 0.0 - - -def test_subdivide(): - coords = np.array([[0, 1, 2, 3]], dtype=np.int32) - feats = np.array([[5.0, 6.0]], dtype=np.float32) - out = subdivide(SparseTensor(mx.array(feats), mx.array(coords))) - c = np.asarray(out.coords) - assert c.shape == (8, 4) and out.feats.shape == (8, 2) - assert set(map(tuple, c[:, 1:].tolist())) == { - (2 + a, 4 + b, 6 + d) for a in (0, 1) for b in (0, 1) for d in (0, 1) - } - assert np.abs(np.asarray(out.feats) - feats).max() < 1e-6 - return 0.0 - - -if __name__ == "__main__": - tests = [ - ("subm k=3 vs reference", lambda: test_subm_conv_matches_reference(3)), - ("subm k=5 vs reference", lambda: test_subm_conv_matches_reference(5, n=140)), - ("k=1 is pointwise", test_kernel1_is_pointwise), - ("isolated voxel", test_isolated_voxel_sees_only_itself), - ("batch isolation", test_batches_do_not_leak), - ("indice map hit rate", test_indice_map_hit_rate), - ("subdivide", test_subdivide), - ] - failed = 0 - for name, fn in tests: - try: - err = fn() - print(f" PASS {name:28s} (max err {err:.2e})") - except AssertionError as e: - print(f" FAIL {name:28s} {e}") - failed += 1 - except Exception as e: # noqa: BLE001 - print(f" ERROR {name:28s} {type(e).__name__}: {e}") - failed += 1 - print(f"\n{len(tests)-failed}/{len(tests)} passed") - sys.exit(1 if failed else 0)