minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Rahul Vijayaraghavan
2026-07-21 09:09:40 +08:00
committed by GitHub
parent bfefdc52d7
commit fa0ced195e
4 changed files with 399 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs_new/docs/hardware-platforms/xpu.mdx
+19 -4
View File
@@ -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`
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 <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
+41 -34
View File
@@ -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:
python/sglang/srt/server_args.py
-10
View File
@@ -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
+339
View File
@@ -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()