diff --git a/tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-XPU-FA2-FP8.yaml b/tests/evals/gsm8k/configs/DiffusionGemma-26B-A4B-it-XPU-FA2-FP8.yaml new file mode 100644 index 000000000000..1b4e090690b0 --- /dev/null +++ b/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 diff --git a/tests/kernels/attention/test_xpu_dynamic_causal.py b/tests/kernels/attention/test_xpu_dynamic_causal.py new file mode 100644 index 000000000000..4a9bbf5e5607 --- /dev/null +++ b/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) diff --git a/tests/models/test_diffusion_gemma_xpu.py b/tests/models/test_diffusion_gemma_xpu.py new file mode 100644 index 000000000000..14573c9b801d --- /dev/null +++ b/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 diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 5c19b014082e..0f03ff4106b4 100644 --- a/vllm/_xpu_ops.py +++ b/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 diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 0c1cc8c55e1a..cf24dc8da93b 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/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. diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index b57af3fd295e..46177505b138 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/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