Metal-native replacement for flex_gemm bottleneck?

[codetiger]

Hey — really appreciate this port. It’s great to have TRELLIS.2 on Apple Silicon.
The README calls out the pure-PyTorch sparse conv in conv_none.py as the main bottleneck (~10x slower than CUDA flex_gemm). Before I start hacking, wanted to check: are you already working on a Metal-native replacement, or is anyone else?

If the path is open, I’d like to take a crack at it. Rough plan: one fused Metal kernel doing neighbor-lookup → gather → tile-GEMM → atomic scatter-add in a single dispatch, called from PyTorch via torch.mps.compile_shader (keeps it a drop-in for conv_none.py, no framework switch). fp16 throughout to match what flex_gemm does. Based on similar gather-matmul kernels, 3-5x on this layer seems realistic.

Happy to open a proper issue or draft PR to discuss the design before writing code. Just didn’t want to duplicate work if you already have something in progress.
Thanks!

[shivampkumar]

Hey, really appreciate you reaching out before diving in. Quick context: I was about to try wiring in mtlgemm by @pedronaugusto (https://github.com/pedronaugusto/mtlgemm) - it's a Metal port of flex_gemm as a C++/Obj-C Metal extension. Just ran into mtlbvh issues on large meshes in a different part of the pipeline, so I haven't confirmed mtlgemm works end-to-end yet.

Two paths:

1. If mtlgemm wires in cleanly, that's probably the least-effort win since it's already built and matches the original flex_gemma API 1:1
2. If it has gaps or perf issues, your torch.mps.compile_shader approach would be great - lighter install, no Xcode toolchain needed by users, in-repo

Would love to coordinate. Want to wait a day while I test mtlgemm? If it ships clean, we'd save your effort for a different gap (there's more - cumesh Metal port also has issues, and sparse attention could use a fused kernel). If mtlgemm falls over, your plan is the path forward.

Open to a PR either way -thanks for asking before writing code, genuinely helpful.

[shivampkumar]

Quick update: ran the tests today. Net: your plan is still the path forward.

mtlgemm end-to-end: builds and installs clean with --no-build-isolation, but sparse_submanifold_conv3d
returns a CPU tensor from MPS inputs. SPARSE_CONV_BACKEND=flex_gemm crashes decode at the next LayerNorm
with RuntimeError: Passed CPU tensor to MPS op, and sampling speeds for shape/texture SLat were within
run-to-run noise - CPU↔MPS bouncing seems to eat any Metal compute win. I kept mtlgemm in setup.sh anyway
because its flex_gemm.ops.grid_sample is the right path for the texture baker (stays on CPU there, clean
output vs. an F.grid_sample fallback that leaves ring artifacts). I was slightly out of my depth, so had to trust some LLMs, so take this with a little pinch of salt.

So conv_none.py is still the diffusion bottleneck. I profiled a 5m 40s fresh-install run (M4 Pro 24GB, cool
machine, pipeline-type 512, 12 sampler steps):

• Sparse structure sampling - 81s / 39% (DiT, SDPA already at ~45 ms/attn call which is the MPS hardware
  limit here, not your target)
• Shape SLat + texture SLat sampling: 34s / 17% (sparse DiT, touches sparse conv)
• Shape + texture decoder VAE forwards: ~38s / 18% (sparse conv heavy, what you were targetting)
• Python flexible_dual_grid_to_mesh: 7s / 4% (confirmed not worth optimizing)

A 3–5× on the sparse conv layer maps to roughly 50–60s off the end-to-end run, which puts it around 4m 30s.
Worth doing.

I also tested reducing sampler steps from 12 to 8 as a cheap win, it saves ~50s but quality drops visibly
(the Nike swoosh on the test shoe shifts from white to peach-cream, features soften). Sticking with 12 so
I'm not trading your upside against quality.

A couple of adjacent gaps if you want a larger scope or want to split the work:

• Sparse attention is currently SDPA-padded, works, but not fused. Same torch.mps.compile_shader idiom
  would apply.
