From 0fadad89339a3b7869dec7b2148defd60dd69584 Mon Sep 17 00:00:00 2001 From: ANSHUMAN TRIPATHY Date: Fri, 11 Sep 2026 08:09:37 +0530 Subject: [PATCH] DFLASH support added for XPU (#32798) Co-authored-by: Ma Mingfei --- python/sglang/srt/arg_groups/choices.py | 1 + .../sglang/srt/arg_groups/speculative_hook.py | 20 +-- .../srt/layers/attention/xpu_backend.py | 17 +- python/sglang/srt/speculative/dflash_utils.py | 1 + .../srt/speculative/dflash_worker_v2.py | 6 +- .../srt/speculative/draft_worker_common.py | 10 +- .../e2e/xpu/test_dflash_intel_xpu.py | 156 ++++++++++++++++++ 7 files changed, 192 insertions(+), 19 deletions(-) create mode 100644 test/registered/e2e/xpu/test_dflash_intel_xpu.py diff --git a/python/sglang/srt/arg_groups/choices.py b/python/sglang/srt/arg_groups/choices.py index 1ec9c21ba..8ce2c00ec 100644 --- a/python/sglang/srt/arg_groups/choices.py +++ b/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 = [ diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index f433f4141..7db89b2c6 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/xpu_backend.py b/python/sglang/srt/layers/attention/xpu_backend.py index 60a2d6901..a689114ed 100644 --- a/python/sglang/srt/layers/attention/xpu_backend.py +++ b/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 = {} diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index f3c12445f..49273b240 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -36,6 +36,7 @@ _DFLASH_VERIFY_SKIP_CUSTOM_MASK_BACKENDS = frozenset( "TritonAttnBackend", "TRTLLMHAAttnBackend", "TRTLLMMLABackend", + "XPUAttentionBackend", } ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 8c537f3bb..a95674399 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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 diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index eb9ddb806..c3e197981 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/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'.", diff --git a/test/registered/e2e/xpu/test_dflash_intel_xpu.py b/test/registered/e2e/xpu/test_dflash_intel_xpu.py new file mode 100644 index 000000000..ce1157310 --- /dev/null +++ b/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 ` (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/). +_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()