Windows/Blackwell (SM 12.0) compatibility: multiple failures preventing shape generation and texturing

[dbollwerk]

I want to be upfront: I'm a 3D Previs Artist, not a software engineer. I've been working through these issues with Claude "vibe coding" I've done my best to document the root causes accurately, but I may have missed nuance in places. If this has any value I'd be happy to provide full diffs of everything we changed if that's useful.

Environment:

OS: Windows 11
GPU: RTX 5090 (Blackwell, SM 12.0 / CC 12.0)
Python: 3.13.12 (embedded)
PyTorch: 2.12.0.dev20260226+cu130
CUDA: 13.0
spconv: latest
flash_attn: not installed (MSVC/CUTLASS incompatible with SM 12.0)
ComfyUI: 0.15.1

Summary:
On Windows with a Blackwell GPU, several incompatibilities prevent the pipeline from running at all. After patching, shape generation works correctly. Texturing runs but output quality is degraded due to the missing grid_sample_3d fallback (see below). This post documents every issue we found and the workarounds applied.

Issue 1: spconv bfloat16 incompatibility

File: trellis2/modules/sparse/conv/conv_spconv.py
Problem: The pipeline runs in bfloat16 on modern GPUs, but spconv's native C++ kernels don't support bfloat16. This causes a KeyError: torch.bfloat16 crash before any generation happens.
Fix applied: Wrapped forward and inverse_forward with a cast:
python

def sparse_conv3d_forward(self, x):
original_dtype = x.feats.dtype
if original_dtype == torch.bfloat16:
x = x.replace(x.feats.to(torch.float16))
out = # ... original forward ...
if original_dtype == torch.bfloat16:
out = out.replace(out.feats.to(original_dtype))
return out

Issue 2: flash_attn mock causes silent garbage attention output

File: trellis2/modules/sparse/attention/full_attn.py
Problem: flash_attn is unavailable on Windows/Blackwell. It gets mocked in sys.modules at startup. However, something in the pipeline sets config.ATTN = 'flash_attn' at runtime after module load (source never identified — may be in pipeline init or model loading). When the flash_attn branch executes, it calls the mock which returns a MagicMock object. The replace() guard in basic.py catches this and returns the original tensor unchanged — meaning qkv [N, 3, H, C] is used as the attention output instead of the attended [N, H, C]. Shape generation produces garbage silently.
Fix applied: Redirected the flash_attn branch to sdpa:
python

elif config.ATTN == 'flash_attn':
out = _sdpa_varlen(q, k, v, q_seqlen, kv_seqlen)

Open question: What is setting config.ATTN = 'flash_attn' at runtime? The default in config.py is already set to 'sdpa' but something overrides it.

Issue 3: Attention and conv backend config defaults

Files: trellis2/modules/sparse/config.py, trellis2/modules/attention/config.py
Problem: 'sdpa' was not in the valid values list for ATTN, causing validation errors. Default values needed to be set explicitly.
Fix applied:

sparse/config.py: set CONV = 'spconv', ATTN = 'sdpa', added 'sdpa' to valid env values
attention/config.py: set BACKEND = 'sdpa'

Issue 4: basic.py SparseTensor crashes when feats is a dict

File: trellis2/modules/sparse/basic.py
Problem: ComfyUI strips custom class wrappers when passing data between nodes. When a SparseTensor passes through the ComfyUI graph, spconv data can arrive as a plain Python dict instead of a C++ spconv object. Accessing .features on a dict crashes.
Fix applied:

feats property: added guard to return _features if it exists and is not None
replace() method: added if not isinstance(feats, torch.Tensor): return self guard
replace() spconv branch: sets both new_data._features = feats and new_data.features = feats to handle both accessor patterns

Issue 5: flex_gemm mock causes silent texture failure (the big one)

Files: trellis2/pipelines/trellis2_image_to_3d.py, nodes.py, trellis2/representations/mesh/base.py
Problem: flex_gemm is unavailable on Windows. It gets mocked in sys.modules at startup. The import guard in each file is:
python

try:
from flex_gemm.ops.grid_sample import grid_sample_3d
HAS_FLEX_GEMM_GRID = True
except Exception:
grid_sample_3d = None
HAS_FLEX_GEMM_GRID = False

Because flex_gemm is already in sys.modules as a MagicMock, the import succeeds — grid_sample_3d becomes a MagicMock, not None. The except branch never fires. Every subsequent if grid_sample_3d is not None: check passes, the mock is called, returns garbage silently, and the texture is near-black with no error.
Fix applied: Added mock detection at import in all three files:
python

try:
from flex_gemm.ops.grid_sample import grid_sample_3d
import inspect
if not inspect.isfunction(grid_sample_3d) and not inspect.isbuiltin(grid_sample_3d):
raise ImportError("grid_sample_3d is mocked")
HAS_FLEX_GEMM_GRID = True
except Exception:
grid_sample_3d = None
HAS_FLEX_GEMM_GRID = False

Current status of fallback: With grid_sample_3d correctly set to None, the texturing pipeline reaches the fallback path. We have been experimenting with sparse trilinear interpolation replacements but have not matched the quality of the original. We're currently seeing ~8-18% hit rate when sampling mesh surface positions against the sparse voxel grid, resulting in degraded texture coverage.