• mtlmesh (the Metal port of cumesh) segfaults on decoder-sized meshes. I stub out fill_holes, simplify,
  remove_faces during decode because of it. That's why output meshes sometimes have small holes. Fused Metal
  mesh-ops would be great too.

Go ahead and open a proper issue or draft PR whenever you're ready. Thanks again for asking before writing
code, genuinely useful.

[pedronaugusto]

Hey Shivam — your bug report was right. All four fixed; mtlgemm and trellis2-apple updated.

Root causes

1. PyBind entries forced inputs to CPU and allocated outputs there — every spconv call was an MPS→CPU→MPS round-trip, the next op got a CPU tensor and crashed. That bandwidth tax is also why your end-to-end timings showed no Metal compute win.
2. Kernels were fp32-only (device float* everywhere) — fp16/bf16 silently misrouted.
3. MASKED_IMPLICIT_GEMM aliased to plain implicit GEMM, dropping the masking arguments without warning.
4. Per-call [cmdbuf waitUntilCompleted] feeding an autotune EMA nothing consumed.

Fixed

• MPS zero-copy. Wrap PyTorch's private MTLBuffer via at::native::mps::getMTLBufferStorage and encode into the active MPSStream through at::mps::dispatch_sync_with_rethrow. Outputs allocated on the input device.
• fp16/bf16 specializations for fwd / bwd-input / bwd-weight. Tile smem narrowed to input dtype, simdgroup accumulator stays fp32 (the recommended Apple Silicon mixed-precision shape).
• Real masked implicit GEMM forward (fp32/fp16/bf16) using the precomputed sorted_idx + valid_kernel + valid_kernel_seg. 26 parity tests vs the dense kernel pass cleanly.
• Skip-empty-V in fwd / bwd-input — per-block atomic OR vote skips the entire Ci loop when no row in the block has a neighbor at v.
• No synchronous waits on the MPS path. PyTorch's stream handles ordering.

trellis2-apple ships a runtime probe in sparse/config.py that falls back to the pytorch backend if the installed mtlgemm pre-dates the device fix — your LayerNorm crash can't recur on stale installs.

Measured (M3 Max, fp16, 27-element kernel volume)

shape	dense	masked	masked vs dense	masked vs old CPU-bounce
res=16 ch=64	0.428 ms	0.205 ms	2.09×	5.74×
res=32 ch=64	0.216 ms	0.213 ms	saturated	5.71×
res=32 ch=128	0.503 ms	0.488 ms	saturated	3.27×
res=64 ch=128	1.993 ms	1.572 ms	1.27×	2.07×
res=64 ch=256	7.339 ms	5.642 ms	1.30×	1.58×
res=128 ch=256	30.05 ms	19.60 ms	1.53×	1.68×

Old CPU-bounce simulates the original behavior (MPS→CPU→MPS per call), which is what your end-to-end timing actually paid.

Open (give me a couple of days)

• Backward-masked GEMM — forward-only is sufficient for inference; bwd is mechanical and just needs its own parity coverage.
• Wide-tile (B2=128) — built, parity-tested, disabled. On M3 Max it regressed at every shape (e.g. res=32 ch=128: 0.50 → 0.67 ms), opposite of the typical NVIDIA pattern. Kernels stay in source for the profiling pass that should precede re-enabling.
• Async prefetch / persistent kernels / fuller autotune / uint64-on-MPS — refinements once the tile-shape question is understood.
• Adjacent: mtlmesh segfaults in fill_holes / simplify / remove_faces — that's what's causing the small holes you see in decoder output (priority 1). Sparse attention as the largest remaining bottleneck per your profile (priority 2). Pattern audit on mtlbvh and mtldiffrast for the same four bug classes above.

I will do a sweep across my repos in an effort to improve their performance and correctness, a lot of this was done with Claude Code and I just try to do my best to guide it with my limited knowledge when i have some time.

— Pedro

[shivampkumar]

Pedro this is cool! Pulled the latest mtlgemm + o-voxel (0ef887d / 69b5622), stripped my profile instrumentation for a clean same-conditions comparison, ran it twice with and without to cross-check.

Fresh-install M4 Pro 24GB, cool machine, pipeline-type 512, 12 steps, weights cached:

	conv_none (before)	flex_gemm (your fix)	Δ
