21c9de4bdd
Pixal3D is built on the same TRELLIS.2 sparse module, so the core belongs in one tested package rather than a copy per port. Verified behaviour-preserving: the encoder still loads 102/102 params and produces a bit-identical latent (mean +0.0508, std 0.9186) after the extraction.
131 lines
6.1 KiB
Markdown
131 lines
6.1 KiB
Markdown
# lato.2_mrp_mlx
|
||
|
||
An MLX port of [LATO.2](https://github.com/LoHhhha/LATO.2) — factorised 3D mesh generation
|
||
(vertex flow, then connectivity flow) — so it runs natively on Apple Silicon.
|
||
|
||
## Why
|
||
|
||
LATO.2 generates meshes to a **controllable vertex budget**, which is the interesting part:
|
||
it sidesteps the generate-dense-then-decimate loop that TRELLIS-style pipelines force on
|
||
you. But upstream inherits TRELLIS.2's `setup.sh` and hard-requires CUDA:
|
||
|
||
```python
|
||
# upstream modules/sparse/__init__.py
|
||
BACKEND = 'spconv' # accepts only ['spconv', 'torchsparse'] — both CUDA-only
|
||
ATTN = 'flash_attn' # accepts only ['xformers', 'flash_attn'] — both CUDA-only
|
||
```
|
||
|
||
No SDPA fallback, no MPS path. Neither sparse backend has a Metal build.
|
||
|
||
## The scope, once measured
|
||
|
||
The CUDA surface is far smaller than `setup.sh --all` implies. Of the seven released
|
||
checkpoints, five are entirely dense. Sparse code touches only `vertex_autoencoder` and
|
||
`vertex_structured_flow`, and **every** `SparseConv3d` in the model is constructed with the
|
||
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` | `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` |
|
||
| `flash_attn` / `xformers` | `attn_mode="full"` everywhere → plain SDPA |
|
||
|
||
## Status
|
||
|
||
- [x] `SparseTensor` container + `subdivide`
|
||
- [x] `SubMConv3d` (k=3 and k=1) — **7/7 correctness tests pass**, max err 3e-7 vs an
|
||
independent naive reference
|
||
- [x] Weight converter, all 7 checkpoints → MLX safetensors (3.3 GB)
|
||
- [x] Remaining sparse ops — `SparseLinear`, `LayerNorm32`, `SparseGroupNorm32`,
|
||
activations, `SparseResBlock`, self/cross attention, transformer blocks.
|
||
**6/6 pass against torch** (real oracle; only spconv was unavailable)
|
||
- [x] `SparseDownsample` (max-pool, not average — upstream passes `reduce="amax"`)
|
||
- [x] **V-VAE encoder runs on the real checkpoint** — 102/102 params loaded, 0 missing,
|
||
0 unmapped. 16,934 voxels -> 56 latent voxels in 135 ms on m3ultra; latent is
|
||
mean +0.05 / std 0.92, i.e. the ~N(0,1) a KL-trained VAE should produce
|
||
- [ ] V-VAE decoder (multi-resolution + pruning heads) -> reconstruction
|
||
- [ ] V-Flow, T-VAE, T-Flow, voxel encoder
|
||
- [ ] End-to-end inference
|
||
- [ ] Fleet benchmark (m1max / m2max / m4pro / m1ultra / m3ultra)
|
||
|
||
## Fleet benchmark
|
||
|
||
`SubMConv3d` at 128³ grid, 128 channels, ~10% occupancy (209,715 voxels) — the hot op
|
||
in the V-VAE. `conv_ms` is GPU work with a warm indice cache.
|
||
|
||
| box | chip | GPU cores | conv_ms | Mvox/s |
|
||
|---|---|---|---|---|
|
||
| m3ultra | M3 Ultra | 80 | **13.8** | 15.25 |
|
||
| m2max | M2 Max | 38 | 27.3 | 7.69 |
|
||
| m1ultra | M1 Ultra | 64 | 39.7 | 5.28 |
|
||
| m1max | M1 Max | 32 | 41.4 | 5.06 |
|
||
|
||
Getting there took two rounds, both of which the fleet data — not local profiling —
|
||
made visible:
|
||
|
||
| version | m3ultra | note |
|
||
|---|---|---|
|
||
| per-offset loop | 81.6 ms | *slower than a 38-core M2 Max* |
|
||
| fused gather + matmul | 39.3 ms | 2.1× |
|
||
| cached device index | **13.8 ms** | 5.9× total |
|
||
|
||
The first version was launch-latency bound: 2·K³ = 54 tiny dispatches per layer, none
|
||
large enough to occupy the GPU. The giveaway was throughput being *anti-correlated* with
|
||
core count — the 80-core Ultra lost to every smaller box, because its fused-die design
|
||
punishes small dispatches hardest. Fusing the taps into one `[N,K³·Cin]×[K³·Cin,Cout]`
|
||
matmul fixed the dispatch count; caching the prepared device index removed ~26 ms of
|
||
per-layer numpy bookkeeping that was hiding behind ~13 ms of real GPU work.
|
||
|
||
Only after both did core count start predicting performance, which is the sign the op is
|
||
finally compute-bound. Note the M1 Ultra (64 cores) still barely beats the M1 Max (32) —
|
||
gather-heavy work scales poorly across UltraFusion.
|
||
|
||
**Next target:** the indice map is now dominant (232–357 ms, CPU numpy) and is only
|
||
amortised because it is cached per coordinate set. It matters whenever coordinates change.
|
||
|
||
## Correctness
|
||
|
||
`spconv` cannot be installed here — that is the reason this port exists — so there is no
|
||
numerical diff against upstream. Instead `tests/test_sparse.py` checks the vectorised
|
||
implementation against a deliberately naive one written straight from the definition
|
||
(dict lookup, per-voxel loop). They share no indexing code, so an off-by-one cannot hide
|
||
in both. Tests also cover batch isolation, isolated voxels, and indice-map hit rates.
|
||
|
||
**One assumption remains unverified**: whether spconv gathers `feats[c+d]`
|
||
(cross-correlation — the deep-learning convention, and what this implements) or
|
||
`feats[c-d]`.
|
||
|
||
Latent statistics were tried as a cheap discriminator and **do not resolve it**. The flip
|
||
is definitely not a no-op (max output delta 3.53), but both orientations yield a
|
||
plausible near-unit-normal latent — std 0.919 as-implemented vs 0.945 flipped, a gap well
|
||
inside the noise of a synthetic input. So the question genuinely needs the decoder:
|
||
reconstruction quality is the discriminator, since a mirrored kernel should produce
|
||
visibly wrong geometry while leaving the statistics intact. `--flip-kernel` builds the
|
||
alternative.
|
||
|
||
## Use
|
||
|
||
```bash
|
||
uv venv --python 3.12 .venv
|
||
VIRTUAL_ENV=.venv uv pip install mlx numpy torch trimesh
|
||
|
||
hf download 0x4c48/LATO.2 --local-dir ckpt # 3.3 GB upstream weights
|
||
.venv/bin/python -m lato_mlx.convert --ckpt ckpt --out weights
|
||
.venv/bin/python tests/test_sparse.py
|
||
```
|
||
|
||
Upstream source is vendored read-only under `upstream/LATO.2` for reference.
|
||
|
||
## Licence
|
||
|
||
Upstream LATO.2 is MIT (Copyright the LATO.2 authors); its sparse module carries
|
||
Microsoft and VAST-AI-Research copyright, also MIT. This port is MIT on the same terms.
|