NAF upsampler crashes with 'CUDA driver error: device not ready' after sampling stages (WSL2, RTX 3090)

[Lulu-IMG]

Bug: NAF upsampler crashes with "CUDA driver error: device not ready" during shape/texture conditioning (WSL2, RTX 3090)

Summary

Inference fails reproducibly with RuntimeError: CUDA driver error: device not ready inside the NAF (Neighborhood Attention Field) upsampler's forward_encoder, specifically at the conv1 layer immediately following torch.cat([self.encoder(x), self.sem_encoder(x)], dim=1). The crash occurs only on conditioning passes where use_naf_upsample=True (i.e. shape_512, shape_1024, tex_1024), never on the ss pass (which has no NAF). It happens consistently at the same point on every run, immediately after the sparse-structure and shape-SLat diffusion sampling stages complete successfully.

Environment

• OS: Windows 11, running Ubuntu 22.04 under WSL2
• GPU: NVIDIA RTX 3090 (24GB VRAM), Ampere (sm_86)
• NVIDIA driver: 610.43.02, CUDA UMD 13.3 (per nvidia-smi)
• CUDA Toolkit installed in WSL: 12.4 (nvcc confirmed working)
• Python: 3.10 (conda env)
• PyTorch: 2.6.0+cu124
• flash-attn: 2.7.3 (built from source, working)
• NATTEN: 0.21.6 (built from source with NATTEN_CUDA_ARCH=8.6, working in isolation — see below)
• flex_gemm / cumesh / o_voxel: all built successfully from the o-voxel install, autotuning runs and completes without error
• nvdiffrast: 0.4.0, working
• Model weights: local checkpoint folder (not HF Hub), pipeline.json-driven load via from_pretrained(local_path), all ckpts/*.safetensors present and loading correctly

Steps to reproduce

1. Set up TRELLIS.2 base env per its README (setup.sh --new-env --basic --flash-attn --nvdiffrast), plus --o-voxel for cumesh/flex_gemm/o_voxel (these names differ from the README's roadmap list, fyi — actual flags are --cumesh --o-voxel --flexgemm --nvdiffrast --nvdiffrec, not --spconv --kaolin etc., which may be worth a docs fix).
2. Build NATTEN from source: NATTEN_CUDA_ARCH="8.6" NATTEN_N_WORKERS=4 pip install natten==0.21.6 --no-build-isolation.
3. Clone Pixal3D, pip install -r requirements.txt.
4. Run:

   python inference.py --image <any image> --output ./out.glb --model_path <local checkpoint dir>
   

5. Pipeline loads, MoGe-2 camera estimation succeeds, sparse structure sampling completes (12/12), shape SLat sampling completes (12/12), then crashes on the very next step (get_proj_cond_shape → image_cond_model.forward → self.naf_model(...)).

Full traceback

Sampling sparse structure (proj): 100%|██████████| 12/12
Sampling shape SLat (proj): 100%|██████████| 12/12
Traceback (most recent call last):
  File "inference.py", line 304, in <module>
    run_inference(...)
  File "inference.py", line 246, in run_inference
    mesh_list, (shape_slat, tex_slat, res) = pipeline.run(...)
  File ".../pixal3d_image_to_3d.py", line 730, in run
    cond_shape_hr = self.get_proj_cond_shape(...)
  File ".../pixal3d_image_to_3d.py", line 271, in get_proj_cond_shape
    z_global, z_proj = image_cond_model(...)
  File ".../image_conditioned_proj.py", line 541, in forward
    hr_features = self.naf_model(
        image_for_naf, lr_features_bchw, self.naf_target_size
    )
  File ".../valeoai_NAF_main/src/model/naf.py", line 105, in forward
    x = self.image_encoder(image, output_size=output_size)
  File ".../valeoai_NAF_main/src/model/naf.py", line 50, in forward
    x = self.forward_encoder(x, o_size)
  File ".../valeoai_NAF_main/src/model/naf.py", line 33, in forward_encoder
    x = torch.cat([self.encoder(x), self.sem_encoder(x)], dim=1)
  File ".../valeoai_NAF_main/src/layers/convolutions.py", line 56, in forward
    x = self.conv1(x)
  File ".../torch/nn/modules/conv.py", line 538, in _conv_forward
    return F.conv2d(...)
RuntimeError: CUDA driver error: device not ready

Diagnostic work already done

I want to save maintainers time, so here's everything I've already ruled out:

1. Not VRAM exhaustion. torch.cuda.mem_get_info() immediately before the crash shows ~24.4GB free out of 24.5GB total.
2. Not a general WSL2/CUDA/conv2d incompatibility. A plain torch.nn.Conv2d on a fresh tensor, run in total isolation (no pipeline state at all), succeeds without issue.
3. Not caused by loading the pipeline itself. Loading the full pipeline (all 4 DinoV3 models on GPU, NAF model loaded, flash-attn/NATTEN/flex_gemm all imported) and then immediately running a plain conv2d afterward — without running any sampling — succeeds fine.
4. Requires the sampling stage to have run first. The crash only manifests after the sparse-structure and shape-SLat sampling stages (which exercise flex_gemm's custom sparse convolution kernels and their autotuner) have executed. This strongly suggests something in flex_gemm's kernel execution leaves the CUDA context/allocator in a state that NAF's plain conv2d then trips on.
5. NATTEN in isolation does not reproduce it. Running a NATTEN na2d op directly, followed immediately by a plain conv2d, both succeed cleanly — so NATTEN alone is not the trigger.
6. The CUDA context is NOT permanently wedged afterward. Catching the RuntimeError and immediately running a fresh, unrelated conv2d in the same process succeeds. This indicates the corruption is transient/contextual rather than a hard device failure.
7. A naive catch-and-retry-once does NOT work. Wrapping the failing self.naf_model(...) call in a try/except RuntimeError, calling torch.cuda.synchronize(), and retrying the same call immediately escalates to a worse, deeper failure:

   RuntimeError: !handles_.at(i) INTERNAL ASSERT FAILED at "/pytorch/c10/cuda/CUDACachingAllocator.cpp":400, please report a bug to PyTorch.
   

   This suggests the first failure leaves the allocator's internal event/stream bookkeeping in a bad state that a simple sync doesn't repair, and that retrying the same already-partially-executed model call doesn't get a clean slate even though a different, fresh op would.
8. Forcing use_naf_upsample=False is not a viable workaround. The shape/texture diffusion models were trained expecting NAF's doubled channel output (proj_channels = embed_dim*2 if use_naf_upsample else embed_dim), so disabling NAF breaks the cross-attention layer shapes downstream (mat1 and mat2 shapes cannot be multiplied (1964x1024 and 2048x1536)).
9. Setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False,garbage_collection_threshold:0.6 (overriding the script's own expandable_segments:True default) had no effect — identical crash, same line.
10. Restarting WSL entirely (wsl --shutdown) between attempts does not help — the crash reproduces identically after a full WSL restart, ruling out simple driver/session staleness.

Suspected root cause

Best guess based on the above: an interaction between flex_gemm's custom Triton-based sparse convolution kernels (used during sparse-structure/shape-SLat sampling) and the WSL2 CUDA virtualization layer, which leaves some piece of CUDA stream or event state in a condition that a subsequent "plain" PyTorch conv2d call (inside NAF, a completely separate, vanilla-PyTorch model with no custom kernels of its own) intermittently/consistently trips on. This may be specific to WSL2's GPU passthrough rather than native Linux, since I have not seen this reported elsewhere, but I have not been able to test on bare-metal Linux to confirm.

Question for maintainers

• Has this been tested on WSL2 specifically, as opposed to native Linux?
• Is there a known-safe way to force a full CUDA context/stream reset between the sampling stage and the NAF-dependent conditioning stage (e.g. torch.cuda.synchronize() + something stronger than empty_cache()) that wouldn't require reloading the whole pipeline?
• Could NAF's encoder be swapped to use o_voxel/flex_gemm kernels consistently rather than mixing in a vanilla PyTorch conv2d, if that's indeed the source of the conflict?

Happy to provide more logs, run additional diagnostic scripts, or test proposed fixes — I have a full reproducible environment set up.

MD copy