[Apple Silicon] Add MPS and Metal inference support#175
Open
Jourloy wants to merge 24 commits into
Open
Conversation
…iate mx.eval syncs
The default Darwin path used to be a try/except ImportError, which only catches build failures. With the mtlgemm round 1 fixes shipped, the more common failure mode for end users will be running an *older* mtlgemm that still returns CPU tensors from MPS calls — that doesn't fail import, it fails with a cryptic LayerNorm crash inside the model on the first conv. The new probe runs a tiny SparseConv3d on MPS and checks the output device. If anything breaks (import, build, dispatch, return-device), fall back to the pure-PyTorch backend rather than crashing inside the model. Tensors in the probe are built on CPU then moved to MPS because some PyTorch builds lack int/fp16 torch.zeros kernels on MPS — that's a PyTorch issue, separate from anything we control here. Plus: end-to-end smoke test (test_flex_gemm_integration.py) that exercises both Algorithm.IMPLICIT_GEMM and Algorithm.MASKED_IMPLICIT_GEMM through a real SparseConv3d → F.layer_norm chain on MPS. Confirms both algorithms return MPS tensors of the right dtype and produce equivalent output. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The original probe only exercised Algorithm.IMPLICIT_GEMM. A stale install where dense works but the masked cache/dispatch path is broken (e.g. the pre-round-2 aliased-to-dense fallback) would select flex_gemm and crash at the first MASKED_IMPLICIT_GEMM call inside the decoder. Probe both. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Wires the new mtlgemm fused sparse attention kernel through the ATTN backend selector. Dispatches to the fused Metal kernel when max(max_q, max_kv) <= 256 (where the naive per-thread-serial-KV kernel beats SDPA-padded on M3 Max), and falls through to an inline SDPA-padded path for larger max_seqlen. Opt in via ATTN_BACKEND=flex_gemm_sparse_attn or SPARSE_ATTN_BACKEND=flex_gemm_sparse_attn. Default on Darwin stays 'sdpa' — the threshold-based fallback doesn't yet prove a universal win across the pipeline's attention shape distribution. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
SparseGroupNorm / SparseLayerNorm used torch.zeros_like which fails with "DispatchStub: missing kernel for mps" on PyTorch builds compiled with both CUDA and MPS backends (the user's local build hit this). Added a _zeros_like_safe helper that builds zeros on CPU and transfers when the reference tensor is on MPS; Apple Silicon unified memory makes the transfer metadata-only, so the overhead vs a working MPS zeros kernel is negligible. On CPU, behaves identically to torch.zeros_like. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…path precedent The 256-cap inside the flex_gemm_sparse_attn branch was a hold-over from the early naive Metal kernel (round 3). The current backend is flash- attention-v2 with simdgroup matmul + simd-shuffle softmax row reductions and wins at every measured shape including max_seqlen=2048. The CUDA backends above (xformers / flash_attn / flash_attn_3) never fork on max_seqlen — match that precedent. Safety-valve preserved as FLEX_GEMM_ATTN_MAX_SEQLEN=N env var: when set, falls back to SDPA-padded above the cap. Useful only on PyTorch builds where the Accelerate-SDPA-CPU-bounce happens to win at a specific shape (measured crossover sits beyond 768 on fp32, higher on fp16).
End-to-end micro-bench for the production decoder block at res=32 ch=64 with seqlens=[256, 192, 128, 64]. Three stages reported (convs-only / attn-only / combined block) at fp16, vs the all-SDPA-padded baseline. Used to track the cumulative effect of the mtlgemm flash-attention-v2 fwd + bwd work on the actual decoder hot path.
…lex_gemm probe passes The flex_gemm_sparse_attn backend is now flash-attention-v2 with simdgroup_matrix_multiply_accumulate for Q@K^T and P@V plus simd-shuffle softmax row reductions, and wins 5–15× over SDPA-padded-CPU-bounce at every measured shape including max_seqlen=2048. Production decoder block on M3 Max: 5.04× wall-clock vs all-SDPA-padded baseline (3.49 ms vs 17.57 ms), entirely from this change being the default. The existing __flex_gemm_works_on_mps() probe covers the same package the attention path lives in, so the gate is identical: if conv probe passes, set both CONV='flex_gemm' and ATTN='flex_gemm_sparse_attn'. SPARSE_ATTN_BACKEND= (or ATTN_BACKEND=) env override is unchanged. Also fix benchmarks/e2e_decoder.py to inject the repo root into sys.path so it runs as `python benchmarks/e2e_decoder.py` from any cwd.
The mtlgemm Metal extension's metal_context.mm calls at::mps::dispatch_sync_with_rethrow, which was added to that namespace in PyTorch 2.11.0 (pytorch/pytorch#167445, merged 2025-11-11). Earlier stable releases (2.6 – 2.10) only expose it under at::native::mps::, so installing mtlgemm against them fails the C++ extension build. torchvision pinned to >=0.26.0 (matches torch 2.11) to keep the torch/torchvision wheel pair ABI-compatible on resolution.
numpy+stdlib module measuring what validation previously asserted from inputs: weld-by-position, boundary/non-manifold edge counts, directed winding consistency, edge-connected components, per-component signed volume. CLI: python -m trellis2.mesh_integrity file.glb --json. Baseline on known-broken ak74m 12k candidate: 4854 boundary edges, 574 non-manifold, winding inconsistent, 127 inverted components. 1.2M-face raw measures in 1.2 s. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Mirrors the to_glb remesh branch geometry-only, with per-op cleanup ablation (--cleanup-ops) and per-stage integrity metrics. Findings on ak74m body raw (308k faces, res 512): remesh alone reaches 119 boundary edges / winding consistent in 2.4 s; repair_non_manifold_edges REGRESSES the result (500 boundary, winding broken) — excluded from the production chain. dedup + small_components + fill_holes + unify ends at 5 boundary edges, winding consistent. Magazine (1.2M, double-walled) needs band 2+. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
to_glb: remesh=True default; validated cleanup chain (dedup, small_components, fill_holes, unify) — repair_non_manifold_edges excluded (measured regression 119->500 boundary edges); no simplify when decimation_target is None; remesh_project default 0 (project>0 drags DC vertices into dirty source sheets, speckle artifacts). generate_asset: --remesh/--remesh-band/--remesh-project flags, pre-simplification only for non-remesh 200k safety attempts, measured integrity for raw+candidate in meta.json, non-raising promotable gate (remeshed && boundary<=8 && winding_consistent). CPU path warns it cannot remesh. Texture-baking section byte-identical (SHA-verified). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Metal DC remesh is not bytewise deterministic (atomic insertion order); the suite asserts the production contract instead: integrity gates and bounds stable across runs, loose metal-vs-cpu bake parity, CPU path warns it cannot remesh. Metal cases skip cleanly without Metal. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Measured on ak74m body: 0.7 restores sub-voxel sharpness lost by DC remeshing (stamped receiver lines legible again) with identical integrity (boundary 5, winding consistent) and no speckles. Dirty double-walled sources (magazines) must still pass 0 explicitly. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
TRELLIS_FP32_DECODE_THRESHOLDS=1 computes the decoder's to_subdiv logits in fp32 inside the fp16 torso and casts the intersected logits before their hard `> 0` thresholds (upstream microsoft#169). Completes the flag whose meta.json recording and tests were swept into 61fb7cf; until now main lacked the symbols those files import. Default stays off: on the WP9 gate (body seed 137, pipeline 512) raw_full.glb is byte-identical with the flag on (boundary_edges 12609 -> 12609, generation 71.8s -> 72.0s). An instrumented decode shows why: MPS fp16 GEMM error on subdiv logits is real (up to ~0.11) yet flips 0 of ~293k decisions on this seed, and the intersected head already runs in fp32 because the latents are fp32. Raw boundary edges are O-Voxel's by-design open topology (microsoft#15/microsoft#21), resolved by the WP3 remesh path (candidate boundary 5, promotable). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Replaces raw DC projection with a guarded step in postprocess.py: project only where the source is within 1.5 voxels AND local normals agree (|dot|>=0.5, orientation-agnostic — source winding untrusted), then iteratively revert any face the projection flips. Measured on the dirty double-walled magazine raw (band 5, strength 0.7): guarded gives boundary 0 / winding consistent / clean surface; unguarded gives speckle confetti. Projection is now safe by default for ALL sources — including hollow objects that previously required --remesh-project 0. Bench gains --no-guard for A/B; cleanup default restored to the validated chain (no repair_nm). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Author
|
@microsoft-github-policy-service agree |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR adds source-native Apple Silicon support to TRELLIS.2.
On arm64 macOS, the automatic backend selects PyTorch MPS together with capability-probed Metal extensions. The existing CUDA/Linux route remains available, while macOS receives explicit fallbacks for operations that are unavailable or unstable on MPS.
This addresses the Apple Silicon / MPS portion of #74. It intentionally does not close the AMD or Intel parts of that issue.
What changed
flex_gemmMetal sparse convolution and sparse attentionmtldiffrastMetal rasterizationmtlmesh/mtlbvhmesh and BVH operationsflash_attndefault on other platforms.pip check, MPS, SDPA, MLX, KDTree, raster, and BVH probesauto,mps, and experimentalmlxbackend selectionraw_full.glbcandidate_pbr.glbmeta.jsonwith revisions, timings, hashes, bounds, backend probes, and fallback attemptsThe PBR fallback order is:
raw_full.glbis never decimated, and every fallback is recorded inmeta.json.Validation
Tested on:
Automated checks:
28 passedin the pytest suitepip checkandcompileallpassedEnd-to-end generation:
512421024raw_full.glb: 3,090,522 trianglescandidate_pbr.glb: 3,025,626 trianglesThe resulting glTF 2.0 files were validated for non-empty geometry, UV0, vertex normals, consistent bounds, and embedded base-color, metallic-roughness, and alpha textures.
The official RMBG-2.0 path was also tested separately with an opaque input and produced a non-trivial alpha mask.
Compatibility and known limitations
segment_reducethrough its CPU fallback on MPS. This is functional but affects performance.Credits
This work builds on and preserves the authorship of the Apple backend work from:
The CLI diagnostics and fallback design were also informed by: