[Bugfix][XPU] Honor full_attn_spans in the XPU FlashAttention path

tests/diffusion/attention/test_flash_attn.py

@@ -728,6 +728,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 = []

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -983,6 +983,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:
@@ -992,6 +993,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):