[bugfix] Preserve masked FlashAttention mask and scale

fastvideo/attention/backends/flash_attn.py

@@ -233,8 +233,13 @@ def _key_padding_mask_from_attn_mask(attn_mask: torch.Tensor, key_len: int) -> t
     elif attn_mask.dim() != 2:
         raise ValueError(f"Unsupported attn_mask shape for FLASH_ATTN: {attn_mask.shape}")
 
-    # SDPA additive mask convention: valid=0, masked=-inf/large negative.
-    key_padding_mask = attn_mask if attn_mask.dtype == torch.bool else attn_mask >= 0
+    # Integer masks follow tokenizer convention; floating masks are additive.
+    if attn_mask.dtype == torch.bool:
+        key_padding_mask = attn_mask
+    elif attn_mask.dtype.is_floating_point:
+        key_padding_mask = attn_mask >= 0
+    else:
+        key_padding_mask = attn_mask != 0
 
     if key_padding_mask.shape[-1] != key_len:
         raise ValueError("Invalid key padding mask length for FLASH_ATTN: "
@@ -506,7 +511,11 @@ def _forward_impl(
                 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=self.causal, dropout_p=0, softmax_scale=None)
+            output = flash_attn_no_pad(qkv,
+                                       attn_mask_padded,
+                                       causal=self.causal,
+                                       dropout_p=0,
+                                       softmax_scale=self.softmax_scale)
         elif self.nvfp4_fa4:
             output = self._forward_nvfp4(query, key, value)
 

fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py

@@ -284,3 +284,46 @@ def test_varlen_qk_opcheck_with_grad_inputs(no_pad_impls):
         torch.ops.fastvideo._flash_attn_varlen_qk_no_pad_forward,
         (q, k, v, qmask, kmask, False, 0.0, None, False),
     )
+
+
+@pytest.mark.parametrize("query_len", [48, 64])
+@pytest.mark.parametrize("mask_kind", ["bool", "integer", "additive"])
+@pytest.mark.parametrize("mask_rank", [2, 3, 4])
+def test_backend_padding_mask_and_custom_scale(no_pad_impls, query_len, mask_kind, mask_rank):
+    """Masked FlashAttention agrees with SDPA for tokenizer and additive masks."""
+    from fastvideo.attention.backends.flash_attn import FlashAttentionImpl
+    from fastvideo.attention.backends.sdpa import SDPAMetadata
+
+    torch.manual_seed(42)
+    batch, key_len, heads, head_dim = 2, 64, 2, 64
+    if query_len != key_len:
+        _fa2_only(no_pad_impls)
+    query = torch.randn(batch, query_len, heads, head_dim, device="cuda", dtype=torch.float16)
+    key = torch.randn(batch, key_len, heads, head_dim, device="cuda", dtype=torch.float16)
+    value = torch.randn_like(key)
+    valid = _padding_mask(batch, key_len, [50, 37], query.device)
+    if mask_kind == "integer":
+        mask = valid.to(torch.int64)
+    elif mask_kind == "additive":
+        mask = torch.where(valid, 0.0, float("-inf"))
+    else:
+        mask = valid
+    if mask_rank == 3:
+        mask = mask[:, None, :]
+    elif mask_rank == 4:
+        mask = mask[:, None, None, :]
+
+    scale = 0.37  # Deliberately different from head_dim**-0.5.
+    impl = FlashAttentionImpl(num_heads=heads, head_size=head_dim, causal=False,
+                              softmax_scale=scale, nvfp4_fa4=False)
+    metadata = SDPAMetadata(current_timestep=0, attn_mask=mask)
+    with torch.inference_mode():
+        actual = impl.forward(query, key, value, metadata)
+        expected = torch.nn.functional.scaled_dot_product_attention(
+            query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2),
+            attn_mask=valid[:, None, None, :], scale=scale,
+        ).transpose(1, 2)
+    if query_len == key_len:
+        # The self-attention path also removes padded query rows.
+        actual, expected = actual[valid], expected[valid]
+    torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2)