[fix]: preserve causal masks in FlashAttention

fastvideo/attention/backends/flash_attn.py

@@ -168,8 +168,12 @@ class FlashAttentionBackend(AttentionBackend):
 def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> torch.Tensor:
     # Normalize attn_mask to [B, key_len] where True means valid token.
     if attn_mask.dim() == 4:
+        if attn_mask.shape[1] != 1 or attn_mask.shape[-2] != 1:
+            raise ValueError("FLASH_ATTN only supports 4D key-padding masks with shape [B, 1, 1, K]")
         attn_mask = attn_mask[:, 0, 0, :]
     elif attn_mask.dim() == 3:
+        if attn_mask.shape[-2] != 1:
+            raise ValueError("FLASH_ATTN only supports 3D key-padding masks with shape [B, 1, K]")
         attn_mask = attn_mask[:, 0, :]
     elif attn_mask.dim() != 2:
         raise ValueError(f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}")
@@ -274,6 +278,14 @@ class FlashAttentionImpl(AttentionImpl):
             )
 
             attn_mask = attn_metadata.attn_mask
+            if getattr(attn_metadata, "is_causal", False):
+                return flash_attn_func_compilable(
+                    query,
+                    key,
+                    value,
+                    softmax_scale=self.softmax_scale,
+                    causal=True,
+                )
 
             # flash_attn_no_pad packs q/k/v as one tensor and assumes equal q/k
             # sequence lengths. Cross-attention can violate this.
@@ -302,7 +314,7 @@ class FlashAttentionImpl(AttentionImpl):
                 raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
                                  f"expected at most {qkv.shape[1]}, got {key_padding_mask.shape[-1]}")
             attn_mask_padded = F.pad(key_padding_mask, (qkv.shape[1] - key_padding_mask.shape[-1], 0), value=True)
-            output = flash_attn_no_pad(qkv, attn_mask_padded, causal=False, dropout_p=0, softmax_scale=None)
+            output = flash_attn_no_pad(qkv, attn_mask_padded, causal=self.causal, dropout_p=0, softmax_scale=None)
         elif self.nvfp4_fa4:
             output = self._forward_nvfp4(query, key, value)
 

fastvideo/attention/backends/sdpa.py

@@ -41,6 +41,8 @@ class SDPABackend(AttentionBackend):
 class SDPAMetadata(AttentionMetadata):
     current_timestep: int
     attn_mask: torch.Tensor | None = None
+    # The mask is exactly native causal attention, with no additional padding.
+    is_causal: bool = False
 
 
 class SDPAMetadataBuilder(AttentionMetadataBuilder):

fastvideo/models/encoders/mmaudio_clip.py

@@ -60,6 +60,7 @@ class MMAudioDFNCLIPTextEncoder(CLIPTextModel):
         metadata = SDPAMetadata(
             current_timestep=0,
             attn_mask=causal_mask[None, None],
+            is_causal=True,
         )
         with set_forward_context(current_timestep=0, attn_metadata=metadata):
             output = super().forward(

fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py

@@ -85,6 +85,50 @@ def test_no_pad_inference_matches_original(no_pad_impls, dtype):
     torch.testing.assert_close(out_test, out_ref, atol=0, rtol=0)
 
 
+def test_backend_translates_explicit_causal_mask(no_pad_impls):
+    """MMAudio's additive mask must route through native causal attention."""
+    from fastvideo.attention.backends.flash_attn import (
+        FlashAttentionImpl,
+        flash_attn_func_compilable,
+    )
+    from fastvideo.attention.backends.sdpa import SDPAMetadata
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    batch, seqlen, heads, head_dim = 1, 32, 2, 64
+    query = torch.randn(batch, seqlen, heads, head_dim, device=device, dtype=torch.bfloat16)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    causal_mask = torch.full(
+        (1, 1, seqlen, seqlen),
+        float("-inf"),
+        device=device,
+        dtype=query.dtype,
+    ).triu_(1)
+    impl = FlashAttentionImpl(
+        num_heads=heads,
+        head_size=head_dim,
+        causal=False,
+        softmax_scale=head_dim**-0.5,
+    )
+
+    with torch.inference_mode():
+        expected = flash_attn_func_compilable(
+            query,
+            key,
+            value,
+            softmax_scale=head_dim**-0.5,
+            causal=True,
+        )
+        actual = impl.forward(
+            query,
+            key,
+            value,
+            SDPAMetadata(current_timestep=0, attn_mask=causal_mask, is_causal=True),
+        )
+    torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+
+
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 def test_no_pad_training_backward_through_registered_autograd(no_pad_impls, dtype):
     """FA2: grads flow through the registered op and match the original."""