diff --git a/fastvideo/attention/backends/flash_attn.py b/fastvideo/attention/backends/flash_attn.py index d53f0d9b0..8aad6a5b0 100644 --- a/fastvideo/attention/backends/flash_attn.py +++ b/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) diff --git a/fastvideo/attention/backends/sdpa.py b/fastvideo/attention/backends/sdpa.py index 5b3c5df50..c461c0491 100644 --- a/fastvideo/attention/backends/sdpa.py +++ b/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): diff --git a/fastvideo/models/encoders/mmaudio_clip.py b/fastvideo/models/encoders/mmaudio_clip.py index 476e2c510..5a0c66f99 100644 --- a/fastvideo/models/encoders/mmaudio_clip.py +++ b/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( diff --git a/fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py b/fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py index a267693d4..56cb8f4a2 100644 --- a/fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py +++ b/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."""