no flex_attention

unison/models/transformers/modules/attention.py

@@ -17,7 +17,6 @@
 import einops
 import torch
 from typing import Optional
-from loguru import logger
 import numpy as np
 import torch.nn.functional as F
 
@@ -32,15 +31,7 @@ from unison.utils.flash_attn_no_pad import (
 )
 from unison.commons import maybe_fallback_attn_mode
 
-try:
-    from torch.nn.attention.flex_attention import flex_attention
-
-    # Left uncompiled: torch.compile(flex_attention, ...) hits a Dynamo bug on this
-    # torch build (AssertionError: Cannot construct ConstantVariable for value of
-    # type torch.device, while tracing _validate_sdpa_input). 
-    flex_mask_cache = {}
-except Exception:
-    logger.warning("Could not load Sliding Tile Attention of FlexAttn.")
+flex_mask_cache = {}
 
 from unison.models.transformers.modules.ssta_attention import ssta_3d_attention
 from unison.commons.infer_state import get_infer_state
@@ -180,17 +171,25 @@ def sequence_parallel_attention(q, k, v,
             elif attn_mask.dtype != torch.bool:
                 attn_mask = attn_mask.to(query.dtype)
                 raise NotImplementedError(f'Float attention mask is not implemented for torch attention.')
-            
+
+            # [B, L] -> [B, 1, L, L]: query token q and key token k must both be
+            # unmasked.
+            attn_mask1 = einops.rearrange(attn_mask, 'b l -> b 1 l 1')
+            attn_mask2 = einops.rearrange(attn_mask1, 'b 1 l 1 -> b 1 1 l')
+            attn_mask = attn_mask1 & attn_mask2
+
         # transpose q,k,v dim to fit scaled_dot_product_attention
         query = query.transpose(1, 2)  # B * Head_num * length * dim
         key = key.transpose(1, 2)      # B * Head_num * length * dim
         value = value.transpose(1, 2)  # B * Head_num * length * dim
 
-        def score_mod(score, b, h, q_idx, kv_idx):
-            return torch.where(attn_mask[b, q_idx] & attn_mask[b, kv_idx], score, float('-inf'))
+        # cut: flex_attention score_mod path, replaced by plain SDPA below —
+        # flex_attention always traces internally, which crashed on this torch build.
+        # def score_mod(score, b, h, q_idx, kv_idx):
+        #     return torch.where(attn_mask[b, q_idx] & attn_mask[b, kv_idx], score, float('-inf'))
+        # hidden_states = flex_attention(query, key, value, score_mod=score_mod)
+        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask)
 
-        hidden_states = flex_attention(query, key, value, score_mod=score_mod)
-        
         # transpose back
         hidden_states = hidden_states.transpose(1, 2)