Total wall-clock	5m 40.8s	5m 13.3s	−27.5s
Generation time	220.9s	191.9s	−29s
Pipeline load	104s	103s	noise
Bake	11s	14s	noise
Quality	reference	matches reference (clean side-view swoosh, no ring artifacts)	—

On matched profile runs isolating the decoder path:

	conv_none	flex_gemm	Speedup
Shape VAE forward	17.92s	7.24s	2.48×
Tex VAE forward	19.94s	6.86s	2.91×

2.5–2.9× on the decoders matches the ratio in your bench table. The zero-copy fix is doing the work - sparse
conv outputs stay on mps:0, and the rest of the pipeline never has to pay the CPU→MPS hop.

Repo default is now SPARSE_CONV_BACKEND=flex_gemm (291a116) with a try/except fallback to conv_none if flex_gemm isn't importable. README updated with the new numbers and credits your fix.

On your priorities list: sparse attention is now the largest single chunk at 81s, about 42% of pipeline.run() or 26% of the full 5m 13s run. That's the DiT on the 32³ sparse structure latent, using an SDPA-padded wrapper; SDPA itself is at ~45 ms/attn call which is where MPS caps out on my machine, so a fused Metal attention kernel is probably where the next big unlock lives. mtlmesh fill_holes / simplify / remove_faces still stubbed out in our decode path, which is why output has small holes.

Note: I'm not a Metal expert. Have been relying on claude code/reddit/stackoverflow for building a lot of my understanding. Happy to keep round-tripping.

Thanks again.

[shivampkumar]

@codetiger

Thanks to @pedronaugusto, repo default is now SPARSE_CONV_BACKEND=flex_gemm (291a116).

That means the conv_none drop-in replacement is less urgent. The two remaining gaps are probably a better fit for your torch.mps.compile_shader approach:

1. Sparse attention is now the largest single chunk - 81s of the 5m 13s run, 42% of pipeline.run(), 26% of the full run. It's a DiT on the 32³ latent going through an SDPA-padded wrapper. SDPA on MPS is at ~45 ms/attn call which is the hardware limit here, so a fused attention kernel is where the next big unlock probably is.

2. mtlmesh (cumesh Metal port) segfaults on decoder-sized meshes. fill_holes / simplify / remove_faces are all stubbed in our decode path, which is why output meshes sometimes have small holes. A fused Metal mesh-ops pass would be a great second target.

Happy to have you open specific issues and work on it!

I genuinely believe all of us coming together is the perfect opportunity to get Trellis awesome on silicon, and create a one command solution for anyone porting these CUDA deps in the future.

[pedronaugusto]

Hey Shivam — picking up on the two items from your follow-up.

mtlmesh segfaults

fill_holes / simplify / remove_faces are fixed in 98047ac (O(F) bounds checks in each op). unify_face_orientations infinite-loop also capped at 64 iters. You can un-stub them in your decode path; that's what was producing the small holes in the GLBs.

Sparse attention as the largest remaining chunk (81s / 42%)

mtlgemm now ships a fused Metal flash-attention-v2 (fwd + bwd, all dtypes) — simdgroup_matrix_multiply_accumulate for both Q@K^T and P@V plus simd-shuffle row reductions for online softmax. Routes when head_dim ≤ 64 and multiple of 8, which fits the trellis2 DiT.

Plugging it into your fork is a small addition to patch_sparse_attention() in mps_compat.py — same shape as your existing sdpa branch, plus 'flex_gemm_sparse_attn' in the allowed-backend list. The equivalent branch already lives at pedronaugusto/trellis2-apple/trellis2/modules/sparse/attention/full_attn.py if a reference is useful. Happy to PR your repo if that helps.

Other increments since round 3

