modelbeast/docs/TRELLIS2_MLX_RECON.md
2026-07-18 23:44:30 +10:00

4.8 KiB
Raw Blame History

TRELLIS.2 on Apple Silicon — recon (2026-07-18)

Verdict up front: we already run TRELLIS.2 locally — vendor/trellis-mac (torch-MPS + Metal kernels) IS TRELLIS.2-4B and is wired as the trellis_mac operator. There is no "MLX interface" shortcut for CUDA code (CUDA kernels must be rewritten, not wrapped), but ~90% of TRELLIS.2 is standard tensor math that MPS/MLX already runs; only four custom CUDA pieces matter, and the community has replaced all four. The open work is speed (unfused sparse ops ≈ 10× slower than CUDA) and quality gaps (hole-filling disabled, forced pre-simplification) — not feasibility.

What we run today (vendor/trellis-mac, shivampkumar fork)

  • torch-MPS + PYTORCH_ENABLE_MPS_FALLBACK, Metal kernels by @pedronaugusto (mtlgemm / mtldiffrast / mtlbvh / mtlmesh), SDPA attention.
  • Full inference: sparse structure → shape SLat → tex SLat → decode → dual-grid mesh extract → simplify → Metal PBR bake → GLB.
  • Our benchmark: 318 s/gen on m3ultra, ~18 GB peak (M4 Pro 24GB: 5m13s cold at 512). H100 does 317 s — the gap is almost entirely the unfused sparse conv + padded attention.
  • Known gaps vs fal's trellis-2: CuMesh skipped → no hole filling, meshes pre-simplified ~858K→200K faces before baking, Metal BVH instability, macOS GPU-watchdog can kill long kernels (detected + workarounds printed by generate.py). macOS 26 needed for the metallib (fleet is on 26.5 ✓).
  • Licensing: DINOv3 is Meta-gated; RMBG-2.0 preprocessing is CC BY-NC.

What TRELLIS.2 actually needs (upstream: Linux, CUDA 12.4, ≥24 GB)

CUDA dep Role Mac status
o-voxel ext (hash/convert/rasterize) the O-Voxel representation pure-Python reimpl (trellis-mac backends/) + CPU fork
FlexGEMM (Triton) all sparse conv the crux — Triton ≠ Metal; Metal mtlgemm or slow pure-torch gather/scatter
flash-attn / xformers 4B flow transformer attention SDPA (padded → unfused, big cost)
CuMesh decimate/remesh/hole-fill/UV skipped → fast_simplification; hole-fill lost
nvdiffrast texture bake mtldiffrast (Metal)
nvdiffrec preview renders only not needed

The MLX question, answered

  • No shim exists or can exist: CUDA kernels are NVIDIA-machine code; "an MLX interface" means rewriting each custom op in MLX/Metal. MLX can express them (mx.fast.metal_kernel JIT-compiles Metal from Python) — but MLX has no sparse-tensor type and no sparse-voxel precedent; SparseTensor/varlen semantics must be hand-rolled. A TRELLIS.2-MLX would be a first.
  • Our own playbook (proven 2×) is adopt-then-patch, not from-scratch: hunyuan3d-mlx = dgrauet's 8.4k-LOC port + our thin packaging layer; corridorkey-mlx = cmoyates/Niko's 5.3k-LOC port + our 8-line mx.compile-in-tiled-mode patch = the 1.47× win. From-scratch dense→MLX ports of this scale are months of solo work (mflux, mlx-video authors).
  • Starting points already exist: upstream PR #175 (Jourloy, 2026-07-17 — MPS + Metal + an experimental mlx backend flag, 28 tests green on M4 Max) and pedronaugusto/trellis2-apple (an mlx_backend/ dir, no benchmarks yet).
  1. Hours — tune what we have: benchmark trellis_mac 1024_cascade vs fal trellis-2 on identical inputs (BENCHMARKS.md format); route MESHGOD's batch/overnight work to the local lane (m3ultra clears 18 GB ~14× over; fal stays for interactive one-offs). $93/mo → mostly $0.
  2. Days — adopt + fleet-patch: vendor PR #175 / trellis2-apple as monster/trellis2-* Gitea forks (house pattern); run the corridorkey ablation moves on them: attention head_dim → pad to 64 fast-path, mx.compile on fixed shapes, sdpa gating by GPU generation, tiled-vs- full sweeps. Prize: M1 Ultra compatibility (MLX-native, like hunyuan3d — today trellis is m3-only in practice) = 2nd free 3D box.
  3. Weeks — the real kernel work (only if we want fal-class speed): fused Metal sparse-conv (gather-GEMM-scatter) + varlen attention via mx.fast.metal_kernel; port CuMesh hole-filling. Closes most of the 10× gap; genuinely novel, upstreamable to PR #175.
  4. Quality parity misc: raise simplification budget on 256 GB boxes, swap RMBG-2.0 → our licensed bg-remove lane, wire trellis_mac into MESHGOD's local/ model list next to hunyuan3d-mlx.

Sources

  • github.com/microsoft/TRELLIS.2 (setup.sh = dep manifest) · PR #175 · issue #74 · shivampkumar/trellis-mac · pedronaugusto/trellis2-apple
  • ml-explore.github.io/mlx custom-metal-kernels docs
  • Local: HANDOFF_HY3D_MLX.md, CORRIDORKEY.md, BENCHMARKS.md, vendor/trellis-mac/README, vendor/corridorkey-mlx/prompts/ (the 6-phase parity-first port template — the blueprint if we ever do ladder step 3)