fix _flash_3_varlen_hub mask handling (#14115)

* fix fa3 mask

* add test case

---------

Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>

Author: zhtmike

src/diffusers/models/attention_dispatch.py

@@ -3024,22 +3024,20 @@ def _flash_attention_3_varlen_hub(
 
     if attn_mask is not None:
         attn_mask = _normalize_attn_mask(attn_mask, batch_size, seq_len_kv)
-
-    (_, seqlens_k), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
-        _prepare_for_flash_attn_or_sage_varlen(
-            batch_size, seq_len_q, seq_len_kv, attn_mask=attn_mask, device=query.device
+        (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
+            _prepare_for_flash_attn_or_sage_varlen_with_mask(batch_size, seq_len_q, attn_mask, query.device)
         )
-    )
-
-    key_valid, value_valid = [], []
-    for b in range(batch_size):
-        valid_len = seqlens_k[b]
-        key_valid.append(key[b, :valid_len])
-        value_valid.append(value[b, :valid_len])
+        indices_k = attn_mask.flatten().nonzero(as_tuple=False).flatten()
+        key_packed = key.reshape(-1, *key.shape[2:])[indices_k]
+        value_packed = value.reshape(-1, *value.shape[2:])[indices_k]
+    else:
+        (_, _), (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = (
+            _prepare_for_flash_attn_or_sage_varlen_without_mask(batch_size, seq_len_q, seq_len_kv, query.device)
+        )
+        key_packed = key.flatten(0, 1)
+        value_packed = value.flatten(0, 1)
 
     query_packed = query.flatten(0, 1)
-    key_packed = torch.cat(key_valid, dim=0)
-    value_packed = torch.cat(value_valid, dim=0)
 
     func = _HUB_KERNELS_REGISTRY[AttentionBackendName._FLASH_3_VARLEN_HUB].kernel_fn
     result = func(

tests/models/testing_utils/attention.py

@@ -70,8 +70,38 @@
     ],
 )
 
+_PARAM_FLASH_VARLEN_HUB = pytest.param(
+    AttentionBackendName.FLASH_VARLEN_HUB,
+    id="flash_varlen_hub",
+    marks=[
+        pytest.mark.skipif(not _CUDA_AVAILABLE, reason="CUDA is required for flash_varlen_hub backend."),
+        pytest.mark.skipif(
+            not is_kernels_available(),
+            reason="`kernels` package is required for flash_varlen_hub backend. Install with `pip install kernels`.",
+        ),
+    ],
+)
+
+_PARAM_FLASH_3_VARLEN_HUB = pytest.param(
+    AttentionBackendName._FLASH_3_VARLEN_HUB,
+    id="flash_3_varlen_hub",
+    marks=[
+        pytest.mark.skipif(not _CUDA_AVAILABLE, reason="CUDA is required for _flash_3_varlen_hub backend."),
+        pytest.mark.skipif(
+            not is_kernels_available(),
+            reason="`kernels` package is required for _flash_3_varlen_hub backend. Install with `pip install kernels`.",
+        ),
+    ],
+)
+
 # All backends under test.
-_ALL_BACKEND_PARAMS = [_PARAM_NATIVE_CUDNN, _PARAM_FLASH_HUB, _PARAM_FLASH_3_HUB]
+_ALL_BACKEND_PARAMS = [
+    _PARAM_NATIVE_CUDNN,
+    _PARAM_FLASH_HUB,
+    _PARAM_FLASH_3_HUB,
+    _PARAM_FLASH_VARLEN_HUB,
+    _PARAM_FLASH_3_VARLEN_HUB,
+]
 
 # Backends that perform non-deterministic operations and therefore cannot run when
 # torch.use_deterministic_algorithms(True) is active (e.g. after enable_full_determinism()).
@@ -267,6 +297,8 @@ class AttentionBackendTesterMixin:
         Use `pytest -m "not attention"` to skip these tests.
     """
 
+    unsupported_attn_backends: list[str] = []
+
     def setup_method(self):
         gc.collect()
         backend_empty_cache(torch_device)
@@ -280,6 +312,8 @@ def teardown_method(self):
     def test_set_attention_backend_matches_context_manager(self, backend):
         """set_attention_backend() and the attention_backend() context manager must yield identical outputs."""
         _skip_if_backend_requires_nondeterminism(backend)
+        if backend.value in self.unsupported_attn_backends:
+            pytest.skip(f"{backend.value} is not supported for this model.")
 
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
@@ -318,6 +352,8 @@ def test_set_attention_backend_matches_context_manager(self, backend):
     def test_output_close_to_native(self, backend, atol=1e-2, rtol=1e-2):
         """All backends should produce model output numerically close to the native SDPA reference."""
         _skip_if_backend_requires_nondeterminism(backend)
+        if backend.value in self.unsupported_attn_backends:
+            pytest.skip(f"{backend.value} is not supported for this model.")
 
         init_dict = self.get_init_dict()
         inputs_dict = self.get_dummy_inputs()
@@ -378,6 +414,8 @@ def test_compile(self, backend, atol=1e-2, rtol=1e-2):
         as opposed to `model.compile`).
         """
         _skip_if_backend_requires_nondeterminism(backend)
+        if backend.value in self.unsupported_attn_backends:
+            pytest.skip(f"{backend.value} is not supported for this model.")
         if getattr(self.model_class, "_repeated_blocks", None) is None:
             pytest.skip("Skipping tests as regional compilation is not supported.")
 

tests/models/testing_utils/utils.py

@@ -8,6 +8,7 @@
     AttentionBackendName.FLASH_HUB,
     AttentionBackendName.FLASH_VARLEN_HUB,
     AttentionBackendName._FLASH_3_HUB,
+    AttentionBackendName._FLASH_3_VARLEN_HUB,
 }
 
 

tests/models/transformers/test_models_transformer_qwenimage.py

@@ -22,6 +22,7 @@
 
 from ...testing_utils import enable_full_determinism, torch_device
 from ..testing_utils import (
+    AttentionBackendTesterMixin,
     AttentionTesterMixin,
     BaseModelTesterConfig,
     BitsAndBytesTesterMixin,
@@ -250,6 +251,21 @@ class TestQwenImageTransformerAttention(QwenImageTransformerTesterConfig, Attent
     """Attention processor tests for QwenImage Transformer."""
 
 
+class TestQwenImageTransformerAttentionBackend(QwenImageTransformerTesterConfig, AttentionBackendTesterMixin):
+    """Attention backend tests for QwenImage Transformer."""
+
+    unsupported_attn_backends = ["flash_hub", "_flash_3_hub"]
+
+    def get_dummy_inputs(self, batch_size: int = 2):
+        inputs = super().get_dummy_inputs(batch_size=batch_size)
+        mask = inputs["encoder_hidden_states_mask"]
+        mask[...] = 0
+        mask[0, :2] = 1
+        mask[1, :6] = 1
+        inputs["encoder_hidden_states_mask"] = mask.bool()
+        return inputs
+
+
 class TestQwenImageTransformerContextParallel(QwenImageTransformerTesterConfig, ContextParallelTesterMixin):
     """Context Parallel inference tests for QwenImage Transformer."""