Fallback TRELLIS attention backends to SDPA when flash-attn missing

app.py

@@ -1,5 +1,6 @@
 import argparse
 import concurrent.futures
+import importlib.util
 import os
 import sys
 import time
@@ -8,7 +9,14 @@ from typing import Any, Dict, Generator, List, Optional, Tuple
 
 os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
-os.environ.setdefault("ATTN_BACKEND", "flash_attn_3")
+if importlib.util.find_spec("flash_attn") is not None:
+    _attn_backend = "flash_attn"
+elif importlib.util.find_spec("xformers") is not None:
+    _attn_backend = "xformers"
+else:
+    _attn_backend = "sdpa"
+os.environ.setdefault("ATTN_BACKEND", _attn_backend)
+os.environ.setdefault("SPARSE_ATTN_BACKEND", _attn_backend)
 os.environ.setdefault("SPCONV_ALGO", "native")
 os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(
     os.path.dirname(os.path.abspath(__file__)), "autotune_cache.json"

trellis/modules/attention/__init__.py

@@ -1,6 +1,6 @@
 from typing import *
 
-BACKEND = 'flash_attn' 
+BACKEND = 'sdpa'
 DEBUG = False
 
 def __from_env():
@@ -12,6 +12,9 @@ def __from_env():
     env_attn_backend = os.environ.get('ATTN_BACKEND')
     env_sttn_debug = os.environ.get('ATTN_DEBUG')
     
+    if env_attn_backend == 'flash_attn_3':
+        env_attn_backend = 'flash_attn'
+
     if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'sdpa', 'naive']:
         BACKEND = env_attn_backend
     if env_sttn_debug is not None:
@@ -25,7 +28,7 @@ def __from_env():
 __from_env()
     
 
-def set_backend(backend: Literal['xformers', 'flash_attn']):
+def set_backend(backend: Literal['xformers', 'flash_attn', 'sdpa', 'naive']):
     global BACKEND
     BACKEND = backend
 

trellis/modules/attention/full_attn.py

@@ -1,12 +1,22 @@
 from typing import *
 import torch
 import math
-from . import DEBUG, BACKEND
+from . import DEBUG, BACKEND, set_backend
 
 if BACKEND == 'xformers':
-    import xformers.ops as xops
+    try:
+        import xformers.ops as xops
+    except ImportError:
+        BACKEND = 'sdpa'
+        set_backend(BACKEND)
+        from torch.nn.functional import scaled_dot_product_attention as sdpa
 elif BACKEND == 'flash_attn':
-    import flash_attn
+    try:
+        import flash_attn
+    except ImportError:
+        BACKEND = 'sdpa'
+        set_backend(BACKEND)
+        from torch.nn.functional import scaled_dot_product_attention as sdpa
 elif BACKEND == 'sdpa':
     from torch.nn.functional import scaled_dot_product_attention as sdpa
 elif BACKEND == 'naive':

trellis/modules/sparse/__init__.py

@@ -2,7 +2,7 @@ from typing import *
 
 BACKEND = 'spconv' 
 DEBUG = False
-ATTN = 'flash_attn'
+ATTN = 'sdpa'
 
 def __from_env():
     import os
@@ -21,7 +21,10 @@ def __from_env():
         BACKEND = env_sparse_backend
     if env_sparse_debug is not None:
         DEBUG = env_sparse_debug == '1'
-    if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn']:
+    if env_sparse_attn == 'flash_attn_3':
+        env_sparse_attn = 'flash_attn'
+
+    if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn', 'sdpa', 'naive']:
         ATTN = env_sparse_attn
         
     print(f"[SPARSE] Backend: {BACKEND}, Attention: {ATTN}")
@@ -38,7 +41,7 @@ def set_debug(debug: bool):
     global DEBUG
     DEBUG = debug
 
-def set_attn(attn: Literal['xformers', 'flash_attn']):
+def set_attn(attn: Literal['xformers', 'flash_attn', 'sdpa', 'naive']):
     global ATTN
     ATTN = attn
     

trellis/modules/sparse/attention/full_attn.py

@@ -1,12 +1,23 @@
 from typing import *
 import torch
+from torch.nn.functional import scaled_dot_product_attention as sdpa
 from .. import SparseTensor
-from .. import DEBUG, ATTN
+from .. import DEBUG, ATTN, set_attn
 
 if ATTN == 'xformers':
