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
docker/xpu.Dockerfile
+12
View File
@@ -83,4 +83,16 @@ RUN echo "Cloning ${SG_LANG_BRANCH} from ${SG_LANG_REPO}" && \
pip install --no-cache-dir ".[dev,diffusion]" --extra-index-url https://download.pytorch.org/whl/xpu && \
pip install --no-cache-dir --no-deps xgrammar==0.1.33
# Install torch_memory_saver for release/resume_memory_occupation ("memory saver").
# XPU ships no prebuilt wheel: it is built from source against the local oneAPI +
# torch-XPU runtime (the .so links libsycl.so.<N>, which must match the installed
# intel-sycl-rt). TMS_PLATFORM=xpu forces the XPU backend; --no-build-isolation
# lets the build import the installed torch (above) so it can match the libsycl
# major to it -- under build isolation torch is absent and the match is skipped.
# Pinned (v0.0.10b2) so image builds are reproducible; bump via --build-arg.
ARG TORCH_MEMORY_SAVER_REF=a5c99f11b18ebb8e9fda71a68812e476ae49e417
RUN . /opt/intel/oneapi/setvars.sh --force >/dev/null 2>&1 && \
TMS_PLATFORM=xpu pip install --no-cache-dir --no-build-isolation \
git+https://github.com/fzyzcjy/torch_memory_saver.git@${TORCH_MEMORY_SAVER_REF}
CMD ["bash", "-c", "source /opt/intel/oneapi/setvars.sh --force && exec bash"]
docs/docs/hardware-platforms/xpu.mdx
+54 -1
View File
@@ -326,10 +326,63 @@ via `--cuda-graph-backend-prefill` or `--cuda-graph-config`.
| Feature | Status |
|---|---|
| Memory saver (`--enable-memory-saver`) | Not yet supported |
| Two-batch overlap (`--enable-two-batch-overlap`) | Not yet supported |
| Speculative decoding | Not yet implemented |
## Memory Saver (release/resume memory occupation) on Intel XPU [Experimental]
SGLang can temporarily release most of the GPU memory it holds — model weights
and/or KV cache — and reclaim it later without restarting the process. This is
the same `release_memory_occupation` / `resume_memory_occupation` feature
available on CUDA, used for RL rollout/training hand-off and for freeing the
device between inference bursts.
This is backed by the [`torch_memory_saver`](https://github.com/fzyzcjy/torch_memory_saver)
package — the same package used on CUDA — which gained an Intel XPU backend
built natively on Level Zero (keeping virtual addresses fixed while
releasing/re-committing physical pages via `zeVirtualMemUnmap` /
`zeVirtualMemMap`).
**Install `torch_memory_saver`.** Unlike CUDA (prebuilt wheel), the XPU backend
is built from source against your local oneAPI + `torch+xpu` runtime (the `.so`
links `libsycl.so.<N>`, which must match the installed `intel-sycl-rt`).
`TMS_PLATFORM=xpu` forces the XPU backend, and `--no-build-isolation` lets the
build import your installed `torch` so it can match the `libsycl` major to it:
The published wheels are CUDA-only, so install from git and let it build. The ref
below is the `v0.0.10b2` release, pinned so the build is reproducible.
```bash
source /opt/intel/oneapi/setvars.sh
TMS_PLATFORM=xpu pip install --no-build-isolation \
git+https://github.com/fzyzcjy/torch_memory_saver.git@a5c99f11b18ebb8e9fda71a68812e476ae49e417
```
**Use it** by launching with `--enable-memory-saver` (the XPU backend is
selected automatically); optionally add `--enable-weights-cpu-backup` to keep
weights in host RAM across a release:
```bash
python -m sglang.launch_server --model-path Qwen/Qwen3-0.6B \
--trust-remote-code --device xpu --enable-memory-saver
```
```bash
# Release GPU memory while idle, then reclaim it (server must be idle).
curl -X POST http://127.0.0.1:30000/release_memory_occupation
curl -X POST http://127.0.0.1:30000/resume_memory_occupation
```
The Python engine API (`engine.release_memory_occupation(tags=...)` /
`engine.resume_memory_occupation(tags=...)`) and the `weights` / `kv_cache` tags
behave the same as on CUDA. Pauseable CUDA-graph capture is not used on XPU, so
the `cuda_graph` tag is a no-op there.
> **Verifying memory was freed:** neither `torch.xpu.memory_allocated()` nor
> `torch.xpu.mem_get_info()` drops when physical pages are released — the first is
> allocator accounting, and the second stays flat because torch keeps the freed block
> cached. Query actual device memory via sysman (`ZES_ENABLE_SYSMAN=1`) instead.
## Prefill-Decode (P/D) Disaggregation on Intel XPU [Experimental]
SGLang supports prefill-decode disaggregation on Intel XPU using the [NIXL](https://github.com/ai-dynamo/nixl) KV-transfer backend.
python/sglang/srt/utils/torch_memory_saver_adapter.py
+53 -1
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,11 +21,26 @@ 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:
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 "
@@ -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()
test/registered/xpu/test_xpu_memory_saver.py
+350
View File
@@ -0,0 +1,350 @@
"""Unit tests for Intel XPU memory saver (release/resume memory occupation).
These exercise the upstream ``torch_memory_saver`` package's Level Zero
pause/resume path on XPU -- the same backend SGLang's ``release_memory_occupation``
/ ``resume_memory_occupation`` use -- without booting a full inference server.
The tests are skipped unless:
* torch XPU is available, and
* ``torch_memory_saver`` is installed with its XPU backend, built from source
against the local oneAPI (see docs_new/docs/hardware-platforms/xpu.mdx).
Pinned to the v0.0.10b2 release; the published wheels are CUDA-only:
TMS_PLATFORM=xpu pip install --no-build-isolation \\
git+https://github.com/fzyzcjy/torch_memory_saver.git@a5c99f11b18ebb8e9fda71a68812e476ae49e417
Physical-memory release is verified via the saver's driver-independent
``tms_xpu_committed_bytes`` (physical bytes it holds ACTIVE on a device, which
pause() releases and resume() re-commits). Neither sysman free-bytes
(``tms_xpu_device_free_bytes`` / ``torch.xpu.mem_get_info``) nor
``torch.xpu.memory_allocated()`` is used: the former is frozen on newer Intel
drivers and the latter is allocator bookkeeping that ignores ``zeVirtualMemUnmap``.
"""
import unittest
import torch
from sglang.srt.constants import (
GPU_MEMORY_TYPE_CUDA_GRAPH,
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
)
from sglang.test.ci.ci_register import register_xpu_ci
register_xpu_ci(est_time=60, suite="stage-b-test-1-gpu-xpu")
def _xpu_saver_available():
if not (hasattr(torch, "xpu") and torch.xpu.is_available()):
return False
try:
import torch_memory_saver # noqa: F401
except ImportError:
return False
return True
_XPU_OK = _xpu_saver_available()
if _XPU_OK:
import torch_memory_saver as _tms
# The package singleton. On XPU it requires hook_mode="torch"; importing the
# adapter sets that exactly once (it must be set before the singleton is
# initialized). We import the adapter here so the mode is configured before
# any region()/pause() in these tests touches the singleton.
import sglang.srt.utils.torch_memory_saver_adapter # noqa: F401
xpu_memory_saver = _tms.torch_memory_saver
_GIB = 1024**3
# 1 GiB tensor: large enough that a real physical release is unmistakable.
_N_FP32 = 256 * 1024 * 1024
# How much committed memory a 1 GiB region should drop on pause, with alignment slack.
_RELEASE_THRESHOLD_GIB = 0.8
@unittest.skipUnless(
_XPU_OK,
"Requires torch XPU and torch_memory_saver with its XPU backend "
"(TMS_PLATFORM=xpu pip install --no-build-isolation git+https://"
"github.com/fzyzcjy/torch_memory_saver.git@a5c99f11b18ebb8e9fda71a68812e476ae49e417).",
)
class TestXpuMemorySaver(unittest.TestCase):
@classmethod
def setUpClass(cls):
# Pin a single device for the whole class. Probing every device with
# set_device() churns SYCL/L0 contexts and can destabilize the runtime
# when many GPU-heavy tests share one process; pick once, by free memory
# (mem_get_info(i) reads a specific device without changing the current one).
best, best_free = 0, -1
for i in range(torch.xpu.device_count()):
free, _ = torch.xpu.mem_get_info(i)
if free > best_free:
best, best_free = i, free
torch.xpu.set_device(best)
cls.device_index = best
cls.device = f"xpu:{best}"
xpu_memory_saver._ensure_initialized()
cls._cdll = xpu_memory_saver._impl._binary_wrapper.cdll
def tearDown(self):
# Ensure no region is left paused across tests (paused state on the
# process-global singleton would leak into the next test).
try:
xpu_memory_saver.resume(None)
except Exception:
pass
torch.xpu.synchronize()
def _committed_gib(self):
# Physical bytes the saver holds ACTIVE on this device. Driver-independent
# (unlike frozen sysman free-bytes); drops on pause(), restored on resume().
torch.xpu.synchronize()
return self._cdll.tms_xpu_committed_bytes(self.device_index) / _GIB
# ------------------------------------------------------------------ basics
def test_pause_releases_and_resume_restores(self):
"""A region tensor's physical memory is freed on pause and re-committed
on resume at the same virtual address; the tensor stays usable."""
from sglang.srt.utils.torch_memory_saver_adapter import (
TorchMemorySaverAdapter,
)
adapter = TorchMemorySaverAdapter.create(enable=True)
self.assertTrue(adapter.enabled)
with adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
t = torch.ones(_N_FP32, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
self.assertEqual(float(t[0]), 1.0)
committed_after_alloc = self._committed_gib()
adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
committed_after_pause = self._committed_gib()
self.assertGreater(
committed_after_alloc - committed_after_pause,
_RELEASE_THRESHOLD_GIB,
"pause should release ~1 GiB of physical device memory",
)
adapter.resume(GPU_MEMORY_TYPE_KV_CACHE)
committed_after_resume = self._committed_gib()
self.assertGreater(
committed_after_resume - committed_after_pause,
_RELEASE_THRESHOLD_GIB,
"resume should re-commit ~1 GiB of physical device memory",
)
# Same VA must be writable again after resume.
t.fill_(5.0)
torch.xpu.synchronize()
self.assertEqual(float(t[0]), 5.0)
def test_region_scoping_isolates_unmanaged_memory(self):
"""pause(tag) must not touch allocations made outside a region()."""
outside = torch.ones(_N_FP32 // 4, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
with xpu_memory_saver.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
inside = torch.ones(_N_FP32, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
xpu_memory_saver.pause(GPU_MEMORY_TYPE_KV_CACHE)
# The unmanaged tensor must remain valid through a pause of another tag.
self.assertEqual(float(outside[0]), 1.0)
xpu_memory_saver.resume(GPU_MEMORY_TYPE_KV_CACHE)
del inside, outside
def test_cpu_backup_preserves_contents(self):
"""With enable_cpu_backup, contents survive a pause/resume cycle."""
with xpu_memory_saver.region(
tag=GPU_MEMORY_TYPE_WEIGHTS, enable_cpu_backup=True
):
w = torch.full((_N_FP32,), 2.0, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
xpu_memory_saver.pause(GPU_MEMORY_TYPE_WEIGHTS)
xpu_memory_saver.resume(GPU_MEMORY_TYPE_WEIGHTS)
torch.xpu.synchronize()
# CPU backup should restore the exact prior contents.
self.assertEqual(float(w[0]), 2.0)
self.assertEqual(float(w[-1]), 2.0)
del w
def test_tag_selectivity(self):
"""pause(weights) must leave a kv_cache region mapped and usable.
Asserted via tensor readability rather than global free memory: the
saver is a process-global singleton and torch's MemPool caches freed
blocks, so other-tagged allocations from earlier tests may still be live
and would perturb a free-memory delta. Readability of *this* tensor is
the precise selectivity property.
"""
with xpu_memory_saver.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
kv = torch.full((_N_FP32,), 4.0, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
# Pausing a DIFFERENT tag must not unmap the kv_cache region.
xpu_memory_saver.pause(GPU_MEMORY_TYPE_WEIGHTS)
torch.xpu.synchronize()
self.assertEqual(
float(kv[0]),
4.0,
"pausing an unrelated tag must leave the kv_cache region mapped",
)
self.assertEqual(float(kv[-1]), 4.0)
# Now pause the matching tag and confirm it frees real memory.
committed_before_kv = self._committed_gib()
xpu_memory_saver.pause(GPU_MEMORY_TYPE_KV_CACHE)
committed_after_kv = self._committed_gib()
self.assertGreater(
committed_before_kv - committed_after_kv,
_RELEASE_THRESHOLD_GIB,
"pausing the matching tag should release ~1 GiB",
)
# Restore everything we paused so we don't leak paused state.
xpu_memory_saver.resume(GPU_MEMORY_TYPE_KV_CACHE)
xpu_memory_saver.resume(GPU_MEMORY_TYPE_WEIGHTS)
del kv
# ----------------------------------------------------------- adapter wiring
def test_adapter_selected_on_xpu(self):
"""TorchMemorySaverAdapter.create() returns the real (enabled) backend on
XPU, and a disabled adapter is a no-op that reports not-enabled."""
from sglang.srt.utils.torch_memory_saver_adapter import (
TorchMemorySaverAdapter,
_TorchMemorySaverAdapterNoop,
_TorchMemorySaverAdapterReal,
)
adapter = TorchMemorySaverAdapter.create(enable=True)
self.assertIsInstance(adapter, _TorchMemorySaverAdapterReal)
self.assertTrue(adapter.enabled)
noop = TorchMemorySaverAdapter.create(enable=False)
self.assertIsInstance(noop, _TorchMemorySaverAdapterNoop)
self.assertFalse(noop.enabled)
def test_adapter_cuda_graph_and_configure_subprocess_are_noops(self):
"""On XPU, configure_subprocess (no LD_PRELOAD) and cuda_graph (no
pauseable graph capture) must be harmless no-op context managers that
do not raise."""
from sglang.srt.utils.torch_memory_saver_adapter import (
TorchMemorySaverAdapter,
)
adapter = TorchMemorySaverAdapter.create(enable=True)
with adapter.configure_subprocess():
pass
with adapter.cuda_graph(tag=GPU_MEMORY_TYPE_CUDA_GRAPH):
pass
# --------------------------------------------------------------- all tags
def test_release_all_tags_with_none(self):
"""pause(None)/resume(None) act on every region regardless of tag,
matching release_memory_occupation() with no tags (all types)."""
with xpu_memory_saver.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
kv = torch.ones(_N_FP32, dtype=torch.float32, device=self.device)
with xpu_memory_saver.region(tag=GPU_MEMORY_TYPE_WEIGHTS):
wt = torch.ones(_N_FP32, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
committed_alloc = self._committed_gib()
# None == all tags: should free both regions (~2 GiB).
xpu_memory_saver.pause(None)
committed_pause = self._committed_gib()
self.assertGreater(
committed_alloc - committed_pause,
2 * _RELEASE_THRESHOLD_GIB,
"pause(None) should release every region (~2 GiB)",
)
xpu_memory_saver.resume(None)
kv.fill_(1.0)
wt.fill_(1.0)
torch.xpu.synchronize()
self.assertEqual(float(kv[-1]), 1.0)
self.assertEqual(float(wt[-1]), 1.0)
del kv, wt
# ------------------------------------------------------------- robustness
def test_double_pause_and_double_resume_are_idempotent(self):
"""Pausing an already-paused tag (or resuming an active one) must be a
safe no-op, not a crash or double-free."""
with xpu_memory_saver.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
t = torch.full((_N_FP32,), 3.0, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
# resume while already active -> no-op
xpu_memory_saver.resume(GPU_MEMORY_TYPE_KV_CACHE)
committed_active = self._committed_gib()
xpu_memory_saver.pause(GPU_MEMORY_TYPE_KV_CACHE)
# pause again while already paused -> no-op (no extra free, no crash)
xpu_memory_saver.pause(GPU_MEMORY_TYPE_KV_CACHE)
committed_paused = self._committed_gib()
self.assertGreater(committed_active - committed_paused, _RELEASE_THRESHOLD_GIB)
xpu_memory_saver.resume(GPU_MEMORY_TYPE_KV_CACHE)
xpu_memory_saver.resume(GPU_MEMORY_TYPE_KV_CACHE) # second resume -> no-op
t.fill_(9.0)
torch.xpu.synchronize()
self.assertEqual(float(t[0]), 9.0)
del t
def test_multiple_regions_same_tag_all_released(self):
"""Several allocations under one tag are all released together and all
usable again after resume."""
tensors = []
with xpu_memory_saver.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
for v in (1.0, 2.0, 3.0):
tensors.append(
torch.full(
(_N_FP32 // 2,), v, dtype=torch.float32, device=self.device
)
)
torch.xpu.synchronize()
committed_alloc = self._committed_gib()
xpu_memory_saver.pause(GPU_MEMORY_TYPE_KV_CACHE)
committed_pause = self._committed_gib()
# 3 x 0.5 GiB = ~1.5 GiB.
self.assertGreater(committed_alloc - committed_pause, _RELEASE_THRESHOLD_GIB)
xpu_memory_saver.resume(GPU_MEMORY_TYPE_KV_CACHE)
for v, t in zip((1.0, 2.0, 3.0), tensors):
# contents are not preserved without cpu_backup; just confirm the VA
# is mapped and writable again.
t.fill_(v)
torch.xpu.synchronize()
for v, t in zip((1.0, 2.0, 3.0), tensors):
self.assertEqual(float(t[0]), v)
del tensors
def test_resume_without_backup_gives_writable_fresh_memory(self):
"""Without cpu_backup, resumed memory need not preserve contents (KV
cache is flushed anyway) but must be freshly writable at the same VA."""
with xpu_memory_saver.region(tag=GPU_MEMORY_TYPE_KV_CACHE):
t = torch.full((_N_FP32,), 7.0, dtype=torch.float32, device=self.device)
torch.xpu.synchronize()
ptr_before = t.data_ptr()
xpu_memory_saver.pause(GPU_MEMORY_TYPE_KV_CACHE)
xpu_memory_saver.resume(GPU_MEMORY_TYPE_KV_CACHE)
# Same virtual address preserved across the cycle.
self.assertEqual(t.data_ptr(), ptr_before)
# Writable again.
t.fill_(42.0)
torch.xpu.synchronize()
self.assertEqual(float(t[0]), 42.0)
self.assertEqual(float(t[-1]), 42.0)
del t
if __name__ == "__main__":
unittest.main()