From 2d4724eb74b26e5a0420d87be2da9c9ba73804db Mon Sep 17 00:00:00 2001 From: vOv <124694036+cr-gao@users.noreply.github.com> Date: Thu, 27 Aug 2026 01:20:59 +0800 Subject: [PATCH] [Bugfix][Diffusion] Fix FLASH_ATTN cross-attention key-padding unpad (#5866) Signed-off-by: cr-gao --- tests/diffusion/attention/test_flash_attn.py | 85 +++++++++++++++++++ .../attention/backends/flash_attn.py | 26 ++++-- 2 files changed, 106 insertions(+), 5 deletions(-) diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 5924beabf..2fba447b0 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -308,6 +308,91 @@ def test_flash_attn_func_preferred_over_varlen(): print("✓ flash_attn_func forward works correctly!") +@pytest.mark.skipif(not is_gpu, reason="FlashAttention requires CUDA or XPU") +@pytest.mark.parametrize("k_len", [40, 256]) +def test_cross_attn_key_padding_vs_sdpa(k_len): + """ + Case 3: Cross-attention with a key-padding mask. + + The mask covers only the key sequence; every query position must attend + to the valid keys. FA and SDPA outputs should be very close. k_len=256 + covers equal Q/K lengths, which must not be mistaken for self-attention, + and valid key lengths differ per batch element. + """ + device = torch.device(current_omni_platform.device_type) + dtype = torch.bfloat16 + + batch_size = 2 + q_len = 256 + valid_lens = (17, 29) + num_heads = 8 + head_dim = 64 + + fa_impl = FlashAttentionImpl( + num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False, role="cross" + ) + sdpa_impl = SDPAImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=1.0 / (head_dim**0.5), causal=False) + + torch.manual_seed(7) + query = torch.randn(batch_size, q_len, num_heads, head_dim, device=device, dtype=dtype) + key = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype) + value = torch.randn(batch_size, k_len, num_heads, head_dim, device=device, dtype=dtype) + + attn_mask = torch.zeros(batch_size, k_len, dtype=torch.bool, device=device) + for i, valid_len in enumerate(valid_lens): + attn_mask[i, :valid_len] = True + attn_metadata = AttentionMetadata(attn_mask=attn_mask) + + output_fa = fa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata) + output_sdpa = sdpa_impl.forward(query=query, key=key, value=value, attn_metadata=attn_metadata) + + max_diff = torch.max(torch.abs(output_fa - output_sdpa)).item() + mean_diff = torch.mean(torch.abs(output_fa - output_sdpa)).item() + + print("\n=== Case 3: Cross-Attention Key-Padding FA vs SDPA ===") + print(f"Max absolute difference: {max_diff:.6f}") + print(f"Mean absolute difference: {mean_diff:.6f}") + + assert max_diff < 0.01, f"Max difference {max_diff} exceeds threshold 0.01" + assert mean_diff < 0.001, f"Mean difference {mean_diff} exceeds threshold 0.001" + + print("✓ Case 3 PASSED: FA and SDPA cross-attention outputs are very close!") + + +def test_varlen_masked_routing_by_role(monkeypatch): + """The unpad route is picked by role, not Q/K length equality; runs without a GPU.""" + calls = [] + + def fake_varlen_func(q, k, v, **kwargs): + calls.append((q.shape[0], kwargs["cu_seqlens_q"], kwargs["cu_seqlens_k"])) + return q + + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func) + + batch_size, seq_len, num_heads, head_dim = 2, 4, 2, 8 + query = torch.randn(batch_size, seq_len, num_heads, head_dim) + metadata = AttentionMetadata(attn_mask=torch.tensor([[True, True, False, False], [True, True, True, False]])) + + cross_impl = FlashAttentionImpl( + num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False, role="cross" + ) + output = cross_impl.forward_cuda(query, query, query, metadata) + assert output.shape == query.shape + num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1] + assert num_q_rows == batch_size * seq_len + assert cu_seqlens_q.tolist() == [0, 4, 8] + assert cu_seqlens_k.tolist() == [0, 2, 5] + + self_impl = FlashAttentionImpl(num_heads=num_heads, head_size=head_dim, softmax_scale=0.5, causal=False) + output = self_impl.forward_cuda(query, query, query, metadata) + assert output.shape == query.shape + num_q_rows, cu_seqlens_q, cu_seqlens_k = calls[-1] + assert num_q_rows == 5 + assert cu_seqlens_q.tolist() == [0, 2, 5] + assert cu_seqlens_k.tolist() == [0, 2, 5] + + def test_piecewise_flash_attn_uses_varlen_fallback(monkeypatch): calls = [] diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 838fd3d6d..dd344e85d 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -93,12 +93,14 @@ class FlashAttentionImpl(AttentionImpl): prefix: str = "", qkv_layout: str | None = None, backend_kwargs: dict | None = None, + role: str = "self", **extra_impl_args, ) -> None: self.num_heads = num_heads self.causal = causal self.softmax_scale = softmax_scale self.qkv_layout = qkv_layout + self.is_cross_attn = role == "cross" cfg = get_current_diffusion_config_or_none() self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False if backend_kwargs: @@ -152,6 +154,7 @@ class FlashAttentionImpl(AttentionImpl): attention_mask: torch.Tensor, ) -> torch.Tensor: from vllm_omni.diffusion.attention.backends.utils.fa import ( + _index_first_axis, _pad_input, _unpad_input, _upad_input, @@ -159,10 +162,21 @@ class FlashAttentionImpl(AttentionImpl): ) assert attention_mask.ndim == 2, "attention_mask must be 2D, (batch_size, seq_len)" - query_length = query.size(1) - q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input( - query, key, value, attention_mask, query_length, _unpad_input - ) + batch_size, query_length = query.shape[:2] + if not self.is_cross_attn and query_length == key.size(1): + q, k, v, indices_q, (cu_seq_lens_q, cu_seq_lens_k), (max_length_q, max_length_k) = _upad_input( + query, key, value, attention_mask, query_length, _unpad_input + ) + else: + # Cross-attention: the mask covers keys only, so keep every query row. + k, indices_k, cu_seq_lens_k, max_length_k, _ = _unpad_input(key, attention_mask) + v = _index_first_axis(value, indices_k) + q = query.flatten(0, 1) + cu_seq_lens_q = torch.arange( + 0, (batch_size + 1) * query_length, query_length, dtype=torch.int32, device=query.device + ) + max_length_q = query_length + indices_q = None out_unpad = flash_attn_varlen_func( q, @@ -178,7 +192,9 @@ class FlashAttentionImpl(AttentionImpl): }, ) out_unpad = self._unwrap_flash_output(out_unpad) - return _pad_input(out_unpad, indices_q, query.size(0), query_length) + if indices_q is None: + return out_unpad.reshape(batch_size, query_length, *out_unpad.shape[1:]) + return _pad_input(out_unpad, indices_q, batch_size, query_length) def _forward_varlen_packed( self,