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

[Feature][Intel XPU] Add memory saver support for Intel XPU via upstream torch_memory_saver (#29935)

Browse Source
This commit is contained in:
Siju Samuel
2026-09-08 09:38:51 +08:00
committed by GitHub
parent b8a81f055d
commit 2358916d5a
4 changed files with 474 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/utils/torch_memory_saver_adapter.py
+58 -6
View File
@@ -2,9 +2,17 @@ import logging
from abc import ABC
from contextlib import contextmanager
from sglang.srt.utils.common import is_xpu
try:
import torch_memory_saver
# Intel XPU requires hook_mode="torch" (in-process pluggable allocator);
# the LD_PRELOAD-based preload mode is CUDA/HIP-only. Set it before the
# singleton is initialized on first use.
if is_xpu():
torch_memory_saver.torch_memory_saver.hook_mode = "torch"
_memory_saver = torch_memory_saver.torch_memory_saver
import_error = None
except ImportError as e:
@@ -13,16 +21,31 @@ except ImportError as e:
logger = logging.getLogger(__name__)
_warned_xpu_cuda_graph = False
class TorchMemorySaverAdapter(ABC):
@staticmethod
def create(enable: bool):
if enable and import_error is not None:
logger.warning(
"enable_memory_saver is enabled, but "
"torch-memory-saver is not installed. Please install it "
"via `pip3 install torch-memory-saver`. "
)
if is_xpu():
# XPU ships no prebuilt wheel; it is built from source against the
# local oneAPI + torch-XPU runtime. TMS_PLATFORM=xpu forces the XPU
# backend; --no-build-isolation lets the build see torch and match
# the libsycl ABI to it.
logger.warning(
"enable_memory_saver is enabled, but torch-memory-saver is "
"not installed. On Intel XPU, build it from source with Intel "
"oneAPI on PATH: `TMS_PLATFORM=xpu pip3 install "
"--no-build-isolation git+https://github.com/fzyzcjy/"
"torch_memory_saver.git@a5c99f11b18ebb8e9fda71a68812e476ae49e417`."
)
else:
logger.warning(
"enable_memory_saver is enabled, but "
"torch-memory-saver is not installed. Please install it "
"via `pip3 install torch-memory-saver`. "
)
raise import_error
return (
_TorchMemorySaverAdapterReal() if enable else _TorchMemorySaverAdapterNoop()
@@ -59,17 +82,46 @@ class TorchMemorySaverAdapter(ABC):
class _TorchMemorySaverAdapterReal(TorchMemorySaverAdapter):
"""Adapter for TorchMemorySaver with tag-based control"""
"""Adapter for TorchMemorySaver with tag-based control.
Backed by the upstream torch_memory_saver package (CUDA VMM, and Intel XPU via
Level Zero). XPU requires the in-process pluggable allocator (hook_mode="torch")
instead of the CUDA LD_PRELOAD path, which is what makes configure_subprocess()
and cuda_graph() no-ops there; region/pause/resume are fully supported.
"""
def configure_subprocess(self):
if is_xpu():
# Nothing to preload: this LD_PRELOADs the preload-mode .so, which the
# upstream setup.py does not build for XPU.
return self._noop_context()
return torch_memory_saver.configure_subprocess()
def region(self, tag: str, enable_cpu_backup: bool = False):
return _memory_saver.region(tag=tag, enable_cpu_backup=enable_cpu_backup)
def cuda_graph(self, **kwargs):
if is_xpu():
# Upstream gates pauseable graph capture on hook_mode="preload" while XPU
# requires hook_mode="torch", so the two are mutually exclusive. Unreachable
# today (XPU routes to FullXPUGraphBackend, which takes no memory saver);
# warn rather than raise, so a future XPU graph backend that does route here
# surfaces that graph memory is not pauseable instead of failing to launch.
global _warned_xpu_cuda_graph
if not _warned_xpu_cuda_graph:
_warned_xpu_cuda_graph = True
logger.warning(
"torch_memory_saver cannot make CUDA-graph memory pauseable on Intel "
"XPU; graph allocations will not be released by "
"release_memory_occupation(tags=['cuda_graph'])."
)
return self._noop_context()
return _memory_saver.cuda_graph(**kwargs)
@contextmanager
def _noop_context(self, **kwargs):
yield
def disable(self):
return _memory_saver.disable()