[XPU][Attention] Enable per-sequence causal masking with FA2

tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-XPU-FA2-FP8.yaml

@@ -0,0 +1,17 @@
+model_name: "google/diffusiongemma-26B-A4B-it"
+accuracy_threshold: 0.84
+num_questions: 1319
+num_fewshot: 5
+max_concurrency: 2
+startup_max_wait_seconds: 1200
+use_chat_completions: true
+# Diffusion models use a fixed internal temperature schedule and do not
+# support per-request temperature or seed overrides.
+temperature: 1.0
+seed: null
+server_args: >-
+  --enforce-eager
+  --max-model-len 4096
+  --tensor-parallel-size 2
+  --attention-backend FLASH_ATTN
+  --quantization fp8

tests/kernels/attention/test_xpu_dynamic_causal.py

@@ -0,0 +1,54 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import patch
+
+import torch
+
+
+def test_xpu_flash_attn_forwards_dynamic_causal():
+    from vllm import _xpu_ops
+
+    q = torch.randn(3, 2, 8)
+    k = torch.randn(3, 1, 8)
+    v = torch.randn_like(k)
+    cu_seqlens = torch.tensor([0, 1, 3], dtype=torch.int32)
+    dynamic_causal = torch.tensor([1, 0], dtype=torch.int32)
+    sentinel = torch.empty_like(q)
+
+    with patch.object(_xpu_ops, "flash_attn_varlen_func", return_value=sentinel) as op:
+        result = _xpu_ops.xpu_ops.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_seqlens,
+            cu_seqlens_k=cu_seqlens,
+            max_seqlen_q=2,
+            max_seqlen_k=2,
+            dynamic_causal=dynamic_causal,
+        )
+
+    assert result is sentinel
+    assert op.call_args.kwargs["dynamic_causal"] is dynamic_causal
+
+
+def test_xpu_fa2_supports_dynamic_causal(monkeypatch):
+    from vllm.v1.attention.backends import flash_attn
+
+    monkeypatch.setattr(flash_attn.current_platform, "is_xpu", lambda: True)
+    assert flash_attn._supports_dynamic_causal(2)
+    assert flash_attn._supports_dynamic_causal(4)
+    assert not flash_attn._supports_dynamic_causal(3)
+
+    monkeypatch.setattr(flash_attn.current_platform, "is_xpu", lambda: False)
+    assert not flash_attn._supports_dynamic_causal(2)
+    assert flash_attn._supports_dynamic_causal(4)
+
+
+def test_dynamic_causal_symmetrizes_causal_window():
+    from vllm.v1.attention.backends.flash_attn import _maybe_symmetrize_window
+
+    dynamic_causal = torch.tensor([1, 0], dtype=torch.int32)
+    assert _maybe_symmetrize_window((32, 0), dynamic_causal) == (32, 32)
+    assert _maybe_symmetrize_window((-1, -1), dynamic_causal) == (-1, -1)
+    assert _maybe_symmetrize_window((32, 0), True) == (32, 0)

tests/models/test_diffusion_gemma_xpu.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+
+def test_xpu_diffusion_compile_is_identity(monkeypatch):
+    from vllm.model_executor.models import diffusion_gemma
+
+    monkeypatch.setattr(diffusion_gemma.current_platform, "is_xpu", lambda: True)
+    fn = lambda value: value  # noqa: E731
+
+    assert diffusion_gemma._compile_diffusion_fn(fn) is fn
+
+
+def test_self_conditioning_buffer_uses_model_dtype():
+    from vllm.model_executor.models.diffusion_gemma import (
+        DiffusionGemmaRequestStates,
+    )
+
+    states = DiffusionGemmaRequestStates(
+        max_num_reqs=2,
+        canvas_length=4,
+        vocab_size=8,
+        max_denoising_steps=3,
+        device=torch.device("cpu"),
+        hidden_size=16,
+        hidden_dtype=torch.bfloat16,
+        stability_threshold=2,
+    )
+
+    assert states.self_conditioning_embeds.dtype == torch.bfloat16

vllm/_xpu_ops.py

@@ -1034,6 +1034,7 @@ def flash_attn_varlen_func(
             q_descale=q_descale,
             k_descale=k_descale,
             v_descale=v_descale,
+            dynamic_causal=dynamic_causal,
         )
 
     @staticmethod

