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

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
+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]