-    import xformers.ops as xops
+    try:
+        import xformers.ops as xops
+    except ImportError:
+        ATTN = 'sdpa'
+        set_attn(ATTN)
 elif ATTN == 'flash_attn':
-    import flash_attn
+    try:
+        import flash_attn
+    except ImportError:
+        ATTN = 'sdpa'
+        set_attn(ATTN)
+elif ATTN in {'sdpa', 'naive'}:
+    pass
 else:
     raise ValueError(f"Unknown attention module: {ATTN}")
 
@@ -16,6 +27,14 @@ __all__ = [
 ]
 
 
+def _sdpa_chunk(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    q = q.permute(1, 0, 2).unsqueeze(0)  # [1, H, Lq, C]
+    k = k.permute(1, 0, 2).unsqueeze(0)  # [1, H, Lk, C]
+    v = v.permute(1, 0, 2).unsqueeze(0)  # [1, H, Lk, C]
+    out = sdpa(q, k, v)
+    return out.squeeze(0).permute(1, 0, 2)  # [Lq, H, C]
+
+
 @overload
 def sparse_scaled_dot_product_attention(qkv: SparseTensor) -> SparseTensor:
     """
@@ -206,6 +225,31 @@ def sparse_scaled_dot_product_attention(*args, **kwargs):
             out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
         elif num_all_args == 3:
             out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+    elif ATTN in {'sdpa', 'naive'}:
+        outs = []
+        q_start, kv_start = 0, 0
+        for q_len, kv_len in zip(q_seqlen, kv_seqlen):
+            if num_all_args == 1:
+                qkv_chunk = qkv[q_start:q_start + q_len]
+                q_i, k_i, v_i = qkv_chunk.unbind(dim=1)
+            elif num_all_args == 2:
+                q_i = q[q_start:q_start + q_len]
+                kv_chunk = kv[kv_start:kv_start + kv_len]
+                k_i, v_i = kv_chunk.unbind(dim=1)
+            else:
+                q_i = q[q_start:q_start + q_len]
+                k_i = k[kv_start:kv_start + kv_len]
+                v_i = v[kv_start:kv_start + kv_len]
+
+            outs.append(_sdpa_chunk(q_i, k_i, v_i))
+            q_start += q_len
+            kv_start += kv_len
+        if outs:
+            out = torch.cat(outs, dim=0)
+        elif num_all_args == 1:
+            out = torch.empty((0, qkv.shape[-2], qkv.shape[-1]), device=device, dtype=qkv.dtype)
+        else:
+            out = torch.empty((0, q.shape[-2], q.shape[-1]), device=device, dtype=q.dtype)
     else:
         raise ValueError(f"Unknown attention module: {ATTN}")
     

trellis/modules/sparse/attention/serialized_attn.py

@@ -2,13 +2,24 @@ from typing import *
 from enum import Enum
 import torch
 import math
+from torch.nn.functional import scaled_dot_product_attention as sdpa
 from .. import SparseTensor
-from .. import DEBUG, ATTN
+from .. import DEBUG, ATTN, set_attn
 
 if ATTN == 'xformers':
-    import xformers.ops as xops
+    try:
+        import xformers.ops as xops
+    except ImportError:
+        ATTN = 'sdpa'
+        set_attn(ATTN)
 elif ATTN == 'flash_attn':
-    import flash_attn
+    try:
+        import flash_attn
+    except ImportError:
+        ATTN = 'sdpa'
+        set_attn(ATTN)
+elif ATTN in {'sdpa', 'naive'}:
+    pass
 else:
     raise ValueError(f"Unknown attention module: {ATTN}")
 
@@ -18,6 +29,21 @@ __all__ = [
 ]
 
 
+def _sdpa_varlen_qkv(qkv_feats: torch.Tensor, seq_lens: List[int]) -> torch.Tensor:
+    outs = []
+    start = 0
+    for seq_len in seq_lens:
+        chunk = qkv_feats[start:start + seq_len]
+        q, k, v = chunk.unbind(dim=1)
+        q = q.permute(1, 0, 2).unsqueeze(0)
+        k = k.permute(1, 0, 2).unsqueeze(0)
+        v = v.permute(1, 0, 2).unsqueeze(0)
+        out = sdpa(q, k, v).squeeze(0).permute(1, 0, 2)
+        outs.append(out)
+        start += seq_len
+    return torch.cat(outs, dim=0) if outs else qkv_feats.new_empty((0, qkv_feats.shape[2], qkv_feats.shape[3]))
+
+
 class SerializeMode(Enum):
     Z_ORDER = 0
     Z_ORDER_TRANSPOSED = 1
@@ -168,6 +194,8 @@ def sparse_serialized_scaled_dot_product_self_attention(
             out = xops.memory_efficient_attention(q, k, v)          # [B, N, H, C]
         elif ATTN == 'flash_attn':
             out = flash_attn.flash_attn_qkvpacked_func(qkv_feats)   # [B, N, H, C]
+        elif ATTN in {'sdpa', 'naive'}:
+            out = _sdpa_varlen_qkv(qkv_feats.reshape(B * N, 3, H, C), [N] * B)
         else:
             raise ValueError(f"Unknown attention module: {ATTN}")
         out = out.reshape(B * N, H, C)                              # [M, H, C]
@@ -183,6 +211,10 @@ def sparse_serialized_scaled_dot_product_self_attention(
             cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
                         .to(qkv.device).int()
             out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
+        elif ATTN in {'sdpa', 'naive'}:
+            out = _sdpa_varlen_qkv(qkv_feats, seq_lens)
+        else:
+            raise ValueError(f"Unknown attention module: {ATTN}")
 
     out = out[bwd_indices]      # [T, H, C]
 

trellis/modules/sparse/attention/windowed_attn.py

@@ -1,13 +1,24 @@
 from typing import *
 import torch
 import math
+from torch.nn.functional import scaled_dot_product_attention as sdpa
 from .. import SparseTensor
-from .. import DEBUG, ATTN
+from .. import DEBUG, ATTN, set_attn
 
 if ATTN == 'xformers':
-    import xformers.ops as xops
+    try:
+        import xformers.ops as xops
+    except ImportError:
+        ATTN = 'sdpa'
+        set_attn(ATTN)
 elif ATTN == 'flash_attn':
-    import flash_attn
+    try:
+        import flash_attn
+    except ImportError:
+        ATTN = 'sdpa'
+        set_attn(ATTN)
+elif ATTN in {'sdpa', 'naive'}:
+    pass
 else:
     raise ValueError(f"Unknown attention module: {ATTN}")
 
@@ -17,6 +28,21 @@ __all__ = [
 ]
 
 
+def _sdpa_varlen_qkv(qkv_feats: torch.Tensor, seq_lens: List[int]) -> torch.Tensor:
+    outs = []
+    start = 0
+    for seq_len in seq_lens:
+        chunk = qkv_feats[start:start + seq_len]
+        q, k, v = chunk.unbind(dim=1)
+        q = q.permute(1, 0, 2).unsqueeze(0)
+        k = k.permute(1, 0, 2).unsqueeze(0)
+        v = v.permute(1, 0, 2).unsqueeze(0)
+        out = sdpa(q, k, v).squeeze(0).permute(1, 0, 2)
+        outs.append(out)
+        start += seq_len
+    return torch.cat(outs, dim=0) if outs else qkv_feats.new_empty((0, qkv_feats.shape[2], qkv_feats.shape[3]))
+
+
 def calc_window_partition(
     tensor: SparseTensor,
     window_size: Union[int, Tuple[int, ...]],
@@ -110,6 +136,8 @@ def sparse_windowed_scaled_dot_product_self_attention(
             out = xops.memory_efficient_attention(q, k, v)          # [B, N, H, C]
         elif ATTN == 'flash_attn':
             out = flash_attn.flash_attn_qkvpacked_func(qkv_feats)   # [B, N, H, C]
+        elif ATTN in {'sdpa', 'naive'}:
+            out = _sdpa_varlen_qkv(qkv_feats.reshape(B * N, 3, H, C), [N] * B)
         else:
             raise ValueError(f"Unknown attention module: {ATTN}")
         out = out.reshape(B * N, H, C)                              # [M, H, C]
@@ -125,6 +153,10 @@ def sparse_windowed_scaled_dot_product_self_attention(
             cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
                         .to(qkv.device).int()
             out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
+        elif ATTN in {'sdpa', 'naive'}:
+            out = _sdpa_varlen_qkv(qkv_feats, seq_lens)
+        else:
+            raise ValueError(f"Unknown attention module: {ATTN}")
 
     out = out[bwd_indices]      # [T, H, C]