diff --git a/docs_new/docs/hardware-platforms/xpu.mdx b/docs_new/docs/hardware-platforms/xpu.mdx index 1fc3fbb5a..ba465931b 100644 --- a/docs_new/docs/hardware-platforms/xpu.mdx +++ b/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 --device xpu \ ### Enable Prefill Graph -Prefill graph capture is **opt-in** on XPU and requires `torch.compile` -and must be enabled explicitly: +Prefill graph capture is **opt-in** on XPU 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 --device xpu \ @@ -178,6 +184,16 @@ python -m sglang.launch_server --model-path --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 --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] diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index df73ed97c..f484e3a90 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/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: - # On ROCm/HIP, cuda-python is unavailable, so use the portable torch API - # (which maps to the HIP runtime). On NVIDIA, keep querying the CUDA runtime - # directly via cuda-python: torch.cuda.is_current_stream_capturing() has - # proven unreliable there, so we preserve the original behavior. - if is_hip(): - with torch.cuda.stream(stream): - return torch.cuda.is_current_stream_capturing() +def _is_stream_capturing(stream: torch.Stream) -> bool: + # On ROCm/HIP and XPU, cuda-python is unavailable, so use the portable torch + # API (which maps to the HIP / XPU runtime). On NVIDIA, keep querying the + # CUDA runtime directly via cuda-python: torch.cuda.is_current_stream_capturing() + # has proven unreliable there, so we preserve the original behavior. + if is_hip() or _is_xpu: + with get_device_module().stream(stream): + 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 - is_self_cap = self is capturing or self.cuda_stream == cap_ptr - is_other_cap = other is capturing or other.cuda_stream == cap_ptr + cap_id = capturing.stream_id + is_self_cap = self is capturing or self.stream_id == cap_id + 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 - torch.cuda.Stream.wait_stream = _hooked_wait_stream # type: ignore[assignment] + _original_wait_stream = get_device_module().Stream.wait_stream + 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( - self, graph: torch.cuda.CUDAGraph, needs_instantiate: bool - ) -> None: + def _append_segment(self, graph, 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( - pool=self._pool, capture_error_mode=self._capture_error_mode - ) + if _is_xpu: + # 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: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 9b50d153f..1b0c733f6 100644 --- a/python/sglang/srt/server_args.py +++ b/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 # ------------------------------------------------------------------ diff --git a/test/registered/xpu/test_breakable_cuda_graph.py b/test/registered/xpu/test_breakable_cuda_graph.py new file mode 100644 index 000000000..680e2a259 --- /dev/null +++ b/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()