vllm/model_executor/models/diffusion_gemma.py

@@ -73,6 +73,14 @@
 logger = init_logger(__name__)
 
 
+def _compile_diffusion_fn(func):
+    if current_platform.is_xpu():
+        return func
+    return torch.compile(
+        func, dynamic=True, backend=current_platform.simple_compile_backend
+    )
+
+
 class DiffusionGemmaSelfConditioning(nn.Module):
     """Gated MLP that processes soft embeddings from the previous denoising step.
 
@@ -130,7 +138,7 @@ def get_mm_max_tokens_per_item(
         return super().get_mm_max_tokens_per_item(seq_len, mm_counts)
 
 
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+@_compile_diffusion_fn
 def _softcap_logits(logits: torch.Tensor, cap: float) -> torch.Tensor:
     # fp32 before tanh for numerical stability (matches HF DiffusionGemma).
     # Compiling fuses the cast/div/tanh/mul into one elementwise kernel over
@@ -354,7 +362,7 @@ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
         raise ValueError(f"Unsupported modality: {modality}")
 
 
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+@_compile_diffusion_fn
 def _compute_num_rejected(
     num_logits: torch.Tensor,
     num_sampled: torch.Tensor,
@@ -366,7 +374,7 @@ def _compute_num_rejected(
     return torch.where(is_denoise, query_lens, num_rejected)
 
 
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
+@_compile_diffusion_fn
 def _compiled_sample_step(
     # Logits from the model [num_decode * CL, vocab]
     logits: torch.Tensor,
@@ -543,7 +551,7 @@ def _compiled_sample_step(
     if tp_size > 1:
         soft_embeds = torch.ops.vllm.all_reduce(soft_embeds, group_name=tp_group_name)
     soft_embeds = soft_embeds * normalizer
-    sc_embeds[decode_slots] = soft_embeds * sc_keep
+    sc_embeds[decode_slots] = (soft_embeds * sc_keep).to(sc_embeds.dtype)
 
     # Overwrite canvas with argmax for newly converged denoise requests
     newly_converged = (converged & is_denoise).unsqueeze(1)
@@ -571,6 +579,7 @@ def __init__(
         max_denoising_steps: int,
         device: torch.device,
         hidden_size: int,
+        hidden_dtype: torch.dtype,
         stability_threshold: int,
     ):
         self.max_num_reqs = max_num_reqs
@@ -629,7 +638,11 @@ def __init__(
         # vocab/hidden (~170x) and moves the matmul to denoise time; the result
         # is identical (SC consumes probs @ embed_weight anyway).
         self.self_conditioning_embeds = torch.zeros(
-            max_num_reqs, canvas_length, hidden_size, dtype=torch.float32, device=device
+            max_num_reqs,
+            canvas_length,
+            hidden_size,
+            dtype=hidden_dtype,
+            device=device,
         )
 
     def init_canvas(self, slot_indices: torch.Tensor) -> None:
@@ -698,6 +711,7 @@ def __init__(
             max_denoising_steps=max_denoising_steps,
             device=device,
             hidden_size=text_config.hidden_size,
+            hidden_dtype=self.model_config.dtype,
             # In Transformers, `stability_threshold=1` (the default) means the current
             # step must match the previous step. In vLLM, the history buffer includes
             # the current step, so we add 1 to match the same behavior.

vllm/v1/attention/backends/flash_attn.py

@@ -540,6 +540,10 @@ def _maybe_symmetrize_window(
     return window
 
 
+def _supports_dynamic_causal(fa_version: int) -> bool:
+    return fa_version == 4 or (current_platform.is_xpu() and fa_version == 2)
+
+
 class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetadata]):
     # FA3:
     # Supports full cudagraphs for all cases.
@@ -1365,9 +1369,10 @@ def forward(
 
                 dynamic_causal = None
                 if isinstance(causal, torch.Tensor):
-                    if self.vllm_flash_attn_version != 4:
+                    if not _supports_dynamic_causal(self.vllm_flash_attn_version):
                         raise NotImplementedError(
-                            "Per-sequence causal requires FA4. Current version: "
+                            "Per-sequence causal requires FA4 or XPU FA2. "
+                            "Current version: "
                             f"FA{self.vllm_flash_attn_version}"
                         )
                     dynamic_causal = causal