Upload BD3LM

modeling_bd3lm.py

@@ -396,10 +396,10 @@ class DDiTBlock(nn.Module):
 
     if mask is None and self.attn_backend == 'flash_attn':
       x = regular_attention_multi_headed(qkv)
-    elif self.attn_backend == 'sdpa':
-      x = self.cross_attn(x, qkv, mask=mask)
     elif self.attn_backend == 'flex' and FLEX_ATTN_AVAILABLE:
       x = self.cross_attn_flex(qkv, mask=mask)
+    elif self.attn_backend == 'sdpa':
+      x = self.cross_attn(x, qkv, mask=mask)
     else:
       raise ValueError('Unknown attention backend')
 
@@ -485,7 +485,7 @@ class DITBackbone(nn.Module):
       config.vocab_size,
       config.cond_dim)
     if self.cross_attn:
-      self.gen_mask(config.model_length, self.block_size)
+      self.gen_mask(config.model_length, self.block_size, attn_backend=config.attn_backend)
     self.precision = torch.float32
 
   def _get_bias_dropout_scale(self):
@@ -496,14 +496,14 @@ class DITBackbone(nn.Module):
     
   def gen_mask(self, seqlen, block_size, attn_backend='sdpa'):
     """Genererates attention mask"""
-    if attn_backend == 'sdpa':
-      self.mask = block_diff_mask(
-        b=None, h=None, q_idx=torch.arange(seqlen*2)[:, None], kv_idx=torch.arange(seqlen*2)[None, :],
-        block_size=block_size, n=seqlen)
-    elif attn_backend == 'flex':
+    if attn_backend == 'flex' and FLEX_ATTN_AVAILABLE:
       self.mask = create_block_mask(
         partial(block_diff_mask, block_size=block_size, n=seqlen),
         B=None, H=None, Q_LEN=seqlen*2, KV_LEN=seqlen*2)
+    elif attn_backend == 'sdpa':
+      self.mask = block_diff_mask(
+        b=None, h=None, q_idx=torch.arange(seqlen*2)[:, None], 
+        kv_idx=torch.arange(seqlen*2)[None, :], block_size=block_size, n=seqlen)
     else:
       raise ValueError('Unknown attention backend')