DFLASH support added for XPU (#32798)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

python/sglang/srt/arg_groups/choices.py

@@ -111,6 +111,7 @@ DRAFT_ATTENTION_BACKEND_CHOICES = [
     "triton",
     "ascend",
     "trtllm_mha",
+    "intel_xpu",
 ]
 
 DETERMINISTIC_ATTENTION_BACKEND_CHOICES = [

python/sglang/srt/arg_groups/speculative_hook.py

@@ -5,6 +5,7 @@ import logging
 import os
 from typing import TYPE_CHECKING, Optional
 
+from sglang.srt.arg_groups.choices import DRAFT_ATTENTION_BACKEND_CHOICES
 from sglang.srt.arg_groups.overrides import (
     _speculative_moe_runner_default,
     attention_backends_of,
@@ -198,9 +199,11 @@ def handle_speculative_decoding(server_args: ServerArgs) -> None:
 def _handle_dflash(server_args: ServerArgs) -> None:
     cfg = resolving_view(server_args)
 
-    if not (cfg.device.startswith("cuda") or cfg.device == "npu"):
+    if not (
+        cfg.device.startswith("cuda") or cfg.device == "npu" or cfg.device == "xpu"
+    ):
         raise ValueError(
-            "DFLASH speculative decoding only supports CUDA and NPU devices."
+            "DFLASH speculative decoding only supports CUDA, NPU and XPU devices."
         )
 
     if resolved_view(server_args).enable_dp_attention:
@@ -722,16 +725,11 @@ def _resolve_dflash_draft_attention_backend(server_args: ServerArgs) -> None:
     """
     cfg = resolving_view(server_args)
 
-    supported_draft_backends = (
-        "flashinfer",
-        "fa3",
-        "fa4",
-        "triton",
-        "trtllm_mha",
-        "ascend",
+    supported_draft_backends = DRAFT_ATTENTION_BACKEND_CHOICES
+    # FlashInfer is CUDA-only; fall back to triton on XPU and ROCm.
+    fallback_backend = (
+        "triton" if (get_platform().is_xpu or get_platform().is_hip) else "flashinfer"
     )
-    # Use triton on ROCm (no FlashInfer), flashinfer on CUDA.
-    fallback_backend = "triton" if get_platform().is_hip else "flashinfer"
 
     draft_backend = cfg.speculative_draft_attention_backend
     if draft_backend is None:

python/sglang/srt/layers/attention/xpu_backend.py

@@ -12,6 +12,7 @@ from sglang.srt.layers.attention.flashattention_backend import (
     merge_state_v2_wrapper,
     prepare_swa_spec_page_table_triton,
 )
+from sglang.srt.layers.radix_attention import AttentionType
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -553,7 +554,13 @@ class XPUAttentionBackend(AttentionBackend):
         #     q = q.to(self.kv_cache_dtype)
         #     q_rope = q_rope.to(self.kv_cache_dtype) if q_rope is not None else None
         #     k_rope = k_rope.to(self.kv_cache_dtype) if k_rope is not None else None
-        causal = not layer.is_cross_attention
+        # Mirror FlashAttentionBackend: ENCODER_ONLY / bidirectional layers
+        # (DFLASH draft full_attention) are non-causal, not just cross-attention.
+        causal = not (
+            layer.is_cross_attention
+            or layer.attn_type
+            in (AttentionType.ENCODER_ONLY, AttentionType.DECODER_BIDIRECTIONAL)
+        )
 
         # Check if we should use local attention
         use_local_attn = (
@@ -974,7 +981,13 @@ class XPUAttentionBackend(AttentionBackend):
             if layer.sliding_window_size is not None and layer.sliding_window_size > -1
             else (-1, -1)
         )
-        causal = not layer.is_cross_attention
+        # Mirror FlashAttentionBackend: ENCODER_ONLY / bidirectional layers
+        # (DFLASH draft full_attention) are non-causal, not just cross-attention.
+        causal = not (
+            layer.is_cross_attention
+            or layer.attn_type
+            in (AttentionType.ENCODER_ONLY, AttentionType.DECODER_BIDIRECTIONAL)
+        )
 
         # For fa3 interface version compatibility, we put new fields into conditional keyword args
         kwargs = {}

python/sglang/srt/speculative/dflash_utils.py

@@ -36,6 +36,7 @@ _DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS = frozenset(
         "TritonAttnBackend",
         "TRTLLMHAAttnBackend",
         "TRTLLMMLABackend",
+        "XPUAttentionBackend",
     }
 )
 

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -82,7 +82,7 @@ from sglang.srt.speculative.spec_utils import (
     assign_req_to_token_pool_func,
     build_grammar_vocab_mask,
 )
-from sglang.srt.utils import is_cuda, is_hip, is_npu
+from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
 
 _is_npu = is_npu()
 
@@ -503,12 +503,12 @@ class DFlashWorkerV2(BaseSpecWorker):
         self._draft_greedy_rank_index_buf: Optional[torch.Tensor] = None
         self._draft_greedy_selected_ids_buf: Optional[torch.Tensor] = None
         self._draft_greedy_index_cap: int = 0
-        self._use_fused_kv_materialize = is_cuda() or is_hip()
+        self._use_fused_kv_materialize = is_cuda() or is_hip() or is_xpu()
         self._fused_kv_helper: Optional[object] = None
         if self._use_fused_kv_materialize:
             self._init_fused_kv_helper()
 
-        supports_gpu_triton = is_cuda() or is_hip()
+        supports_gpu_triton = is_cuda() or is_hip() or is_xpu()
         self._use_triton_prepare_block = supports_gpu_triton
         self._use_triton_accept_bonus = supports_gpu_triton
         # The legacy compact-rebuild path host-syncs twice per step (masked

python/sglang/srt/speculative/draft_worker_common.py

@@ -9,7 +9,7 @@ import torch
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode
-from sglang.srt.runtime_context import attention_backends, get_spec
+from sglang.srt.runtime_context import attention_backends, get_platform, get_spec
 from sglang.srt.server_args import DRAFT_ATTENTION_BACKEND_CHOICES, ServerArgs
 from sglang.srt.speculative.dflash_info import DFlashVerifyInput
 from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
@@ -36,13 +36,17 @@ def _resolve_draft_attention_backend_fallback(*, algo_label: str) -> str:
     otherwise the process's prefill backend. Both are resolution's answers, so
     they come from the bags.
     """
+    # FlashInfer is CUDA-only; fall back to triton on XPU and ROCm.
+    platform_fallback = (
+        "triton" if (get_platform().is_xpu or torch.version.hip) else "flashinfer"
+    )
     draft_backend = get_spec().speculative_draft_attention_backend
     if draft_backend is None:
         draft_backend, _ = attention_backends()
     if draft_backend is None:
-        return "triton" if torch.version.hip else "flashinfer"
+        return platform_fallback
     if draft_backend not in DRAFT_ATTENTION_BACKEND_CHOICES:
-        fallback = "triton" if torch.version.hip else "flashinfer"
+        fallback = platform_fallback
         logger.warning(
             "%s draft worker only supports attention_backend in %s for now, "
             "but got %r. Falling back to '%s'.",

test/registered/e2e/xpu/test_dflash_intel_xpu.py

@@ -0,0 +1,156 @@
+"""DFLASH speculative decoding on Intel XPU."""
+
+import os
+import sys
+import unittest
+
+from sglang.srt.utils.common import is_xpu
+
+# Put the `test/` root on sys.path so `registered.<...>` resolves regardless of
+# cwd: CI runs each file as `python3 <full_path>` (only the file's own dir is on
+# the path), and pytest inserts only the file's dir too. `test/` is three levels
+# up from this file's dir (test/registered/xpu/e2e/<this>).
+_TEST_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+if _TEST_ROOT not in sys.path:
+    sys.path.insert(0, _TEST_ROOT)
+
+# Reference the base as a module attribute rather than importing the `Test*`
+# name directly: pytest collects by class __name__, so a bare
+# `from ... import TestDFlashServerBase` (even aliased) would make it re-collect
+# the base's CUDA/flashinfer config here. Only the XPU subclass should run.
+from registered.spec.dflash import test_dflash as _dflash_base
+
+from sglang.test.ci.ci_register import register_xpu_ci
+
+register_xpu_ci(est_time=600, suite="nightly-xpu-1-gpu", nightly=True)
+
+# Appended after the base launch_args by setUpClass: the trailing
+# --mem-fraction-static overrides the base 0.7, and --device selects the Intel
+# GPU. Variants that need extra flags must prepend these (the base does not
+# merge other_launch_args — the subclass value replaces it wholesale).
+_XPU_LAUNCH_ARGS = ["--device", "xpu", "--mem-fraction-static", "0.75"]
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPU(_dflash_base.TestDFlashServerBase):
+    """Full DFLASH suite on device=xpu with the triton attention backend."""
+
+    max_running_requests = 8
+    attention_backend = "triton"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUPage256(_dflash_base.TestDFlashServerPage256):
+    """page_size=256 + radix-attention smoke test on XPU."""
+
+    max_running_requests = 8
+    attention_backend = "triton"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUChunkedPrefill(_dflash_base.TestDFlashServerChunkedPrefill):
+    """Chunked prefill (size 4) on XPU."""
+
+    max_running_requests = 8
+    attention_backend = "triton"
+    # XPU args first, then the variant's own --chunked-prefill-size.
+    other_launch_args = _XPU_LAUNCH_ARGS + ["--chunked-prefill-size", "4"]
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUDecodeGraph(_dflash_base.TestDFlashServerBase):
+    """Decode CUDA-graph enabled on XPU (opt-in via --cuda-graph-backend-decode)."""
+
+    max_running_requests = 8
+    attention_backend = "triton"
+    other_launch_args = _XPU_LAUNCH_ARGS + ["--cuda-graph-backend-decode", "full"]
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUOverlap(_dflash_base.TestDFlashServerOverlap):
+    """Overlap schedule enabled on XPU."""
+
+    max_running_requests = 8
+    attention_backend = "triton"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUOverlapPlanStream(
+    _dflash_base.TestDFlashServerOverlapPlanStream
+):
+    """Overlap schedule with the plan stream on XPU."""
+
+    max_running_requests = 8
+    attention_backend = "triton"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+# --- Native XPUAttentionBackend (intel_xpu) variants --------------------------
+# Same configs as above, but exercising the native intel_xpu backend rather than
+# triton. intel_xpu is not the XPU default (triton is), so it must be selected
+# explicitly; these guard the DFLASH draft/verify path through XPUAttentionBackend.
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUBackend(_dflash_base.TestDFlashServerBase):
+    """Full DFLASH suite on device=xpu with the intel_xpu attention backend."""
+
+    max_running_requests = 8
+    attention_backend = "intel_xpu"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUBackendPage256(_dflash_base.TestDFlashServerPage256):
+    """page_size=256 + radix-attention smoke test on XPU (intel_xpu backend)."""
+
+    max_running_requests = 8
+    attention_backend = "intel_xpu"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUBackendChunkedPrefill(
+    _dflash_base.TestDFlashServerChunkedPrefill
+):
+    """Chunked prefill (size 128) on XPU (intel_xpu backend)."""
+
+    max_running_requests = 8
+    attention_backend = "intel_xpu"
+    other_launch_args = _XPU_LAUNCH_ARGS + ["--chunked-prefill-size", "128"]
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUBackendNoCudaGraph(_dflash_base.TestDFlashServerNoCudaGraph):
+    """CUDA-graph disabled on XPU (intel_xpu backend)."""
+
+    max_running_requests = 8
+    attention_backend = "intel_xpu"
+    other_launch_args = _XPU_LAUNCH_ARGS + ["--disable-cuda-graph"]
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUBackendOverlap(_dflash_base.TestDFlashServerOverlap):
+    """Overlap schedule enabled on XPU (intel_xpu backend)."""
+
+    max_running_requests = 8
+    attention_backend = "intel_xpu"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+@unittest.skipUnless(is_xpu(), "Intel XPU not available")
+class TestDFlashIntelXPUBackendOverlapPlanStream(
+    _dflash_base.TestDFlashServerOverlapPlanStream
+):
+    """Overlap schedule with the plan stream on XPU (intel_xpu backend)."""
+
+    max_running_requests = 8
+    attention_backend = "intel_xpu"
+    other_launch_args = _XPU_LAUNCH_ARGS
+
+
+if __name__ == "__main__":
+    unittest.main()