From e92137f96821f1fd5ac8ebc75261708f63752c62 Mon Sep 17 00:00:00 2001 From: Iryna Boiko Date: Mon, 19 Oct 2026 21:00:58 +0300 Subject: [PATCH] [Bugfix][XPU] Honor full_attn_spans in the XPU FlashAttention path Signed-off-by: Iryna Boiko --- tests/diffusion/attention/test_flash_attn.py | 25 +++++++++++++++++++ .../attention/backends/flash_attn.py | 13 ++++++++++ 2 files changed, 38 insertions(+) diff --git a/tests/diffusion/attention/test_flash_attn.py b/tests/diffusion/attention/test_flash_attn.py index 2542a956a26..2b8293f00df 100644 --- a/tests/diffusion/attention/test_flash_attn.py +++ b/tests/diffusion/attention/test_flash_attn.py @@ -629,6 +629,31 @@ def fake_varlen_func(q, k, v, **kwargs): assert calls[0]["softmax_scale"] == 0.5 +@pytest.mark.cpu +def test_xpu_piecewise_spans_ignore_4d_mask(monkeypatch): + """XPU must honor full_attn_spans instead of feeding a 4D mask to the 2D-only varlen-masked path.""" + calls = [] + + def fake_varlen_func(q, k, v, **kwargs): + calls.append((kwargs["max_seqlen_q"], kwargs["max_seqlen_k"], kwargs["causal"])) + return q + + monkeypatch.setattr(fa, "HAS_FLASH_ATTN", True) + monkeypatch.setattr(fa, "flash_attn_varlen_func", fake_varlen_func) + + impl = FlashAttentionImpl(num_heads=2, head_size=4, softmax_scale=0.5, causal=False) + query = torch.randn(1, 6, 2, 4) + metadata = AttentionMetadata( + attn_mask=torch.ones((1, 1, 6, 6), dtype=torch.bool).tril(), + full_attn_spans=[[(2, 4)]], + ) + + output = impl.forward_xpu(query, query, query, metadata) + + assert output.shape == query.shape + assert calls == [(2, 2, True), (2, 4, False), (2, 6, True)] + + @pytest.mark.cpu def test_packed_varlen_metadata_bypasses_mask_unpadding(monkeypatch): calls = [] diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index 814d6298253..29f87f1e8ce 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -979,6 +979,7 @@ def forward_xpu( """XPU flash attention implementation.""" from vllm_omni.diffusion.attention.backends.utils.fa import ( HAS_FLASH_ATTN, + flash_attn_varlen_func, ) if not HAS_FLASH_ATTN: @@ -988,6 +989,18 @@ def forward_xpu( "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA" ) + # Spans take precedence over attn_mask, matching forward_cuda. + if attn_metadata is not None and attn_metadata.full_attn_spans is not None: + return piecewise_attn( + query, + key, + value, + attn_metadata.full_attn_spans, + self.softmax_scale, + partial(FlashAttentionImpl._flash_varlen_wrapper, attn_func=flash_attn_varlen_func), + query_ranges=attn_metadata.query_ranges, + ) + attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None if attention_mask is not None and torch.any(~attention_mask):