• mtlbvh — opt-in output_dtype=torch.{float16,bfloat16} on the three query entry points; default None → fp32 (bit-exact backward-compat, your existing call sites unaffected).
• mtldiffrast — rasterize now returns MPS-native outputs via MTLBlitCommandEncoder copyFromTexture:toBuffer: (replaces a getBytes: call that triggered Obj-C runtime corruption when the target tensor was MPS-allocated). Async-dispatch helpers across all four kernel files.
• trellis2-apple also has an MLX-native backend (mlx_backend/ + api_server.py) — full DiT + SLAT + structure decoder + VAE + DINOv3 in MLX, with the texture-baker still routing through flex_gemm conv. Worth flagging since MLX wins at single-call dense attention: at the trellis2 sparse decoder shape (1, H=8, N=640, D=64, fp16), mx.fast.scaled_dot_product_attention clocks 0.648 ms vs our flex_gemm flash-attn's 1.442 ms. Haven't bench'd full pipeline.run() head-to-head between the two backends yet.

Tested

gate	result
mtlgemm sparse-attn parity (24 fwd, 24 bwd × 3 dtypes, 44 masked-fwd, 21 masked-bwd)	113/113 pass; max fp16 bwd err vs CPU reference 5.95e-05
mtlbvh dtype parity	14/14 pass; narrow output bit-exact equals fp32 result cast to narrow dtype
mtldiffrast MPS↔CPU output parity	5/5 bit-exact at fp32
sparse-attn output parity vs SDPA at trellis2 decoder shapes (T=640, H=8, C=64, fp16, [256,192,128,64])	max abs err 6.1e-05 (~2 fp16 ULPs); output ranges identical to 4 decimal places

Decoder-block microbench (trellis2-apple/benchmarks/e2e_decoder.py, M3 Max, fp16, fresh PyTorch 2.11 venv): attn-only 5.4× over SDPA-padded, combined block 3.3×. Microbench, not a full pipeline.run() — your re-profile is the real number.

PyTorch pin

mtlgemm's metal_context.mm calls at::mps::dispatch_sync_with_rethrow, added to the at::mps:: namespace in PyTorch 2.11.0 (pytorch/pytorch#167445). Older stable PyTorch (2.6 – 2.10) only exposes it under at::native::mps::, so a fresh pip install against pre-2.11 fails the C++ extension build. Pinned torch>=2.11.0 in mtlgemm/pyproject.toml and trellis2-apple/requirements_macos.txt. Your setup.sh does plain pip install torch so resolves to latest stable already.

Let me know if anything looks off, i might have messed something up, not exactly doing things by the book on this, just Clauding a lot. Happy to keep collaborating.

[lastowl]

Hi, I've been working on a fused Metal attention kernel, and saw your note in here that the SDPA-padded sparse attention is the biggest remaining chunk now that flex_gemm landed. Wanted to see if it might be useful here.

It's a small library, mtlattn (https://github.com/lastowl/mtlattn), fused varlen flash attention for PyTorch MPS, forward only. I just put it on PyPI so it's a plain pip install mtlattn with prebuilt Apple Silicon wheels, no Xcode/Metal toolchain needed by users, which I think was one of the things you were wary of.

For the sparse dispatch it'd slot in as an opt-in backend (SPARSE_ATTN_BACKEND=mtlattn), same shape as the sdpa one: build cu_seqlens from the seqlens and pass the ragged q/k/v straight in, no padding. A few things it changes vs the padded path:

• It doesn't build the [B, H, max, max] score tensor, so the big sparse blocks that currently run out of memory just run. At 32K tokens that tensor is around 69GB in fp32, which is roughly why the larger configs fall over.
• The sdpa branch runs without a mask, so in a ragged batch the short sequences end up attending the zero padding. mtlattn does it per-sequence via cu_seqlens like the xformers/flash_attn backends. On a ragged test batch I get about 4e-4 max error against a per-sequence reference, where the padded path is off by ~1.4 on the short ones.
• Speed depends on the machine and I don't want to oversell it. I'm on an M5 where it uses the Metal 4 matmul2d path and comes out faster, but on M1-M4 it falls back to a simdgroup kernel that's roughly on par with SDPA. You're right that SDPA is about at the MPS limit there, so on those machines the real value is the memory headroom and the masking rather than speed. Didn't want to quote you a number you wouldn't see on the M4.

Happy to put together a PR if it's useful, but didn't want to assume — figured I'd see if it's something you'd want first.