add _compile = True to create_block_mask

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "transfusion-pytorch"
-version = "0.9.4"
+version = "0.9.5"
 description = "Transfusion in Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

transfusion_pytorch/transfusion.py

@@ -984,7 +984,7 @@ def forward(
         if needs_masking:
             if causal_mask:
                 if should_use_flex_attn:
-                    block_mask = create_block_mask(causal, B = None, H = None, Q_LEN = seq_len, KV_LEN = seq_len, device = device)
+                    block_mask = create_block_mask(causal, B = None, H = None, Q_LEN = seq_len, KV_LEN = seq_len, _compile = True, device = device)
                     attn_mask_kwargs.update(block_mask = block_mask)
                 else:
                     attn_mask_kwargs.update(causal = True)
@@ -994,7 +994,7 @@ def forward(
 
                 if should_use_flex_attn:
                     transfusion_mask_fn = transfusion_attn_mask(modality_positions)
-                    block_mask = create_block_mask(transfusion_mask_fn, B = None, H = None, Q_LEN = seq_len, KV_LEN = seq_len, device = device)
+                    block_mask = create_block_mask(transfusion_mask_fn, B = None, H = None, Q_LEN = seq_len, KV_LEN = seq_len, _compile = True, device = device)
                     attn_mask_kwargs.update(block_mask = block_mask)
                 else:
                     attn_mask = naive_attn_mask(seq_len, modality_positions, device = device)