attn bf16

trellis/modules/attention/full_attn.py

@@ -122,12 +122,16 @@ def scaled_dot_product_attention(*args, **kwargs):
             k, v = kv.unbind(dim=2)
         out = xops.memory_efficient_attention(q, k, v)
     elif BACKEND == 'flash_attn':
+        _fa_dtype = torch.bfloat16
         if num_all_args == 1:
-            out = flash_attn.flash_attn_qkvpacked_func(qkv)
+            _orig_dtype = qkv.dtype
+            out = flash_attn.flash_attn_qkvpacked_func(qkv.to(_fa_dtype)).to(_orig_dtype)
         elif num_all_args == 2:
-            out = flash_attn.flash_attn_kvpacked_func(q, kv)
+            _orig_dtype = q.dtype
+            out = flash_attn.flash_attn_kvpacked_func(q.to(_fa_dtype), kv.to(_fa_dtype)).to(_orig_dtype)
         elif num_all_args == 3:
-            out = flash_attn.flash_attn_func(q, k, v)
+            _orig_dtype = q.dtype
+            out = flash_attn.flash_attn_func(q.to(_fa_dtype), k.to(_fa_dtype), v.to(_fa_dtype)).to(_orig_dtype)
     elif BACKEND == 'sdpa':
         if num_all_args == 1:
             q, k, v = qkv.unbind(dim=2)

trellis/modules/sparse/attention/full_attn.py

@@ -219,12 +219,16 @@ def sparse_scaled_dot_product_attention(*args, **kwargs):
         cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
         if num_all_args in [2, 3]:
             cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+        _fa_dtype = torch.bfloat16
         if num_all_args == 1:
-            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
+            _orig_dtype = qkv.dtype
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv.to(_fa_dtype), cu_seqlens_q, max(q_seqlen)).to(_orig_dtype)
         elif num_all_args == 2:
-            out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+            _orig_dtype = q.dtype
+            out = flash_attn.flash_attn_varlen_kvpacked_func(q.to(_fa_dtype), kv.to(_fa_dtype), cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen)).to(_orig_dtype)
         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))
+            _orig_dtype = q.dtype
+            out = flash_attn.flash_attn_varlen_func(q.to(_fa_dtype), k.to(_fa_dtype), v.to(_fa_dtype), cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen)).to(_orig_dtype)
     elif ATTN in {'sdpa', 'naive'}:
         outs = []
         q_start, kv_start = 0, 0

trellis/modules/sparse/attention/serialized_attn.py

@@ -193,7 +193,8 @@ def sparse_serialized_scaled_dot_product_self_attention(
             q, k, v = qkv_feats.unbind(dim=2)                       # [B, N, H, C]
             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]
+            _orig_dtype = qkv_feats.dtype
+            out = flash_attn.flash_attn_qkvpacked_func(qkv_feats.to(torch.bfloat16)).to(_orig_dtype)   # [B, N, H, C]
         elif ATTN in {'sdpa', 'naive'}:
             out = _sdpa_varlen_qkv(qkv_feats.reshape(B * N, 3, H, C), [N] * B)
         else:
@@ -210,7 +211,8 @@ def sparse_serialized_scaled_dot_product_self_attention(
         elif ATTN == 'flash_attn':
             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]
+            _orig_dtype = qkv_feats.dtype
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats.to(torch.bfloat16), cu_seqlens, max(seq_lens)).to(_orig_dtype) # [M, H, C]
         elif ATTN in {'sdpa', 'naive'}:
             out = _sdpa_varlen_qkv(qkv_feats, seq_lens)
         else:

trellis/modules/sparse/attention/windowed_attn.py

@@ -135,7 +135,8 @@ def sparse_windowed_scaled_dot_product_self_attention(
             q, k, v = qkv_feats.unbind(dim=2)                       # [B, N, H, C]
             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]
+            _orig_dtype = qkv_feats.dtype
+            out = flash_attn.flash_attn_qkvpacked_func(qkv_feats.to(torch.bfloat16)).to(_orig_dtype)   # [B, N, H, C]
         elif ATTN in {'sdpa', 'naive'}:
             out = _sdpa_varlen_qkv(qkv_feats.reshape(B * N, 3, H, C), [N] * B)
         else:
@@ -152,7 +153,8 @@ def sparse_windowed_scaled_dot_product_self_attention(
         elif ATTN == 'flash_attn':
             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]
+            _orig_dtype = qkv_feats.dtype
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats.to(torch.bfloat16), cu_seqlens, max(seq_lens)).to(_orig_dtype) # [M, H, C]
         elif ATTN in {'sdpa', 'naive'}:
             out = _sdpa_varlen_qkv(qkv_feats, seq_lens)
         else: