lato.2_mrp_mlx

An MLX port of 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:

# 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

• SparseTensor container + subdivide
• SubMConv3d (k=3 and k=1) — 7/7 correctness tests pass, max err 3e-7 vs an independent naive reference
• Weight converter, all 7 checkpoints → MLX safetensors (3.3 GB)
• Remaining sparse ops — SparseLinear, LayerNorm32, SparseGroupNorm32, activations, SparseResBlock, self/cross attention, transformer blocks. 6/6 pass against torch (real oracle; only spconv was unavailable)
• SparseDownsample (max-pool, not average — upstream passes reduce="amax")
• 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

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.