70cd436eb9
Extracted from lato.2_mrp_mlx. The same sparse module underlies LATO.2, Pixal3D and the rest of the TRELLIS.2 family, and all of them are blocked on Apple Silicon by the same single op - submanifold conv - so it belongs in one tested package rather than vendored per port. 13/13 tests: 7 for the conv against a hand-written reference (spconv is uninstallable here so there is no upstream oracle), 6 for the remaining layers against torch. SubMConv3d runs 13.8ms at 128^3/128ch on m3ultra, 5.9x faster than the obvious per-offset loop.
79 lines
3.8 KiB
Markdown
79 lines
3.8 KiB
Markdown
# 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).
|