Question:
Is there an official fallback planned for platforms where flex_gemm is unavailable?
What coordinate system does grid_sample_3d expect for the grid parameter? Is it normalized [-1,1] or voxel-space [0, N]?
Is pbr_voxel a surface shell (voxels only at the mesh surface) or a filled volume? We're confused about why ~90% of mesh surface query points fall in empty voxels.
Would a pure F.grid_sample on a downsampled dense grid be a viable fallback, or is the sparse format essential for memory reasons?

Issue 6: Blackwell SM 12.0 compute capability spoofing

File: Custom blackwell_fix.py injected at startup
Problem: Several CUDA kernels fail or are unavailable on SM 12.0 (Blackwell). spconv uses static Ampere kernels that work correctly when CC is spoofed to (8, 6).
Fix applied: Early in init.py, spoof compute capability before any CUDA kernel compilation:
python

import torch.cuda
torch.cuda.get_device_capability = lambda device=None: (8, 6)

Also set SPCONV_DISABLE_JIT=1 and CUMM_DISABLE_JIT=1 environment variables to prevent NVRTC JIT compilation failures.

[Slion]

It sounds like you got it working on Windows with your Blackwell. It would be great if you could provide a link to your working branch?

[dbollwerk]

@Slion not sure my Branch would really help anyone tbh. My folder structure is bespoke and I've hardcoded some patches that might not work on a different setup. I do have the mesh creation working on Blackwell with flex_gemm and flash attention but I still get a mostly black texture using base trellis -2. I'm vibe coding these patches so I don't really understand the ins and outs of everything. Here's a summery of what claude did to patch it.

Following up on this with what we worked through after posting — hoping it saves others some time. Quick context: I'm a 3D previs artist, not a software engineer. All of this was done with Claude doing the actual coding, which is why I framed it the way I did in the original post.

---

What we confirmed and fixed

All six issues in the original post are real and the fixes work as described. A few additions:

Issue 2 & 3 — The env var approach is unreliable. Setting ATTN_BACKEND in the bat file fails because ComfyUI imports the Trellis2 modules before your process environment is fully visible — so __from_env() runs with the wrong state regardless of what the bat sets. The permanent fix is to hardcode directly in the source files:

# trellis2/modules/sparse/config.py
ATTN = 'flash_attn'  # or 'sdpa' if flash_attn unavailable

# trellis2/modules/attention/config.py
BACKEND = 'flash_attn'  # or 'sdpa'

For those who can't get flash_attn installed (our original situation), the sdpa redirect in full_attn.py works fine.

Issue 5 — The mock detection fix is necessary but not sufficient. Detecting the mock and setting grid_sample_3d = None correctly stops the silent garbage output. But if flex_gemm genuinely isn't installed, you still hit the fallback path, which gives degraded texture coverage (~8–18% hit rate). The patch prevents the completely black texture — it doesn't solve the underlying texture quality issue.

For us, texture generation works correctly because we rebuilt the environment around Bruno's bundled cp312/cu128/torch2.8 wheels which include flex_gemm. On cp313/cu130 stacks without those wheels, the fallback path is the ceiling for now.

Critical: do not use flex_gemm as CONV_BACKEND. Even with it installed, setting CONV_BACKEND to flex_gemm crashes the shape SLat decoder with 'SparseConvTensor' object is not subscriptable. The decoder uses spconv internally regardless of what the flow model uses — the format mismatch kills it. spconv as CONV_BACKEND is the only stable path for the standard pipeline.

Hard limit: sparse_structure_resolution max = 64. At 128, spconv hits an int32 overflow in the decoder. This is a known spconv issue requiring a 2.2/2.3 fix upstream. 32³ ≈ 9 min end-to-end, 64³ ≈ 11 min.

Answering the open questions from the original post:

1. grid_sample_3d expects normalized [-1,1] coordinates, not voxel space.
2. pbr_voxel IS a surface shell — voxels only at the mesh surface, not a filled volume. The ~90% miss rate is expected and normal even when the real function is working. Our confirmed hit rate with actual grid_sample_3d was ~10.8%, which is enough for the pipeline to produce output.
3. F.grid_sample on a dense downsampled grid is possible but the sparse format exists for memory reasons — at high resolution a dense grid is prohibitively large.

On the a524247 rollback Fix Blackwell -> does not work commit: worth knowing that Bruno already attempted a Blackwell fix in the repo and rolled it back. The CC spoofing + config hardcoding approach is the only currently working path.

---

Where we ended up

After getting the base pipeline functional, we moved to the Pixal3D workflow (TencentARC). The proj conditioning approach handles the feature extraction differently and produces meaningfully better shape quality for our use case. The Pixal3D env uses flex_gemm natively so the texture issue doesn't apply.

We're now looking at the Trellis2 + multi-view projection path (Bruno's Projection_Blender_Qwen_XViews workflow) as the next step — using projected views to improve texture coverage rather than relying solely on the SLat texture stage. This seems like the more sustainable direction for production use while the upstream issues (spconv int32, flex_gemm fallback quality) remain open.

Happy to share diffs of everything we changed if useful. The CC spoof, bfloat16 cast, config hardcodes, and mock detection are all straightforward patches.

[Slion]

Thanks, meanwhile I found someone actually provide pull requests to get it working on Blackwell so I'm hoping that will do the trick #116. Though it seems the whole thing still needs a lot of work, especially on Blackwell apparently, so I'm not sure when I'll find enough time to get to it myself.