[XPU] Enable XPU graph support (decode full-graph + prefill tc_piecewise) (#29053)

docs_new/docs/hardware-platforms/xpu.mdx

@@ -141,6 +141,124 @@ Additionally, the requests can be formed with
 [OpenAI Completions API](../basic_usage/openai_api_completions)
 and sent via the command line (e.g. using `curl`) or via your own script.
 
+## XPU Graph [Experimental]
+
+SGLang enables XPU graph capture to reduce per-step kernel-launch overhead.
+
+| Phase | Backend | Mechanism | Default |
+|---|---|---|---|
+| Decode | `full` | One `torch.xpu.XPUGraph` per batch size, captured on startup | **On** |
+| Prefill | `tc_piecewise` | `torch.compile` + XPU graph, one graph segment per token-length bucket | **Off** (opt-in) |
+
+### Enable Prefill Graph
+
+Prefill graph capture is **opt-in** on XPU and requires `torch.compile`
+and must be enabled explicitly:
+
+```bash
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --cuda-graph-backend-prefill tc_piecewise
+```
+
+By default the prefill subgraphs are compiled with `eager` mode. Switch to
+`inductor` for higher-quality generated code at the cost of longer startup:
+
+```bash
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --cuda-graph-backend-prefill tc_piecewise \
+    --cuda-graph-tc-compiler inductor
+```
+
+You can also configure both phases together with a single `--cuda-graph-config` JSON argument:
+
+```bash
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --cuda-graph-config '{"decode":{"backend":"full"},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}'
+```
+
+### Enable torch.compile for Decode
+
+`--enable-torch-compile` adds a `torch.compile` pass on top of the decode
+XPU graph: the model forward is compiled first, and the compiled forward is
+then captured as an `XPUGraph`. This can reduce per-kernel overhead further
+but increases startup time.
+
+```bash
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --enable-torch-compile
+```
+
+> **Note:** `--enable-torch-compile` is mutually exclusive with the prefill
+> `tc_piecewise` graph (the compatibility rules auto-disable it). Use them
+> separately or lock the prefill backend explicitly via `--cuda-graph-config`
+> if you need both.
+
+### Disable XPU Graph
+
+To opt out of one or both phases:
+
+```bash
+# Disable decode graph
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --cuda-graph-backend-decode=disabled
+
+# Disable prefill graph (already off by default; explicit form)
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --cuda-graph-backend-prefill=disabled
+
+# Disable both phases
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --cuda-graph-backend-decode=disabled \
+    --cuda-graph-backend-prefill=disabled
+```
+
+### Customize Capture Buckets
+
+By default, prefill capture sizes are derived from `--chunked-prefill-size`.
+To specify explicit token-length buckets:
+
+```bash
+python -m sglang.launch_server \
+    --model-path <MODEL> --device xpu \
+    --cuda-graph-backend-prefill tc_piecewise \
+    --cuda-graph-bs-prefill 64 128 256 512
+```
+
+To specify explicit decode graph batch sizes:
+
+```bash
+python -m sglang.launch_server \
+    --model-path <MODEL> --device xpu \
+    --cuda-graph-bs-decode 1 2 4 8
+```
+
+### Server Args
+
+| Argument | XPU allowed values | Default | Description |
+|---|---|---|---|
+| `--cuda-graph-backend-decode` | `full`, `disabled` | `full` | Backend for the decode phase. Only `full` is supported on XPU. |
+| `--cuda-graph-backend-prefill` | `tc_piecewise`, `disabled` | `disabled`* | Backend for the prefill phase. Must be set to `tc_piecewise` explicitly to enable. |
+| `--cuda-graph-tc-compiler` | `eager`, `inductor` | `eager` | Compiler for `tc_piecewise` prefill subgraphs. `inductor` produces more optimized code but has longer startup. |
+| `--cuda-graph-bs-prefill` | list of ints | auto | Explicit token-length buckets to capture for prefill. |
+| `--cuda-graph-bs-decode` | list of ints | auto | Explicit batch sizes to capture for decode. |
+| `--cuda-graph-config` | JSON string | — | One-shot JSON config for both phases, e.g. `'{"decode":{"backend":"full"},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}'`. Overrides all per-phase flags. |
+| `--disable-decode-cuda-graph` | — | `False` | Shorthand for `--cuda-graph-backend-decode=disabled`. |
+| `--disable-prefill-cuda-graph` | — | `False` | Shorthand for `--cuda-graph-backend-prefill=disabled`. |
+| `--enable-torch-compile` | — | `False` | Apply `torch.compile` on top of the decode XPU graph for further kernel optimization. |
+| `--torch-compile-max-bs` | int | `32` | Maximum batch size compiled by `torch.compile` when `--enable-torch-compile` is set. |
+
+\* Prefill graph is auto-disabled on XPU unless you lock the backend explicitly
+via `--cuda-graph-backend-prefill` or `--cuda-graph-config`.
+
+### Limitations
+
+| Feature | Status |
+|---|---|
+| Memory saver (`--enable-memory-saver`) | Not yet supported |
+| Two-batch overlap (`--enable-two-batch-overlap`) | Not yet supported |
+| Breakable CUDA graph | Not yet supported |
+| Speculative decoding | Not yet implemented |
+
 ## Prefill-Decode (P/D) Disaggregation on Intel XPU [Experimental]
 
 SGLang supports prefill-decode disaggregation on Intel XPU using the [NIXL](https://github.com/ai-dynamo/nixl) KV-transfer backend.

python/sglang/srt/compilation/backend.py

@@ -23,9 +23,10 @@ from sglang.srt.compilation.compiler_interface import EagerAdapter, InductorAdap
 from sglang.srt.compilation.cuda_piecewise_backend import CUDAPiecewiseBackend
 from sglang.srt.compilation.npu_piecewise_backend import NPUPiecewiseBackend
 from sglang.srt.compilation.pass_manager import PostGradPassManager
+from sglang.srt.compilation.xpu_piecewise_backend import XPUPiecewiseBackend
 from sglang.srt.environ import envs
 from sglang.srt.platforms import current_platform
-from sglang.srt.utils.common import is_npu
+from sglang.srt.utils.common import is_npu, is_xpu
 
 logger = logging.getLogger(__name__)
 
@@ -53,6 +54,8 @@ def make_backend(
 
     if current_platform.is_out_of_tree():
         backend_cls = current_platform.get_piecewise_backend_cls()
+    elif is_xpu():
+        backend_cls = XPUPiecewiseBackend
     elif is_npu():
         backend_cls = NPUPiecewiseBackend
     else:

python/sglang/srt/compilation/weak_ref_tensor.py

@@ -2,14 +2,16 @@ from typing import Any, Union
 
 import torch
 
-from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu
+from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu, is_xpu
 
-if is_cuda() or is_hip() or is_musa():
+if is_cuda() or is_hip() or is_musa() or is_xpu():
     from sgl_kernel import weak_ref_tensor
 elif is_npu():
     from torch_npu._C import _weak_ref_tensor as weak_ref_tensor
 else:
-    raise NotImplementedError("weak_ref_tensor is implemented only for CUDA and NPU.")
+    raise NotImplementedError(
+        "weak_ref_tensor is implemented only for CUDA, XPU, and NPU."
+    )
 
 
 def weak_ref_tensors(

python/sglang/srt/compilation/xpu_piecewise_backend.py

@@ -0,0 +1,99 @@
+from contextlib import ExitStack
+from typing import Any
+from unittest.mock import patch
+
+import torch
+
+from sglang.srt.compilation.compilation_counter import compilation_counter
+from sglang.srt.compilation.compile_phase import (
+    get_pcg_capture_stream,
+    is_in_torch_compile_warmup,
+)
+from sglang.srt.compilation.cuda_piecewise_backend import (
+    CUDAPiecewiseBackend,
+    weak_ref_tensors,
+)
+
+
+class XPUPiecewiseBackend(CUDAPiecewiseBackend):
+    def __call__(self, *args) -> Any:
+        if not self.first_run_finished:
+            self.first_run_finished = True
+            self.check_for_ending_compilation()
+            return self.compiled_graph_for_general_shape(*args)
+
+        if len(self.sym_shape_indices) == 0:
+            return self.compiled_graph_for_general_shape(*args)
+
+        runtime_shape = args[self.sym_shape_indices[0]]
+        if runtime_shape not in self.concrete_size_entries:
+            return self.compiled_graph_for_general_shape(*args)
+
+        entry = self.concrete_size_entries[runtime_shape]
+
+        if entry.runnable is None:
+            entry.runnable = self.compiled_graph_for_general_shape
+
+        if entry.need_to_compile and not entry.compiled:
+            entry.compiled = True
+            self.to_be_compiled_sizes.remove(runtime_shape)
+            entry.runnable = self.sglang_backend.compiler_manager.compile(
+                self.graph,
+                args,
+                self.inductor_config,
+                graph_index=self.piecewise_compile_index,
+                num_graphs=self.total_piecewise_compiles,
+                runtime_shape=runtime_shape,
+            )
+
+            if self.is_last_graph and not self.to_be_compiled_sizes:
+                self.check_for_ending_compilation()
+
+        if is_in_torch_compile_warmup():
+            return entry.runnable(*args)
+
+        if entry.cudagraph is None:
+            if entry.num_finished_warmup < 1:  # noqa
+                entry.num_finished_warmup += 1
+                return entry.runnable(*args)
+
+            stream = get_pcg_capture_stream()
+            assert (
+                stream is not None
+            ), "PCG capture stream is not set, please check if runtime recompilation happened"
+
+            if self.compile_config.get_enable_debug_mode():
+                entry.input_addresses = [
+                    x.data_ptr() for x in args if isinstance(x, torch.Tensor)
+                ]
+
+            xpugraph = torch.xpu.XPUGraph()
+
+            with ExitStack() as stack:
+                if not self.is_first_graph:
+                    stack.enter_context(patch("gc.collect", lambda: None))
+                    stack.enter_context(patch("torch.xpu.empty_cache", lambda: None))
+
+                with torch.xpu.graph(
+                    xpu_graph=xpugraph, pool=self.graph_pool, stream=stream
+                ):
+                    output = entry.runnable(*args)
+                    if self.is_last_graph:
+                        output = weak_ref_tensors(output)
+
+            entry.output = weak_ref_tensors(output)
+            entry.cudagraph = xpugraph
+
+            compilation_counter.num_cudagraph_captured += 1
+            return output
+
+        if self.compile_config.get_enable_debug_mode():
+            new_input_addresses = [
+                x.data_ptr() for x in args if isinstance(x, torch.Tensor)
+            ]
+            assert new_input_addresses == entry.input_addresses, (
+                "Input addresses for cudagraphs are different during replay."
+                f" Expected {entry.input_addresses}, got {new_input_addresses}"
+            )
+        entry.cudagraph.replay()
+        return entry.output

python/sglang/srt/distributed/parallel_state.py

@@ -517,7 +517,7 @@ class GroupCoordinator:
     def graph_capture(
         self,
         graph_capture_context: Optional[GraphCaptureContext] = None,
-        stream: Optional[torch.cuda.Stream] = None,
+        stream=None,
     ):
         if graph_capture_context is None:
             if stream is None:
@@ -606,7 +606,15 @@ class GroupCoordinator:
             return self.hpu_communicator.all_reduce(input_)
 
         if self.xpu_communicator is not None and not self.xpu_communicator.disabled:
-            return self.xpu_communicator.all_reduce(input_)
+            # Route through inplace_all_reduce custom op so Dynamo treats this as
+            # an opaque call and does not decompose it into _c10d_functional primitives
+            # (which invoke sycl_event.wait() and break XPU graph capture).
+            # Keeps the operation in-place; the all-reduce is performed by
+            # _all_reduce_in_place, which for XPU falls through to
+            # torch.distributed.all_reduce on self.device_group (the same group
+            # used by xpu_communicator).
+            inplace_all_reduce(input_, group_name=self.unique_name)
+            return input_
 
         if self.npu_communicator is not None and not self.npu_communicator.disabled:
             return self.npu_communicator.all_reduce(input_)
@@ -1002,9 +1010,13 @@ class GroupCoordinator:
         return envs.SGLANG_ENABLE_DETERMINISTIC_INFERENCE.get()
 
     def all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor):
-        if _is_npu or _is_xpu:
+        if _is_npu:
             self._all_gather_into_tensor(output, input)
         else:
+            # XPU and CUDA both go through reg_all_gather_into_tensor (custom_op) to
+            # stay opaque to Dynamo. Calling torch.distributed.all_gather_into_tensor
+            # directly causes Dynamo to rewrite it as _c10d_functional.all_gather_into_tensor
+            # + wait_tensor, which invokes sycl_event.wait() and breaks XPU graph capture.
             reg_all_gather_into_tensor(output, input, group_name=self.unique_name)
 
     def cp_all_gather_into_tensor_async(
@@ -1762,7 +1774,7 @@ def get_mooncake_transfer_engine():
 
 
 @contextmanager
-def graph_capture(stream: Optional[torch.cuda.Stream] = None):
+def graph_capture(stream=None):
     """
     `graph_capture` is a context manager which should surround the code that
     is capturing the CUDA graph. Its main purpose is to ensure that the

python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_full_graph_backend.py

@@ -0,0 +1,100 @@
+"""FullXPUGraphBackend — Intel XPU full-graph capture (torch.xpu.XPUGraph).
+
+Mirrors FullCudaGraphBackend with XPU-specific differences:
+  - Captures via torch.xpu.graph(xpu_graph=...) into torch.xpu.XPUGraph.
+  - Shares the global graph memory pool with the prefill backend so that
+    decode + prefill graphs occupy max(decode, prefill) rather than their sum.
+  - No set_graph_pool_id: SymmetricMemoryContext is never triggered on XPU
+    (oneCCL has no ncclMemAlloc equivalent; enable_symm_mem defaults False).
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import torch
+
+from sglang.srt.model_executor.runner.shape_key import ShapeKey
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+
+
+class FullXPUGraphBackend(BaseCudaGraphBackend):
+    """One torch.xpu.XPUGraph per shape for Intel XPU devices."""
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+    ) -> None:
+        self._graphs: Dict[Any, torch.xpu.XPUGraph] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = cuda_graph_runner.model_runner.tp_group
+        self._capture_stream: Optional[torch.xpu.Stream] = None
+
+    @contextmanager
+    def capture_session(self, stream: torch.xpu.Stream):
+        if self._pool is None:
+            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = torch.xpu.XPUGraph()
+
+        with self._device_module.graph(
+            xpu_graph=graph, pool=self._pool, stream=self._capture_stream
+        ):
+            out = forward_fn()
+
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None

python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py

@@ -0,0 +1,173 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Run the model with xpu graph and torch.compile."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import torch
+from torch.profiler import ProfilerActivity, profile
+
+from sglang.srt.model_executor.runner import DecodeCudaGraphRunner
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.utils import register_xpu_device_properties_for_dynamo
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+_fake_ops_registered = False
+
+
+def register_fake_ops():
+    """Register fake/abstract implementations for XPU sgl_kernel ops so that
+    torch.compile (Dynamo) can trace through them using FakeTensors for shape
+    and dtype propagation, without executing the real GPU kernels.
+    """
+    global _fake_ops_registered
+    if _fake_ops_registered:
+        return
+    _fake_ops_registered = True
+
+    @torch.library.register_fake("sgl_kernel::fwd")
+    def _(
+        q,
+        k,
+        v,
+        q_v,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlen_q,
+        max_seqlen_k,
+        page_table,
+        kv_batch_idx,
+        leftpad_k,
+        rotary_cos,
+        rotary_sin,
+        seqlens_rotary,
+        q_descale,
+        k_descale,
+        v_descale,
+        softmax_scale,
+        sinks,
+        is_causal,
+        window_size_left,
+        window_size_right,
+        softcap,
+        is_rotary_interleaved,
+        scheduler_metadata,
+        num_kv_splits,
+        pack_gqa,
+        sm_margin,
+        out=None,
+    ):
+        total_q = q.shape[0]
+        num_heads_q = q.shape[1]
+        head_size_v = v.shape[-1]
+        if out is None:
+            out = q.new_empty(total_q, num_heads_q, head_size_v)
+        softmax_lse = q.new_empty(num_heads_q, total_q, dtype=torch.float32)
+        # out_accum and softmax_lse_accum are intermediate split-kv buffers;
+        # they are only read when num_kv_splits > 1, which is determined at
+        # runtime.  Return empty tensors with correct rank so downstream ops
+        # that index into the list do not fail shape propagation.
+        out_accum = q.new_empty(0)
+        softmax_lse_accum = q.new_empty(0, dtype=torch.float32)
+        return (out, softmax_lse, out_accum, softmax_lse_accum)
+
+    @torch.library.register_fake("sgl_kernel::flash_mla_decode")
+    def _(
+        out,
+        q_nope,
+        q_pe,
+        kv_c_and_k_pe_cache,
+        seq_lens,
+        page_table,
+        workspace,
+        sm_scale,
+        num_kv_splits,
+    ):
+        return
+
+
+class XPUGraphRunner(DecodeCudaGraphRunner):
+    """A XPUGraphRunner runs the forward pass of a model with xpu graph and torch.compile."""
+
+    @staticmethod
+    def _apply_xpu_compile_config() -> None:
+        """Apply XPU-specific torch.compile / dynamo settings.
+
+        Called unconditionally before super().__init__() so that the settings
+        are in place regardless of whether --enable-torch-compile is passed.
+        The critical flag is suppress_errors: when the Intel IGC compiler
+        crashes with SIGFPE on certain reduction kernels (ocloc -device bmg
+        returns exit code 245), dynamo falls back to eager for that subgraph
+        instead of propagating the crash.
+        """
+        import torch._dynamo.config
+
+        torch._dynamo.config.suppress_errors = True
+
+    def __init__(self, model_runner: ModelRunner):
+        assert (
+            not model_runner.server_args.enable_memory_saver
+        ), "XPUGraphRunner does not support Torch Memory Saver yet."
+        register_fake_ops()
+        self._apply_xpu_compile_config()
+        register_xpu_device_properties_for_dynamo()
+        super().__init__(model_runner)
+
+        assert (
+            not self.enable_two_batch_overlap
+        ), "XPUGraphRunner does not support two batch overlap yet."
+        assert (
+            not self.require_mlp_tp_gather
+        ), "XPUGraphRunner does not support MLP TP gather yet."
+        assert (
+            not self.require_mlp_sync
+        ), "XPUGraphRunner does not support MLP sync yet."
+        assert (
+            not self.require_gathered_buffer
+        ), "XPUGraphRunner does not support gathered buffer yet."
+        assert (
+            model_runner.spec_algorithm == SpeculativeAlgorithm.NONE
+        ), "XPUGraphRunner does not support speculative inference yet."
+
+    def _init_profile_context_and_memory_record(self):
+        profile_context = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.XPU],
+            record_shapes=True,
+        )
+        torch.xpu.memory._record_memory_history()
+        return profile_context
+
+    def _post_process_after_profile(self, prof_context):
+        torch.xpu.memory._dump_snapshot("xpu_graph_runner_memory_usage.pickle")
+        torch.xpu.memory._record_memory_history(enabled=None)
+        log_message = (
+            "Sorted by XPU Time:\n"
+            + prof_context.key_averages(group_by_input_shape=True).table(
+                sort_by="self_xpu_time_total"
+            )
+            + "\n\nSorted by CPU Time:\n"
+            + prof_context.key_averages(group_by_input_shape=True).table(
+                sort_by="self_cpu_time_total"
+            )
+            + "\n\nMemory Usage is saved to xpu_graph_runner_memory_usage.pickle\n"
+        )
+        logger.info(log_message)

python/sglang/srt/layers/attention/fla/layernorm_gated.py

@@ -6,6 +6,7 @@
 # The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.
 
 
+from contextlib import nullcontext
 from functools import lru_cache
 
 import torch
@@ -254,7 +255,21 @@ def _layer_norm_fwd(
     # Update grid to use rows_per_block
     grid = (cdiv(M, rows_per_block), ngroups)
     pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
-    with device_context(x.device):
+    # Workaround for PyTorch <= 2.12: torch.xpu.device is not Dynamo-compatible
+    # in that release — it creates a DynamoConfigPatchProxy that
+    # SourcelessBuilder cannot wrap, causing a hard error under
+    # torch.compile(fullgraph=True).  The device context is a functional no-op
+    # for Triton kernel launches (device is determined by the tensor, not the
+    # surrounding context), so we simply skip it when Dynamo is tracing.
+    # PyTorch main already has the proper fix (XPUDeviceVariable registered in
+    # torch/_dynamo/variables/ctx_manager.py analogous to CUDADeviceVariable).
+    # TODO: remove this branch once we upgrade from PyTorch 2.12.
+    device_ctx = (
+        nullcontext()
+        if x.device.type == "xpu" and torch.compiler.is_compiling()
+        else device_context(x.device)
+    )
+    with device_ctx:
         _layer_norm_fwd_1pass_kernel[grid](
             x,
             out,

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

@@ -104,6 +104,7 @@ class XPUAttentionBackend(AttentionBackend):
         self.has_swa = (
             self.sliding_window_size is not None and self.sliding_window_size > -1
         )
+        self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder
 
     def init_forward_metadata(self, forward_batch: ForwardBatch):
         """Initialize forward metadata hence all layers in the forward pass can reuse it."""
@@ -411,22 +412,6 @@ class XPUAttentionBackend(AttentionBackend):
                     workspace_size, device=self.device, dtype=torch.uint8
                 )
 
-        # Translate full-pool indices to SWA-pool indices for hybrid models
-        if self.use_sliding_window_kv_pool:
-            # flash_attn_with_kvcache requires int32 page tables; the SWA index
-            # mapping is int64, so cast (matches flashattention_backend.py).
-            metadata.swa_page_table = (
-                self.token_to_kv_pool.translate_loc_from_full_to_swa(
-                    metadata.page_table
-                ).to(torch.int32)
-            )
-            if forward_batch.out_cache_loc is not None:
-                metadata.swa_out_cache_loc = (
-                    self.token_to_kv_pool.translate_loc_from_full_to_swa(
-                        forward_batch.out_cache_loc
-                    )
-                )
-
         # Convert the page table to a strided format which is needed by FA3 API
         if self.page_size > 1:
             self.strided_indices = torch.arange(
@@ -477,7 +462,10 @@ class XPUAttentionBackend(AttentionBackend):
                 if not self.use_mla:
                     self.token_to_kv_pool.set_kv_buffer(
                         layer,
-                        KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
+                        KVWriteLoc(
+                            cache_loc,
+                            self.forward_metadata.swa_out_cache_loc,
+                        ),
                         k,
                         v,
                         layer.k_scale,
@@ -599,6 +587,19 @@ class XPUAttentionBackend(AttentionBackend):
                 k_descale=k_descale,
                 v_descale=v_descale,
                 return_softmax_lse=use_cascade_attn,
+                # Piecewise XPU graph for prefill requires a pre-allocated
+                # output buffer at a stable device address so the graph can
+                # record writes to the same storage on every replay.
+                # _attn_output is that fixed buffer; None falls back to a
+                # freshly allocated tensor (eager / cascade-attn path).
+                out=(
+                    forward_batch._attn_output.view(
+                        -1, layer.tp_q_head_num, layer.v_head_dim
+                    )
+                    if not use_cascade_attn
+                    and getattr(forward_batch, "_attn_output", None) is not None
+                    else None
+                ),
                 **kwargs,
             )
 
@@ -783,7 +784,10 @@ class XPUAttentionBackend(AttentionBackend):
                 if not self.use_mla:
                     self.token_to_kv_pool.set_kv_buffer(
                         layer,
-                        KVWriteLoc(cache_loc, self.forward_metadata.swa_out_cache_loc),
+                        KVWriteLoc(
+                            cache_loc,
+                            self.forward_metadata.swa_out_cache_loc,
+                        ),
                         k,
                         v,
                         layer.k_scale,
@@ -997,6 +1001,208 @@ class XPUAttentionBackend(AttentionBackend):
         """Get the fill value for sequence length in CUDA graph."""
         return 1
 
+    def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
+        """Pre-allocate fixed-size tensors reused across XPU graph captures."""
+        max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size
+        self.decode_cuda_graph_metadata = {
+            "cache_seqlens": torch.zeros(max_bs, dtype=torch.int32, device=self.device),
+            "cu_seqlens_q": torch.arange(
+                0, max_bs + 1, dtype=torch.int32, device=self.device
+            ),
+            "cu_seqlens_k": torch.zeros(
+                max_bs + 1, dtype=torch.int32, device=self.device
+            ),
+            "page_table": torch.zeros(
+                max_bs, max_num_pages, dtype=torch.int32, device=self.device
+            ),
+            "strided_indices": torch.arange(
+                0, self.max_context_len, self.page_size, device=self.device
+            ),
+        }
+        if self.use_sliding_window_kv_pool:
+            self.decode_cuda_graph_metadata["swa_page_table"] = torch.zeros(
+                max_bs, max_num_pages, dtype=torch.int32, device=self.device
+            )
+            self.decode_cuda_graph_metadata["swa_out_cache_loc"] = torch.zeros(
+                max_num_tokens, dtype=torch.int64, device=self.device
+            )
+        if self.is_encoder_decoder:
+            self.encoder_metadata = {
+                "encoder_page_table": torch.zeros(
+                    max_bs, self.max_context_len, dtype=torch.int32, device=self.device
+                ),
+                "encoder_lens_int32": torch.zeros(
+                    max_bs, dtype=torch.int32, device=self.device
+                ),
+                "encoder_cu_seqlens_k": torch.zeros(
+                    max_bs + 1, dtype=torch.int32, device=self.device
+                ),
+            }
+        else:
+            self.encoder_metadata = {}
+
+    def init_forward_metadata_out_graph(
+        self,
+        forward_batch: ForwardBatch,
+        in_capture: bool = False,
+    ):
+        """New unified graph capture/replay entry point (replaces the legacy
+        init_forward_metadata_capture_cuda_graph /
+        init_forward_metadata_replay_cuda_graph pair).
+
+        Called by DecodeCudaGraphRunner:
+          - capture: in_capture=True  → bind static metadata buffer slices, then fill
+          - replay:  in_capture=False → update pre-allocated buffers in-place
+          - eager:   via init_forward_metadata() default wrapper
+        """
+        bs = forward_batch.batch_size
+        req_pool_indices = forward_batch.req_pool_indices
+        seq_lens = forward_batch.seq_lens
+        seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
+        forward_mode = forward_batch.forward_mode
+        spec_info = forward_batch.spec_info
+
+        assert (
+            spec_info is None
+        ), "XPUAttentionBackend does not support speculative decoding in XPU graph"
+        assert (
+            forward_mode.is_decode_or_idle()
+        ), "XPUAttentionBackend XPU graph only supports decode mode"
+
+        if in_capture:
+            # Bind static-shape slices of the pre-allocated buffers so the
+            # captured graph always reads from the same storage addresses.
+            metadata = FlashAttentionMetadata()
+            metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[
+                "cache_seqlens"
+            ][:bs]
+            metadata.cu_seqlens_q = self.decode_cuda_graph_metadata["cu_seqlens_q"][
+                : bs + 1
+            ]
+            metadata.cu_seqlens_k = self.decode_cuda_graph_metadata["cu_seqlens_k"][
+                : bs + 1
+            ]
+            metadata.page_table = self.decode_cuda_graph_metadata["page_table"][:bs, :]
+            if self.use_sliding_window_kv_pool:
+                # Bind SWA page table slice so the graph captures the right tensor.
+                metadata.swa_page_table = self.decode_cuda_graph_metadata[
+                    "swa_page_table"
+                ][:bs, :]
+            if self.is_encoder_decoder and forward_batch.encoder_lens is not None:
+                encoder_bs = forward_batch.encoder_lens.numel()
+                metadata.encoder_lens_int32 = self.encoder_metadata[
+                    "encoder_lens_int32"
+                ][:encoder_bs]
+                metadata.encoder_cu_seqlens_k = self.encoder_metadata[
+                    "encoder_cu_seqlens_k"
+                ][: encoder_bs + 1]
+                metadata.encoder_page_table = self.encoder_metadata[
+                    "encoder_page_table"
+                ][:bs, :]
+            self.decode_cuda_graph_metadata[bs] = metadata
+
+        # Both capture and replay: fill data into the pre-allocated buffers.
+        seq_lens = seq_lens[:bs]
+        seq_lens_cpu = seq_lens_cpu[:bs] if seq_lens_cpu is not None else None
+        req_pool_indices = req_pool_indices[:bs]
+
+        metadata = self.decode_cuda_graph_metadata[bs]
+        max_len = (
+            seq_lens_cpu.max().item()
+            if seq_lens_cpu is not None
+            else seq_lens.max().item()
+        )
+        metadata.max_seq_len_k = max_len
+
+        metadata.cache_seqlens_int32.copy_(seq_lens.to(torch.int32))
+
+        metadata.cu_seqlens_k[0] = 0
+        metadata.cu_seqlens_k[1 : bs + 1].copy_(
+            torch.cumsum(seq_lens.to(torch.int32), dim=0)
+        )
+
+        if self.is_encoder_decoder and forward_batch.encoder_lens is not None:
+            encoder_lens = forward_batch.encoder_lens[:bs].to(torch.int32)
+            metadata.encoder_max_seq_len_k = int(encoder_lens.max().item())
+            metadata.encoder_lens_int32.copy_(encoder_lens)
+            metadata.encoder_cu_seqlens_k[0] = 0
+            metadata.encoder_cu_seqlens_k[1 : bs + 1].copy_(
+                torch.cumsum(encoder_lens, dim=0, dtype=torch.int32)
+            )
+            metadata.encoder_page_table[:bs, : metadata.encoder_max_seq_len_k].copy_(
+                self.req_to_token[
+                    req_pool_indices, : metadata.encoder_max_seq_len_k
+                ].to(torch.int32)
+            )
+            # Self-attention (text) page_table: decoder tokens start after encoder tokens.
+            text_max = metadata.max_seq_len_k
+            arange_text = torch.arange(text_max, device=req_pool_indices.device)
+            text_col = encoder_lens[:bs].long().unsqueeze(1) + arange_text.unsqueeze(0)
+            text_row = req_pool_indices.unsqueeze(1).expand(-1, text_max)
+            metadata.page_table[:bs, :text_max].copy_(
+                self.req_to_token[text_row, text_col].to(torch.int32)
+            )
+            metadata.page_table[:bs, text_max:].zero_()
+        else:
+            raw_page = self.req_to_token[
+                req_pool_indices[:, None],
+                self.decode_cuda_graph_metadata["strided_indices"][
+                    : ((metadata.max_seq_len_k + self.page_size - 1) // self.page_size)
+                ][None, :],
+            ]
+            if self.page_size > 1:
+                raw_page = raw_page // self.page_size
+            metadata.page_table[:bs, : raw_page.shape[1]].copy_(
+                raw_page.to(torch.int32)
+            )
+            metadata.page_table[:bs, raw_page.shape[1] :].zero_()
+
+        if self.use_sliding_window_kv_pool:
+            if forward_batch.out_cache_loc is None:
+                raise ValueError(
+                    f"out_cache_loc is None for hybrid SWA model in graph "
+                    f"{'capture' if in_capture else 'replay'} "
+                    f"(forward_mode={forward_batch.forward_mode}). This should not happen."
+                )
+            swa_out_cache_loc = self.decode_cuda_graph_metadata["swa_out_cache_loc"]
+            n = forward_batch.out_cache_loc.shape[0]
+            swa_out_cache_loc[n:].zero_()
+            swa_out_cache_loc[:n].copy_(
+                self.token_to_kv_pool.translate_loc_from_full_to_swa(
+                    forward_batch.out_cache_loc
+                )
+            )
+            metadata.swa_out_cache_loc = swa_out_cache_loc[:n]
+
+            if not (self.is_encoder_decoder and forward_batch.encoder_lens is not None):
+                max_seq_pages = (
+                    metadata.max_seq_len_k + self.page_size - 1
+                ) // self.page_size
+                swa_page_table = self.decode_cuda_graph_metadata["swa_page_table"]
+                swa_page_table[:bs, max_seq_pages:].zero_()
+                swa_page_table[:bs, :max_seq_pages].copy_(
+                    (
+                        self.token_to_kv_pool.translate_loc_from_full_to_swa(raw_page)
+                        if self.page_size == 1
+                        else self.token_to_kv_pool.translate_loc_from_full_to_swa(
+                            self.req_to_token[
+                                req_pool_indices[:, None],
+                                self.decode_cuda_graph_metadata["strided_indices"][
+                                    :max_seq_pages
+                                ][None, :],
+                            ]
+                        )
+                        // self.page_size
+                    ).to(torch.int32)
+                )
+                metadata.swa_page_table = swa_page_table[:bs, :]
+
+        self.forward_metadata = metadata
+
+    def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
+        """Graph-recordable ops for XPU graph (no-op: all metadata setup is
+        host-side and lives in init_forward_metadata_out_graph)."""
+
     def _init_local_attn_metadata(
         self,
         forwardbatch: ForwardBatch,

python/sglang/srt/model_executor/model_runner.py

@@ -112,6 +112,7 @@ from sglang.srt.eplb.lplb_solver import (
     set_global_lplb_solver,
 )
 from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
+from sglang.srt.hardware_backend.xpu.graph_runner.xpu_graph_runner import XPUGraphRunner
 from sglang.srt.kv_canary.api import install_canary
 from sglang.srt.kv_canary.runner.canary_manager import context_tuple
 from sglang.srt.kv_canary.token_oracle.install import install_token_oracle_from_env
@@ -879,7 +880,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         if self.device == "cuda" or self.device == "musa":
             self.init_cublas()
             self.init_attention_backend()
-        elif self.device == "cpu":
+        elif self.device in ["cpu", "xpu"]:
             self.init_attention_backend()
         elif self.device == "npu":
             self.init_attention_backend()
@@ -922,7 +923,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         self.graph_mem_usage = 0
 
         if capture_decode_cuda_graph:
-            if self.device in ("cuda", "musa", "cpu", "npu"):
+            if self.device in ("cuda", "musa", "cpu", "npu", "xpu"):
                 self.init_decode_cuda_graph()
             elif (
                 current_platform.is_out_of_tree()
@@ -2601,6 +2602,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 "musa": "CUDA graph",
                 "cpu": "CPU graph",
                 "npu": "NPU graph",
+                "xpu": "XPU graph",
             },
         )
         role = "draft" if self.is_draft_worker else "target"
@@ -2636,6 +2638,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 {
                     "cpu": CPUGraphRunner,
                     "npu": NPUGraphRunner,
+                    "xpu": XPUGraphRunner,
                 },
             )
             self.decode_cuda_graph_runner = graph_runners[self.device](self)

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -148,12 +148,17 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
         enable_torch_compile_warmup to drive FX / inductor through
         every shape without capturing cuda graphs yet."""
         language_model = self._language_model
+        # Some multimodal models (e.g. Gemma4) store the inner transformer
+        # directly as `language_model` rather than wrapping it in a
+        # ForCausalLM that has a `.model` child.  Fall back to the module
+        # itself when `.model` is absent.
+        inner_model = getattr(language_model, "model", language_model)
         compiler = self._compile_config.compiler
         with enable_tc_piecewise_cuda_graph():
             try:
                 if compiler != "eager":
                     _toggle_multi_platform_ops(
-                        language_model.model, reverse=False, num_tokens=16
+                        inner_model, reverse=False, num_tokens=16
                     )
 
                 cuda_graph_runner._run_dummy_forward(
@@ -167,7 +172,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
                 set_graph_pool_id(self._pool)
 
                 self.install_compile(
-                    language_model.model,
+                    inner_model,
                     compile_config=self._compile_config,
                     graph_pool=self._pool,
                 )
@@ -196,9 +201,7 @@ class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
                                 )
                             cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens)
             finally:
-                _toggle_multi_platform_ops(
+                _toggle_multi_platform_ops(inner_model, reverse=True, num_tokens=16)
-                    language_model.model, reverse=True, num_tokens=16
-                )
 
     @contextmanager
     def capture_session(self, stream: torch.cuda.Stream):

python/sglang/srt/model_executor/runner_backend/utils.py

@@ -78,6 +78,17 @@ def resolve_decode_backend(
 
         return XPUCudaGraphBackend(cuda_graph_runner)
 
+    if model_runner.device == "xpu":
+        if backend_name not in (Backend.FULL, Backend.DISABLED):
+            raise ValueError(
+                f"XPU only supports cuda_graph_config decode backend 'full', got '{backend_name}'"
+            )
+        from sglang.srt.hardware_backend.xpu.graph_runner.xpu_full_graph_backend import (
+            FullXPUGraphBackend,
+        )
+
+        return FullXPUGraphBackend(cuda_graph_runner)
+
     if backend_name == Backend.BREAKABLE:
         return BreakableCudaGraphBackend(
             cuda_graph_runner,

python/sglang/srt/server_args.py

@@ -3054,10 +3054,13 @@ class ServerArgs:
 
     def _handle_xpu_backends(self):
         if self.device == "xpu":
-            if self.cuda_graph_config.prefill.backend != Backend.DISABLED:
+            if self.cuda_graph_config.prefill.backend not in (
+                Backend.DISABLED,
+                Backend.TC_PIECEWISE,
+            ):
                 logger.warning(
-                    "XPU platform does not support piecewise CUDA graph, "
+                    "XPU platform currently only supports prefill tc_piecewise CUDA graph; "
-                    "disabling prefill cuda graph."
+                    "disabling unsupported prefill backend."
                 )
                 self.cuda_graph_config.prefill.backend = Backend.DISABLED
 

python/sglang/srt/utils/common.py

@@ -179,6 +179,17 @@ def is_xpu() -> bool:
     return hasattr(torch, "xpu") and torch.xpu.is_available()
 
 
+def register_xpu_device_properties_for_dynamo() -> None:
+    if not is_xpu():
+        return
+
+    import torch._dynamo.utils as dynamo_utils
+
+    xpu_props_type = getattr(torch.xpu, "_XpuDeviceProperties", None)
+    if xpu_props_type is not None:
+        dynamo_utils.common_constant_types.add(xpu_props_type)
+
+
 @lru_cache(maxsize=1)
 def is_npu() -> bool:
     if not hasattr(torch, "npu"):
@@ -610,9 +621,9 @@ def get_available_gpu_memory(
 
         if empty_cache:
             empty_device_cache(torch.xpu)
-        used_memory = torch.xpu.memory_allocated()
+        # Use mem_get_info() to reflect true OS-level free memory
-        total_gpu_memory = torch.xpu.get_device_properties(gpu_id).total_memory
+        # including graph pool reservations; avoids KV-cache over-allocation.
-        free_gpu_memory = total_gpu_memory - used_memory
+        free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(gpu_id)
 
     elif device == "hpu":
         num_gpus = torch.hpu.device_count()

test/registered/xpu/test_deepseek_ocr.py

@@ -38,6 +38,7 @@ class TestDeepSeekOCR(CustomTestCase):
             "xpu",
             "--attention-backend",
             "intel_xpu",
+            "--disable-decode-cuda-graph",
         ]
         os.environ["SGLANG_USE_SGL_XPU"] = "1"
         cls.process = popen_launch_server(

test/registered/xpu/test_deepseek_ocr_triton.py

@@ -41,6 +41,7 @@ class TestDeepSeekOCRTriton(TestDeepSeekOCR):
             "xpu",
             "--attention-backend",
             "intel_xpu",
+            "--disable-decode-cuda-graph",
         ]
         os.environ["SGLANG_USE_SGL_XPU"] = "0"
         cls.process = popen_launch_server(

test/registered/xpu/test_encoder_attention_backend.py

@@ -39,6 +39,7 @@ class TestEncoderAttention(CustomTestCase):
             "xpu",
             "--mm-attention-backend",
             "xpu_attn",
+            "--disable-decode-cuda-graph",
         ]
         os.environ["SGLANG_USE_SGL_XPU"] = "1"
         cls.process = popen_launch_server(
@@ -127,6 +128,7 @@ class TestEncoderAttention_Triton(TestEncoderAttention):
             "xpu",
             "--mm-attention-backend",
             "triton_attn",
+            "--disable-decode-cuda-graph",
         ]
         os.environ["SGLANG_USE_SGL_XPU"] = "0"
         cls.process = popen_launch_server(

test/registered/xpu/test_gemma_4_e2b.py

@@ -53,6 +53,7 @@ XPU_SERVER_ARGS = [
     "intel_xpu",
     "--model-impl",
     "sglang",
+    "--disable-decode-cuda-graph",
 ]
 
 # Standard sglang e2e Q&A prompt (see test_openai_server.py::run_chat_completion).

test/registered/xpu/test_intel_xpu_backend.py

@@ -34,6 +34,7 @@ def intel_xpu_benchmark(
                 "1",
                 "--device",
                 "xpu",
+                "--disable-decode-cuda-graph",
             ]
             ci_args = ["--input", "64", "--output", "4"] if is_in_ci() else []
             full_args = common_args + ci_args + (extra_args or [])

test/registered/xpu/test_triton_attention_backend.py

@@ -31,6 +31,7 @@ def triton_attention_benchmark(extra_args=None, mem_fraction_static="0.84"):
                 "2050",
                 "--attention-backend",
                 "triton",
+                "--disable-decode-cuda-graph",
             ]
             full_args = common_args + (extra_args or [])
 

test/registered/xpu/test_xpu_basic.py

@@ -33,6 +33,7 @@ class TestXPUBasic(CustomTestCase):
             "0.6",
             "--batch-size",
             "1",
+            "--disable-decode-cuda-graph",
         ]
         if is_in_ci():
             args += ["--input", "64", "--output", "4"]

test/registered/xpu/test_xpu_embedding.py

@@ -34,7 +34,12 @@ class TestXPUEmbedding(CustomTestCase):
             cls.model,
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=["--is-embedding", "--device", "xpu"],
+            other_args=[
+                "--is-embedding",
+                "--device",
+                "xpu",
+                "--disable-decode-cuda-graph",
+            ],
         )
         cls.openai_url = cls.base_url + "/v1"
 

test/registered/xpu/test_xpu_graph.py

@@ -0,0 +1,71 @@
+"""
+XPU graph tests: verifies decode full-graph and prefill tc_piecewise graph
+on Intel XPU produce valid outputs.
+
+  - TestXPUGraph : decode full-graph and prefill tc_piecewise graph enabled
+    together in a single bench_one_batch invocation.
+
+Usage:
+    python3 -m unittest test_xpu_graph.TestXPUGraph
+"""
+
+import unittest
+
+from sglang.test.ci.ci_register import register_xpu_ci
+from sglang.test.test_utils import (
+    DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN,
+    CustomTestCase,
+    is_in_ci,
+    run_bench_one_batch,
+)
+
+register_xpu_ci(est_time=600, suite="stage-b-test-1-gpu-xpu")
+
+_COMMON_ARGS = [
+    "--device",
+    "xpu",
+    "--attention-backend",
+    "triton",
+    "--disable-radix-cache",
+    "--mem-fraction-static",
+    "0.6",
+    "--batch-size",
+    "1",
+]
+
+_CI_IO_ARGS = ["--input", "64", "--output", "4"]
+_FULL_IO_ARGS = ["--input", "128", "--output", "16"]
+
+
+class TestXPUGraph(CustomTestCase):
+    """Decode full-graph + prefill tc_piecewise together."""
+
+    def test_full_graph_runs(self):
+        args = [
+            *_COMMON_ARGS,
+            "--cuda-graph-config",
+            '{"decode":{"backend":"full"},"prefill":{"backend":"tc_piecewise","tc_compiler":"eager"}}',
+            "--cuda-graph-bs-prefill",
+            "64",
+            "128",
+        ]
+        if is_in_ci():
+            args += _CI_IO_ARGS
+        else:
+            args += _FULL_IO_ARGS
+
+        prefill_latency, decode_throughput, _ = run_bench_one_batch(
+            DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, args
+        )
+        self.assertGreater(
+            prefill_latency,
+            0,
+            "prefill latency must be > 0 with tc_piecewise XPU graph",
+        )
+        self.assertGreater(
+            decode_throughput, 0, "decode throughput must be > 0 with full XPU graph"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/xpu/test_xpu_serving_features.py

@@ -43,7 +43,7 @@ class TestXPUServingFeatures(CustomTestCase):
             cls.model,
             cls.base_url,
             timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=["--device", "xpu"],
+            other_args=["--device", "xpu", "--disable-decode-cuda-graph"],
         )
         cls.openai_url = cls.base_url + "/v1"