[XPU] Enable breakable prefill CUDA graph on XPU (#30273)

docs_new/docs/hardware-platforms/xpu.mdx

@@ -149,6 +149,7 @@ SGLang enables XPU graph capture to reduce per-step kernel-launch overhead.
 |---|---|---|---|
 | Decode | `full` | One `torch.xpu.XPUGraph` per batch size, captured on startup | **Off** (opt-in) |
 | Prefill | `tc_piecewise` | `torch.compile` + XPU graph, one graph segment per token-length bucket | **Off** (opt-in) |
+| Prefill | `breakable` | Segmented `torch.xpu.XPUGraph` capture/replay (no `torch.compile`); eager break points at attention / MoE boundaries | **Off** (opt-in) |
 
 ### Enable Decode Graph
 
@@ -161,8 +162,13 @@ python -m sglang.launch_server --model-path <MODEL> --device xpu \
 
 ### Enable Prefill Graph
 
-Prefill graph capture is **opt-in** on XPU and requires `torch.compile`
+Prefill graph capture is **opt-in** on XPU and must be enabled explicitly.
-and must be enabled explicitly:
+Two backends are available: `tc_piecewise` and `breakable`.
+
+#### tc_piecewise
+
+Uses `torch.compile` plus an XPU graph, one graph segment per token-length
+bucket:
 
 ```bash
 python -m sglang.launch_server --model-path <MODEL> --device xpu \
@@ -178,6 +184,16 @@ python -m sglang.launch_server --model-path <MODEL> --device xpu \
     --cuda-graph-tc-compiler inductor
 ```
 
+#### breakable
+
+Captures the transformer stack as segmented `XPUGraph`s with eager break points
+at attention / MoE boundaries, without `torch.compile`:
+
+```bash
+python -m sglang.launch_server --model-path <MODEL> --device xpu \
+    --cuda-graph-backend-prefill breakable
+```
+
 You can also configure both phases together with a single `--cuda-graph-config` JSON argument:
 
 ```bash
@@ -246,7 +262,7 @@ python -m sglang.launch_server \
 | Argument | XPU allowed values | Default | Description |
 |---|---|---|---|
 | `--cuda-graph-backend-decode` | `full`, `disabled` | `disabled` | Backend for the decode phase. Only `full` is supported on XPU. Set to `full` to enable. |
-| `--cuda-graph-backend-prefill` | `tc_piecewise`, `disabled` | `disabled`* | Backend for the prefill phase. Must be set to `tc_piecewise` explicitly to enable. |
+| `--cuda-graph-backend-prefill` | `tc_piecewise`, `breakable`, `disabled` | `disabled`* | Backend for the prefill phase. Set to `tc_piecewise` or `breakable` 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. |
@@ -265,7 +281,6 @@ via `--cuda-graph-backend-prefill` or `--cuda-graph-config`.
 |---|---|
 | 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]

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -37,10 +37,12 @@ except ImportError:
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
     checkCudaErrors,
 )
-from sglang.srt.utils import is_hip
+from sglang.srt.utils import get_device_module, is_hip, is_xpu
 
 logger = logging.getLogger(__name__)
 
+_is_xpu = is_xpu()
+
 __all__ = [
     "eager_on_graph",
     "BreakableCUDAGraph",
@@ -63,18 +65,18 @@ def _check_cuda_bindings():
 _current_capture_var: ContextVar["BreakableCUDAGraphCapture | None"] = ContextVar(
     "current_capture", default=None
 )
-_current_stream_var: ContextVar[torch.cuda.Stream | None] = ContextVar(
+_current_stream_var: ContextVar[torch.Stream | None] = ContextVar(
     "current_stream", default=None
 )
-_forked_streams_var: ContextVar[set[torch.cuda.Stream] | None] = ContextVar(
+_forked_streams_var: ContextVar[set[torch.Stream] | None] = ContextVar(
     "forked_streams", default=None
 )
 
 
-def get_current_stream(device: torch.device | None = None) -> torch.cuda.Stream:
+def get_current_stream(device: torch.device | None = None) -> torch.Stream:
     stream = _current_stream_var.get()
     if stream is None:
-        return torch.cuda.current_stream(device)
+        return get_device_module().current_stream(device)
     return stream
 
 
@@ -84,14 +86,14 @@ def _capture_status(stream_ptr: int) -> "rt.cudaStreamCaptureStatus":
     return status
 
 
-def _is_stream_capturing(stream: torch.cuda.Stream) -> bool:
+def _is_stream_capturing(stream: torch.Stream) -> bool:
-    # On ROCm/HIP, cuda-python is unavailable, so use the portable torch API
+    # On ROCm/HIP and XPU, cuda-python is unavailable, so use the portable torch
-    # (which maps to the HIP runtime). On NVIDIA, keep querying the CUDA runtime
+    # API (which maps to the HIP / XPU runtime). On NVIDIA, keep querying the
-    # directly via cuda-python: torch.cuda.is_current_stream_capturing() has
+    # CUDA runtime directly via cuda-python: torch.cuda.is_current_stream_capturing()
-    # proven unreliable there, so we preserve the original behavior.
+    # has proven unreliable there, so we preserve the original behavior.
-    if is_hip():
+    if is_hip() or _is_xpu:
-        with torch.cuda.stream(stream):
+        with get_device_module().stream(stream):
-            return torch.cuda.is_current_stream_capturing()
+            return get_device_module().is_current_stream_capturing()
     return (
         _capture_status(stream.cuda_stream)
         == rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
@@ -107,7 +109,7 @@ _hook_lock = threading.Lock()
 _hook_refcount = 0
 
 
-def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream):
+def _hooked_wait_stream(self: torch.Stream, other: torch.Stream):
     assert _original_wait_stream is not None
     forked = _forked_streams_var.get()
     if forked is None:
@@ -118,9 +120,9 @@ def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream):
         _original_wait_stream(self, other)
         return
 
-    cap_ptr = capturing.cuda_stream
+    cap_id = capturing.stream_id
-    is_self_cap = self is capturing or self.cuda_stream == cap_ptr
+    is_self_cap = self is capturing or self.stream_id == cap_id
-    is_other_cap = other is capturing or other.cuda_stream == cap_ptr
+    is_other_cap = other is capturing or other.stream_id == cap_id
 
     if is_self_cap and not is_other_cap:
         if not _is_stream_capturing(other):
@@ -138,8 +140,8 @@ def _install_wait_stream_hook():
     global _original_wait_stream, _hook_refcount
     with _hook_lock:
         if _hook_refcount == 0:
-            _original_wait_stream = torch.cuda.Stream.wait_stream
+            _original_wait_stream = get_device_module().Stream.wait_stream
-            torch.cuda.Stream.wait_stream = _hooked_wait_stream  # type: ignore[assignment]
+            get_device_module().Stream.wait_stream = _hooked_wait_stream  # type: ignore[assignment]
         _hook_refcount += 1
 
 
@@ -149,7 +151,7 @@ def _uninstall_wait_stream_hook():
         _hook_refcount -= 1
         if _hook_refcount == 0:
             assert _original_wait_stream is not None, "wait_stream hook not installed"
-            torch.cuda.Stream.wait_stream = _original_wait_stream  # type: ignore[assignment]
+            get_device_module().Stream.wait_stream = _original_wait_stream  # type: ignore[assignment]
             _original_wait_stream = None
 
 
@@ -270,7 +272,7 @@ class BreakableCUDAGraph:
         self._deduped_cuda_graph = deduped_cuda_graph
 
     def replay(self) -> None:
-        stream = torch.cuda.current_stream()
+        stream = get_device_module().current_stream()
         token = _current_stream_var.set(stream)
         try:
             for i, seg in enumerate(self._segments):
@@ -280,9 +282,7 @@ class BreakableCUDAGraph:
         finally:
             _current_stream_var.reset(token)
 
-    def _append_segment(
+    def _append_segment(self, graph, needs_instantiate: bool) -> None:
-        self, graph: torch.cuda.CUDAGraph, needs_instantiate: bool
-    ) -> None:
         if self._deduped_cuda_graph is not None:
             self._segments.append(self._deduped_cuda_graph.register(graph))
             return
@@ -306,7 +306,7 @@ class BreakableCUDAGraphCapture:
         self,
         cuda_graph: BreakableCUDAGraph,
         pool=None,
-        stream: torch.cuda.Stream | None = None,
+        stream: torch.Stream | None = None,
         capture_error_mode: str = "global",
     ):
         assert isinstance(
@@ -320,17 +320,17 @@ class BreakableCUDAGraphCapture:
         self._capture_token = None
         self._stream_token = None
         self._forked_token = None
-        self._current_graph: torch.cuda.CUDAGraph | None = None
+        self._current_graph = None
         self._current_graph_needs_instantiate = False
 
     def __enter__(self):
         _install_wait_stream_hook()
         if self._stream is not None:
-            self._stream_ctx = torch.cuda.stream(self._stream)
+            self._stream_ctx = get_device_module().stream(self._stream)
             self._stream_ctx.__enter__()
         self._capture_token = _current_capture_var.set(self)
         self._stream_token = _current_stream_var.set(
-            self._stream or torch.cuda.current_stream()
+            self._stream or get_device_module().current_stream()
         )
         self._forked_token = _forked_streams_var.set(set())
         self._begin_new_segment()
@@ -350,20 +350,27 @@ class BreakableCUDAGraphCapture:
         return False
 
     def _begin_new_segment(self) -> None:
+        graph_cls = torch.xpu.XPUGraph if _is_xpu else torch.cuda.CUDAGraph
         # keep_graph retains the raw graph for dedup; skip it on the plain path.
+        # Dedup is CUDA-only (it introspects the raw graph via cuda-python), so
+        # XPU always takes the plain path below.
         if self.cuda_graph._deduped_cuda_graph is not None:
             try:
-                graph = torch.cuda.CUDAGraph(keep_graph=True)
+                graph = graph_cls(keep_graph=True)
                 self._current_graph_needs_instantiate = True
             except TypeError:
-                graph = torch.cuda.CUDAGraph()
+                graph = graph_cls()
                 self._current_graph_needs_instantiate = False
         else:
-            graph = torch.cuda.CUDAGraph()
+            graph = graph_cls()
             self._current_graph_needs_instantiate = False
-        graph.capture_begin(
+        if _is_xpu:
-            pool=self._pool, capture_error_mode=self._capture_error_mode
+            # torch.xpu.XPUGraph.capture_begin takes only an optional pool.
-        )
+            graph.capture_begin(pool=self._pool)
+        else:
+            graph.capture_begin(
+                pool=self._pool, capture_error_mode=self._capture_error_mode
+            )
         self._current_graph = graph
 
     def _end_current_segment(self) -> None:

python/sglang/srt/server_args.py

@@ -3516,16 +3516,6 @@ class ServerArgs:
                 )
                 self.cuda_graph_config.decode.backend = Backend.DISABLED
 
-            if self.cuda_graph_config.prefill.backend not in (
-                Backend.DISABLED,
-                Backend.TC_PIECEWISE,
-            ):
-                logger.warning(
-                    "XPU platform currently only supports prefill tc_piecewise CUDA graph; "
-                    "disabling unsupported prefill backend."
-                )
-                self.cuda_graph_config.prefill.backend = Backend.DISABLED
-
     # ------------------------------------------------------------------
     # CUDA graph configuration resolution
     # ------------------------------------------------------------------

test/registered/xpu/test_breakable_cuda_graph.py

@@ -0,0 +1,339 @@
+"""Tests for the breakable CUDA graph (BCG) runner on XPU.
+
+Two test classes:
+- TestBreakableCUDAGraphBasic / TestCopyOutput / TestBreakGraphHelper:
+  unit tests for the core capture / replay mechanism (simple tensor ops).
+- TestXPUBreakableGraph: integration test — run a small Qwen model with the
+  breakable prefill CUDA graph backend via a single bench_one_batch invocation.
+"""
+
+import unittest
+
+import torch
+
+from sglang.srt.utils import get_device, get_device_module
+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 TestBreakableCUDAGraphBasic(CustomTestCase):
+    """Test basic breakable CUDA graph capture and replay."""
+
+    @classmethod
+    def setUpClass(cls):
+        if not get_device_module().is_available():
+            raise unittest.SkipTest(f"{get_device()} not available")
+
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
+            BreakableCUDAGraph,
+            BreakableCUDAGraphCapture,
+            eager_on_graph,
+        )
+
+        cls.BreakableCUDAGraph = BreakableCUDAGraph
+        cls.BreakableCUDAGraphCapture = BreakableCUDAGraphCapture
+        cls.eager_on_graph = staticmethod(eager_on_graph)
+        cls.device = torch.device(f"{get_device()}:0")
+
+    def test_no_break_capture_replay(self):
+        """Capture and replay without any graph breaks should work like normal CUDA graph."""
+        x = torch.zeros(4, device=self.device)
+        y = torch.zeros(4, device=self.device)
+
+        graph = self.BreakableCUDAGraph()
+        stream = get_device_module().Stream(self.device)
+        with self.BreakableCUDAGraphCapture(graph, stream=stream):
+            y.copy_(x + 1.0)
+
+        # Replay with new input
+        x.fill_(5.0)
+        graph.replay()
+        get_device_module().synchronize()
+        self.assertTrue(torch.allclose(y, torch.full((4,), 6.0, device=self.device)))
+
+    def test_single_break(self):
+        """A single graph break should split capture into two segments."""
+        x = torch.zeros(4, device=self.device)
+        intermediate = torch.zeros(4, device=self.device)
+        y = torch.zeros(4, device=self.device)
+
+        @self.eager_on_graph(enable=True)
+        def eager_op(src):
+            return src * 2.0
+
+        graph = self.BreakableCUDAGraph()
+        stream = get_device_module().Stream(self.device)
+        with self.BreakableCUDAGraphCapture(graph, stream=stream):
+            intermediate.copy_(x + 1.0)
+            broken = eager_op(intermediate)
+            y.copy_(broken + 3.0)
+
+        # Replay with new input
+        x.fill_(10.0)
+        graph.replay()
+        get_device_module().synchronize()
+        # x=10 -> intermediate=11 -> eager: 11*2=22 -> y=22+3=25
+        self.assertTrue(torch.allclose(y, torch.full((4,), 25.0, device=self.device)))
+
+    def test_multiple_breaks(self):
+        """Multiple graph breaks should produce correct chained results."""
+        x = torch.zeros(4, device=self.device)
+        y = torch.zeros(4, device=self.device)
+
+        @self.eager_on_graph(enable=True)
+        def add_one(src):
+            return src + 1.0
+
+        @self.eager_on_graph(enable=True)
+        def double(src):
+            return src * 2.0
+
+        graph = self.BreakableCUDAGraph()
+        stream = get_device_module().Stream(self.device)
+        with self.BreakableCUDAGraphCapture(graph, stream=stream):
+            t1 = x + 1.0  # graph segment 1
+            t2 = add_one(t1)  # break 1: eager
+            t3 = t2 + 1.0  # graph segment 2
+            t4 = double(t3)  # break 2: eager
+            y.copy_(t4)  # graph segment 3
+
+        # Replay: x=5 -> +1=6 -> add_one=7 -> +1=8 -> double=16
+        x.fill_(5.0)
+        graph.replay()
+        get_device_module().synchronize()
+        self.assertTrue(torch.allclose(y, torch.full((4,), 16.0, device=self.device)))
+
+    def test_eager_on_graph_disabled(self):
+        """@eager_on_graph(enable=False) should be a no-op passthrough."""
+
+        @self.eager_on_graph(enable=False)
+        def my_fn(x):
+            return x + 1.0
+
+        # Should just be the original function
+        t = torch.tensor([1.0, 2.0], device=self.device)
+        result = my_fn(t)
+        self.assertTrue(
+            torch.allclose(result, torch.tensor([2.0, 3.0], device=self.device))
+        )
+
+    def test_eager_on_graph_outside_capture(self):
+        """@eager_on_graph called outside capture should run the function directly."""
+
+        @self.eager_on_graph(enable=True)
+        def my_fn(x):
+            return x + 1.0
+
+        t = torch.tensor([1.0, 2.0], device=self.device)
+        result = my_fn(t)
+        self.assertTrue(
+            torch.allclose(result, torch.tensor([2.0, 3.0], device=self.device))
+        )
+
+    def test_replay_updates_output(self):
+        """Replay should produce different results when input buffers change."""
+        x = torch.zeros(4, device=self.device)
+        y = torch.zeros(4, device=self.device)
+
+        @self.eager_on_graph(enable=True)
+        def scale(src):
+            return src * 3.0
+
+        graph = self.BreakableCUDAGraph()
+        stream = get_device_module().Stream(self.device)
+        with self.BreakableCUDAGraphCapture(graph, stream=stream):
+            t = x + 1.0
+            t2 = scale(t)
+            y.copy_(t2)
+
+        # First replay: x=0 -> 0+1=1 -> 1*3=3
+        graph.replay()
+        get_device_module().synchronize()
+        self.assertTrue(torch.allclose(y, torch.full((4,), 3.0, device=self.device)))
+
+        # Second replay: x=10 -> 10+1=11 -> 11*3=33
+        x.fill_(10.0)
+        graph.replay()
+        get_device_module().synchronize()
+        self.assertTrue(torch.allclose(y, torch.full((4,), 33.0, device=self.device)))
+
+    def test_eager_output_is_held_strongly_for_replay_bridge(self):
+        """The replay closure must keep the eager output bridge buffer alive."""
+        x = torch.zeros(4, device=self.device)
+        y = torch.zeros(4, device=self.device)
+
+        @self.eager_on_graph(enable=True)
+        def scale(src):
+            return src * 3.0
+
+        graph = self.BreakableCUDAGraph()
+        stream = get_device_module().Stream(self.device)
+        with self.BreakableCUDAGraphCapture(graph, stream=stream):
+            t = x + 1.0
+            broken = scale(t)
+            y.copy_(broken)
+
+        replay_closure = graph._break_fns[0].__closure__ or ()
+        self.assertTrue(
+            any(cell.cell_contents is broken for cell in replay_closure),
+            "eager output bridge buffer must be strongly captured",
+        )
+
+
+class TestCopyOutput(CustomTestCase):
+    """Test the _copy_output helper for structured output writeback."""
+
+    @classmethod
+    def setUpClass(cls):
+        if not get_device_module().is_available():
+            raise unittest.SkipTest(f"{get_device()} not available")
+
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
+            _copy_output,
+        )
+
+        cls._copy_output = staticmethod(_copy_output)
+        cls.device = torch.device(f"{get_device()}:0")
+
+    def test_tensor_copy(self):
+        dst = torch.zeros(4, device=self.device)
+        src = torch.ones(4, device=self.device) * 5.0
+        result = self._copy_output(dst, src)
+        self.assertIs(result, dst)
+        self.assertTrue(torch.allclose(dst, src))
+
+    def test_dict_copy(self):
+        dst = {
+            "a": torch.zeros(4, device=self.device),
+            "b": torch.zeros(4, device=self.device),
+        }
+        src = {
+            "a": torch.ones(4, device=self.device),
+            "b": torch.ones(4, device=self.device) * 2.0,
+        }
+        result = self._copy_output(dst, src)
+        self.assertIs(result, dst)
+        self.assertTrue(torch.allclose(dst["a"], torch.ones(4, device=self.device)))
+        self.assertTrue(
+            torch.allclose(dst["b"], torch.ones(4, device=self.device) * 2.0)
+        )
+
+    def test_object_copy(self):
+        class FakeOutput:
+            def __init__(self, t, label):
+                self.tensor = t
+                self.label = label
+
+        dst = FakeOutput(torch.zeros(4, device=self.device), "old")
+        src = FakeOutput(torch.ones(4, device=self.device) * 3.0, "new")
+        result = self._copy_output(dst, src)
+        self.assertIs(result, dst)
+        self.assertTrue(
+            torch.allclose(dst.tensor, torch.ones(4, device=self.device) * 3.0)
+        )
+        self.assertEqual(dst.label, "new")
+
+    def test_non_tensor_fallback(self):
+        result = self._copy_output(42, 99)
+        self.assertEqual(result, 99)
+
+
+class TestBreakGraphHelper(CustomTestCase):
+    """Test the break_graph() convenience function."""
+
+    @classmethod
+    def setUpClass(cls):
+        if not get_device_module().is_available():
+            raise unittest.SkipTest(f"{get_device()} not available")
+
+        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
+            BreakableCUDAGraph,
+            BreakableCUDAGraphCapture,
+            break_graph,
+        )
+
+        cls.BreakableCUDAGraph = BreakableCUDAGraph
+        cls.BreakableCUDAGraphCapture = BreakableCUDAGraphCapture
+        cls.break_graph = staticmethod(break_graph)
+        cls.device = torch.device(f"{get_device()}:0")
+
+    def test_break_graph_inserts_segment(self):
+        """break_graph() should insert a graph break even though it does nothing."""
+        x = torch.zeros(4, device=self.device)
+        y = torch.zeros(4, device=self.device)
+
+        graph = self.BreakableCUDAGraph()
+        stream = get_device_module().Stream(self.device)
+        with self.BreakableCUDAGraphCapture(graph, stream=stream):
+            t = x + 1.0
+            self.break_graph()
+            y.copy_(t + 2.0)
+
+        x.fill_(10.0)
+        graph.replay()
+        get_device_module().synchronize()
+        # x=10 -> +1=11 -> break -> +2=13
+        self.assertTrue(torch.allclose(y, torch.full((4,), 13.0, device=self.device)))
+
+
+class TestXPUBreakableGraph(CustomTestCase):
+    """Integration: breakable prefill CUDA graph on XPU via bench_one_batch.
+
+    The prefill graph shapes are pinned with --cuda-graph-bs-prefill; capturing
+    the full default shape range exhausts the level-zero backend on the current
+    XPU stack, so a small explicit set keeps capture within device limits.
+    """
+
+    def test_breakable_graph_runs(self):
+        args = [
+            *_COMMON_ARGS,
+            "--cuda-graph-config",
+            '{"prefill":{"backend":"breakable"}}',
+            "--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 breakable XPU prefill graph",
+        )
+        self.assertGreater(
+            decode_throughput,
+            0,
+            "decode throughput must be > 0 with breakable XPU prefill graph",
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()