# trellis_sparse_mlx

The TRELLIS-lineage sparse module, in MLX, so the models built on it run on Apple Silicon.

## Why this is a package and not vendored code

Microsoft's TRELLIS.2 sparse stack has been inherited near-verbatim by a growing family
of 3D generation models — [LATO.2](https://github.com/LoHhhha/LATO.2) and TencentARC's
[Pixal3D](https://github.com/TencentARC/Pixal3D) among them. Every one of them hard-requires
`spconv` or `torchsparse`:

```python
BACKEND = 'spconv'    # accepts only ['spconv', 'torchsparse'] — both CUDA-only
ATTN = 'flash_attn'   # accepts only ['xformers', 'flash_attn'] — both CUDA-only
```

Neither sparse backend has a Metal build, and there is no SDPA fallback. That single
dependency is what keeps the entire lineage off Apple Silicon.

The blocker turns out to be **one operation**. Every `SparseConv3d` in these models is
constructed `stride=1, padding=None`, which spconv dispatches to `SubMConv3d`. Nothing
instantiates strided or inverse sparse conv. Implement submanifold convolution and the
remainder is ordinary linear / norm / attention work.

Since each new model in this family needs the same core, it lives here once — tested —
rather than being copy-pasted per port.

## What's in it

| | |
|---|---|
| `SubMConv3d` | submanifold 3×3×3 (and 1×1×1) sparse conv — the actual blocker |
| `SparseTensor` | coords `[N,4]` + feats `[N,C]`, batch-contiguous, with an indice-map cache |
| `subdivide` / `downsample` | ×2 coord expansion; ×2 **max**-pool (upstream says "average" in its docstring but passes `reduce="amax"`) |
| `SparseLinear`, `LayerNorm32`, `SparseGroupNorm32`, `SparseSiLU/GELU` | |
| `SparseResBlock`, `SparseFeedForwardNet` | |
| `SparseMultiHeadAttention` (self + cross), `SparseTransformerBlock`, `SparseTransformerCrossBlock` | `attn_mode="full"`, which is all these models instantiate |

## Performance

`SubMConv3d`, 128³ grid / 128 channels / ~10% occupancy (209,715 voxels), warm indice cache:

| box | chip | GPU cores | conv_ms |
|---|---|---|---|
| m3ultra | M3 Ultra | 80 | **13.8** |
| m2max | M2 Max | 38 | 27.3 |
| m1ultra | M1 Ultra | 64 | 39.7 |
| m1max | M1 Max | 32 | 41.4 |
| m4pro | M4 Pro | ~20 | 55.1 |

The naive implementation — a loop accumulating `gather(i) @ W[i]` over the 27 offsets —
was **5.9× slower**, and the giveaway was that throughput *anti-correlated* with GPU core
count: the 80-core Ultra lost to a 38-core M2 Max, because 2·K³ = 54 tiny dispatches per
layer never fill the machine and UltraFusion punishes small dispatches hardest.

Two fixes: fuse the taps into a single `[N, K³·Cin] × [K³·Cin, Cout]` matmul, and cache the
*prepared device* gather index (not just the raw indice map — rebuilding the sentinel
substitution and re-uploading cost ~26 ms against ~13 ms of real GPU work). Only after
both does core count predict performance, which is the sign it is finally compute-bound.

## Correctness

`spconv` cannot be installed on Apple Silicon — that is the entire reason this exists — so
there is no upstream oracle for the conv. `tests/test_sparse.py` instead checks the
vectorised implementation against a deliberately naive one written from the definition
(dict lookup, per-voxel loop) sharing no indexing code, so an off-by-one cannot hide in
both. `tests/test_ops.py` compares every other layer against **torch**, which *is*
available. 13/13 pass.

**One assumption is unverified**: whether spconv gathers `feats[c+d]` (cross-correlation,
the deep-learning convention and what this implements) or `feats[c-d]`. A flipped kernel is
numerically silent — latent statistics were tried as a discriminator and do not separate
them. It needs an end-to-end reconstruction to settle.

## Licence

MIT. Derived from the TRELLIS/TRELLIS.2 sparse module (Copyright Microsoft Corporation and
VAST-AI-Research contributors, MIT).