[XPU] Remove redundant xpu graph backend and make xpu graph opt-in by default (#29911)

This commit is contained in:
Cao E
2026-07-03 15:58:56 +08:00
committed by GitHub
parent 67697fb891
commit 9df16b5ba9
17 changed files with 68 additions and 140 deletions
+13 -4
View File
@@ -147,9 +147,18 @@ 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** |
| 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) |
### Enable Decode Graph
Decode graph capture is **opt-in** on XPU. Enable it explicitly:
```bash
python -m sglang.launch_server --model-path <MODEL> --device xpu \
--cuda-graph-backend-decode full
```
### Enable Prefill Graph
Prefill graph capture is **opt-in** on XPU and requires `torch.compile`
@@ -195,10 +204,10 @@ python -m sglang.launch_server --model-path <MODEL> --device xpu \
### Disable XPU Graph
To opt out of one or both phases:
Both phases are disabled by default. To explicitly disable them anyway:
```bash
# Disable decode graph
# Disable decode graph (already off by default; explicit form)
python -m sglang.launch_server --model-path <MODEL> --device xpu \
--cuda-graph-backend-decode=disabled
@@ -236,7 +245,7 @@ python -m sglang.launch_server \
| 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-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-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. |
@@ -103,11 +103,22 @@ class XpuPlatform(Platform):
if empty_cache:
torch.xpu.empty_cache()
# Use mem_get_info() with a sanity cap to avoid KV-cache over-allocation
# on drivers that incorrectly return total memory as free memory.
# Consistent with the fallback: free = max(0, total - allocated).
try:
free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(device_id)
used_memory = float(torch.xpu.memory_allocated(device_id))
free_gpu_memory = min(
float(free_gpu_memory),
max(0.0, float(total_gpu_memory) - used_memory),
)
except Exception:
# Fallback for devices/drivers that do not support querying free memory
used_memory = float(torch.xpu.memory_allocated(device_id))
total_gpu_memory = float(
torch.xpu.get_device_properties(device_id).total_memory
)
free_gpu_memory = max(0.0, total_gpu_memory - used_memory)
if distributed:
@@ -22,7 +22,6 @@ 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__)
@@ -144,9 +143,6 @@ class XPUGraphRunner(DecodeCudaGraphRunner):
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(
@@ -1,106 +0,0 @@
# Copyright 2023-2026 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.
# ==============================================================================
"""XPUCudaGraphBackend — Intel XPU full-graph capture (torch.xpu.XPUGraph)."""
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_backend.base_cuda_graph_backend import (
BaseCudaGraphBackend,
)
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 XPUCudaGraphBackend(BaseCudaGraphBackend):
"""One torch.xpu.XPUGraph per shape; attention metadata is
captured inside the graph.
"""
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 = self._device_module.graph_pool_handle()
self._capture_stream = stream
try:
yield
finally:
self._capture_stream = None
def capture_one(
self,
shape_key: Any,
forward_fn: Callable[[], Any],
dummies: Optional[Any] = None,
post_warmup_hook: Optional[Callable[[], None]] = None,
) -> None:
# Two warmups so kernels are loaded and one-time setup is paid before capture.
# post_warmup_hook lets the attention backend reset state that warmup mutated.
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()
# graph_ctx: Callable[..., AbstractContextManager]
graph_ctx = self._device_module.graph
with graph_ctx(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: Any) -> bool:
return shape_key in self._graphs
@contextmanager
def replay_session(self):
yield
def replay(
self,
shape_key: Any,
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
@@ -72,13 +72,6 @@ def resolve_decode_backend(
cuda_graph_runner, enable_memory_saver=enable_memory_saver
)
elif model_runner.device == "xpu":
from sglang.srt.hardware_backend.xpu.xpu_cudagraph_backend import (
XPUCudaGraphBackend,
)
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}'"
+16
View File
@@ -3054,6 +3054,22 @@ class ServerArgs:
def _handle_xpu_backends(self):
if self.device == "xpu":
# Decode graph is opt-in on XPU: unless the user explicitly set
# --cuda-graph-backend-decode (or --cuda-graph-config), keep it
# disabled so the default startup doesn't require graph capture.
if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked:
self.cuda_graph_config.decode.backend = Backend.DISABLED
elif self.cuda_graph_config.decode.backend not in (
Backend.DISABLED,
Backend.FULL,
):
logger.warning(
"XPU platform only supports decode backend 'full'; "
"disabling unsupported decode backend '%s'.",
self.cuda_graph_config.decode.backend,
)
self.cuda_graph_config.decode.backend = Backend.DISABLED
if self.cuda_graph_config.prefill.backend not in (
Backend.DISABLED,
Backend.TC_PIECEWISE,
+16 -2
View File
@@ -621,9 +621,23 @@ def get_available_gpu_memory(
if empty_cache:
empty_device_cache(torch.xpu)
# Use mem_get_info() to reflect true OS-level free memory
# including graph pool reservations; avoids KV-cache over-allocation.
# Use mem_get_info() with a sanity cap to avoid KV-cache over-allocation
# on drivers that incorrectly return total memory as free memory.
# Consistent with the fallback: free = max(0, total - allocated).
try:
free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(gpu_id)
used_memory = float(torch.xpu.memory_allocated(gpu_id))
free_gpu_memory = min(
float(free_gpu_memory),
max(0.0, float(total_gpu_memory) - used_memory),
)
except Exception:
# Fallback for devices/drivers that do not support querying free memory
used_memory = float(torch.xpu.memory_allocated(gpu_id))
total_gpu_memory = float(
torch.xpu.get_device_properties(gpu_id).total_memory
)
free_gpu_memory = max(0.0, total_gpu_memory - used_memory)
elif device == "hpu":
num_gpus = torch.hpu.device_count()
@@ -42,6 +42,10 @@ class TestEagle3ParityXPU(SpecParityKit, _Eagle3ParityBase):
disable_overlap = False
attention_backend = "triton"
# Decode full-graph was active by default when this test was added
# (via XPUCudaGraphBackend). Opt in explicitly now that it is disabled
# by default so the coverage is preserved.
extra_args = ("--cuda-graph-config", '{"decode":{"backend":"full"}}')
if __name__ == "__main__":
-1
View File
@@ -38,7 +38,6 @@ class TestDeepSeekOCR(CustomTestCase):
"xpu",
"--attention-backend",
"intel_xpu",
"--disable-decode-cuda-graph",
]
os.environ["SGLANG_USE_SGL_XPU"] = "1"
cls.process = popen_launch_server(
@@ -41,7 +41,6 @@ class TestDeepSeekOCRTriton(TestDeepSeekOCR):
"xpu",
"--attention-backend",
"intel_xpu",
"--disable-decode-cuda-graph",
]
os.environ["SGLANG_USE_SGL_XPU"] = "0"
cls.process = popen_launch_server(
@@ -39,7 +39,6 @@ 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(
@@ -128,7 +127,6 @@ 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(
-1
View File
@@ -53,7 +53,6 @@ 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).
@@ -34,7 +34,6 @@ 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 [])
@@ -31,7 +31,6 @@ 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 [])
-1
View File
@@ -33,7 +33,6 @@ class TestXPUBasic(CustomTestCase):
"0.6",
"--batch-size",
"1",
"--disable-decode-cuda-graph",
]
if is_in_ci():
args += ["--input", "64", "--output", "4"]
@@ -38,7 +38,6 @@ class TestXPUEmbedding(CustomTestCase):
"--is-embedding",
"--device",
"xpu",
"--disable-decode-cuda-graph",
],
)
cls.openai_url = cls.base_url + "/v1"
@@ -43,7 +43,7 @@ class TestXPUServingFeatures(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=["--device", "xpu", "--disable-decode-cuda-graph"],
other_args=["--device", "xpu"],
)
cls.openai_url = cls.base_url + "/v1"