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

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

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/platforms/xpu.py
+17 -6
View File
@@ -103,12 +103,23 @@ class XpuPlatform(Platform):
if empty_cache:
torch.xpu.empty_cache()
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)
# 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:
import torch.distributed as dist
python/sglang/srt/hardware_backend/xpu/graph_runner/xpu_graph_runner.py
-4
View File
@@ -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(
python/sglang/srt/hardware_backend/xpu/xpu_cudagraph_backend.py
-106
View File
@@ -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
python/sglang/srt/model_executor/runner_backend/utils.py
-7
View File
@@ -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}'"
python/sglang/srt/server_args.py
+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,
python/sglang/srt/utils/common.py
+17 -3
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.
free_gpu_memory, total_gpu_memory = torch.xpu.mem_get_info(gpu_id)
# 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()