Fix real OOM: 4-D SDPA layout; cu130 is ZeroGPU-native, drop cu128 cruft

GPU-DIAG showed torch 2.11.0+cu130 persists and diff-gaussian compiles
against the system CUDA 13.0 toolkit -- i.e. ZeroGPU natively ships
cu130 + a matching CUDA 13.0 toolkit. So cu130 was never the problem
and the cu128/nvcc machinery was unnecessary: remove the pip-nvcc
CUDA_HOME logic and the nvidia-cuda-nvcc/cccl requirements; compile
against the matching system toolkit.

The actual failure (42 GiB free yet OOM at first attention) was the
SDPA call receiving 3-D tensors, which forces SDPA's math backend to
materialize the full multi-view attention matrix (tens of GB). Reshape
q/k/v to 4-D (batch, heads, seq, head_dim) so SDPA uses its flash /
memory-efficient backend -- the efficient path xformers used to provide.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -99,41 +99,10 @@ except ImportError:
     print("Installing diff-gaussian-rasterization (compiling for detected CUDA arch)...")
     env = os.environ.copy()
 
-    # Compile against the pip-installed CUDA 12.8 toolkit (nvidia-cuda-nvcc-cu12),
-    # NOT the container's CUDA 13.0 toolkit. torch here is the cu128 build ZeroGPU
-    # requires, and torch's build check errors on a major CUDA mismatch (nvcc 13 vs
-    # torch 12.8). Point CUDA_HOME at the pip nvcc and feed the runtime/cub headers
-    # and cudart lib via CPATH/LIBRARY_PATH so nvcc 12.8 matches torch 12.8.
-    import site, glob
-    _search_dirs = list(site.getsitepackages())
-    try:
-        _search_dirs.append(site.getusersitepackages())
-    except Exception:
-        pass
-    _nvcc_bins = []
-    for _d in _search_dirs:
-        _nvcc_bins += glob.glob(os.path.join(_d, "nvidia", "cuda_nvcc*", "bin", "nvcc"))
-    if _nvcc_bins:
-        _nvcc_home = os.path.dirname(os.path.dirname(_nvcc_bins[0]))  # .../nvidia/cuda_nvcc
-        _nv = os.path.dirname(_nvcc_home)                              # .../nvidia
-        env["CUDA_HOME"] = _nvcc_home
-        env["PATH"] = os.path.join(_nvcc_home, "bin") + os.pathsep + env.get("PATH", "")
-        _incs = [p for p in (os.path.join(_nv, "cuda_runtime", "include"),
-                             os.path.join(_nv, "cuda_cccl", "include")) if os.path.isdir(p)]
-        if _incs:
-            env["CPATH"] = os.pathsep.join(_incs + ([env["CPATH"]] if env.get("CPATH") else []))
-        _rt_lib = os.path.join(_nv, "cuda_runtime", "lib")
-        if os.path.isdir(_rt_lib):
-            env["LIBRARY_PATH"] = _rt_lib + os.pathsep + env.get("LIBRARY_PATH", "")
-            env["LD_LIBRARY_PATH"] = _rt_lib + os.pathsep + env.get("LD_LIBRARY_PATH", "")
-        print(f"[BUILD] using pip CUDA toolkit: CUDA_HOME={_nvcc_home}")
-    else:
-        print("⚠️  pip CUDA nvcc (nvidia-cuda-nvcc-cu12) not found under site-packages/nvidia; "
-              "falling back to system nvcc (will mismatch cu128 torch). "
-              "If this persists, Factory Reboot the Space to clear the stale build cache.")
-
-    # Build a cross-arch set incl. Blackwell sm_120 (the ZeroGPU GPU is cap 12.0);
-    # the build stage may not see a GPU, so don't rely on get_device_capability.
+    # ZeroGPU ships a cu130 torch and a matching CUDA 13.0 toolkit, so compile
+    # against the system toolkit (no version mismatch). Include Blackwell sm_120
+    # (the ZeroGPU GPU is capability 12.0); the build stage may not see a GPU, so
+    # use a fixed cross-arch set rather than get_device_capability.
     env["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6;8.9;9.0;12.0+PTX"
 
     # --no-build-isolation: the package's setup.py imports torch to detect the

mvdiffusion/models/transformer_mv2d_image.py

@@ -725,20 +725,25 @@ class MVAttnProcessor:
                 # value = rearrange(values, 'b t f d c -> (b t) (f d) c')
 
 
-        query = attn.head_to_batch_dim(query).contiguous()
-        key = attn.head_to_batch_dim(key).contiguous()
-        value = attn.head_to_batch_dim(value).contiguous()
+        # Memory-efficient attention via PyTorch SDPA. Reshape to 4-D
+        # (batch, heads, seq, head_dim): SDPA only dispatches to its flash /
+        # memory-efficient backends for 4-D inputs -- with the 3-D head_to_batch_dim
+        # layout it falls back to the math backend, which materializes the full
+        # multi-view attention matrix (tens of GB at 512-res) and OOMs. xformers was
+        # the original efficient path but doesn't support the Blackwell GPU.
+        # reshape (not view): the multi-view key/value are non-contiguous.
+        bsz = query.shape[0]
+        head_dim = query.shape[-1] // attn.heads
+        query = query.reshape(bsz, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.reshape(bsz, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.reshape(bsz, -1, attn.heads, head_dim).transpose(1, 2)
 
-        # Memory-efficient attention via PyTorch SDPA. Mathematically equivalent to
-        # the get_attention_scores (baddbmm) + softmax + bmm full-attention path, but
-        # it never materializes the full attention matrix -- that full matrix OOMs on
-        # multi-view 512-res inputs. xformers' kernels were the original efficient
-        # path but don't support the Blackwell GPU, so use SDPA (flash/mem-efficient
-        # backend), which does and runs on Blackwell.
         hidden_states = F.scaled_dot_product_attention(
             query, key, value, attn_mask=attention_mask, scale=attn.scale
         )
-        hidden_states = attn.batch_to_head_dim(hidden_states)
+        hidden_states = hidden_states.transpose(1, 2).reshape(
+            bsz, -1, attn.heads * head_dim
+        ).to(query.dtype)
 
         # linear proj
         hidden_states = attn.to_out[0](hidden_states)

requirements.txt

@@ -1,13 +1,9 @@
 # torch 2.11.0 is ZeroGPU-supported and the version line where torchvision 0.26.0
 # and xformers 0.0.35 agree (xformers 0.0.33 hard-pinned torch==2.9.0, which ZeroGPU
-# rejects; 0.0.35 relaxed to >=2.10). This is the cu128 (CUDA 12.8) build ZeroGPU
-# ships -- required for ZeroGPU's GPU memory management to work.
+# rejects; 0.0.35 relaxed to >=2.10). ZeroGPU provides the cu130 build + a matching
+# CUDA 13.0 toolkit, so diff-gaussian-rasterization compiles against the system toolkit.
 torch==2.11.0
 torchvision==0.26.0
-# CUDA 12.8 compiler + headers so diff-gaussian-rasterization compiles against a
-# toolkit matching cu128 torch (the container only ships CUDA 13.0). See app.py.
-nvidia-cuda-nvcc-cu12~=12.8.0
-nvidia-cuda-cccl-cu12~=12.8.0
 numpy==1.26.4
 pillow==10.4.0
 diffusers[torch]==0.30.3