[Feature][Intel XPU] Add memory saver support for Intel XPU via upstream torch_memory_saver - #29935
[Feature][Intel XPU] Add memory saver support for Intel XPU via upstream torch_memory_saver#29935siju-samuel wants to merge 3 commits into
Conversation
Enable release/resume_memory_occupation on Intel XPU using the upstream
torch_memory_saver package (Level Zero VMM backend). Replaces in-repo
implementation with pip-installable upstream dependency.
Features:
- Pause/resume physical GPU memory while preserving virtual addresses
- Tag-based selective memory management (kv_cache, weights, cuda_graph)
- Optional CPU backup for content preservation
- Multi-device support (verified via the package's multi-device test)
- Sysman-based memory verification (torch accounting doesn't reflect unmapped pages)
Changes:
- Add TorchMemorySaverAdapter for unified memory saver API
- XPU uses hook_mode='torch' (in-process pluggable allocator, not LD_PRELOAD)
- configure_subprocess() and cuda_graph() are no-ops on XPU
- Update Dockerfile to build the upstream package for XPU
- Add test suite (10 tests, all passing)
Installation (XPU builds from source; TMS_PLATFORM=xpu forces the XPU backend
and --no-build-isolation lets the build match libsycl to the installed torch):
source /opt/intel/oneapi/setvars.sh
TMS_PLATFORM=xpu pip install --no-build-isolation \
git+https://github.com/fzyzcjy/torch_memory_saver.git
Usage:
python -m sglang.launch_server --model-path <model> --device xpu --enable-memory-saver
Depends on: fzyzcjy/torch_memory_saver#78
github-actions
Bot
added
the
documentation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for the experimental Memory Saver feature on Intel XPU, allowing the engine to temporarily release and reclaim GPU memory. The changes include updating the Dockerfile to install torch_memory_saver from source, adding documentation, adapting the TorchMemorySaverAdapter for XPU specifics (such as using the in-process pluggable allocator), updating memory utility functions, and introducing a comprehensive suite of unit tests. The reviewer feedback suggests pinning the torch_memory_saver repository to a specific commit hash in the Dockerfile to ensure reproducible builds, and removing an unused variable cur in the test setup.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| 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 |
There was a problem hiding this comment.
To ensure reproducible builds and prevent future breaking changes on the main branch of torch_memory_saver from breaking the Docker image build, it is highly recommended to pin the installation to a specific commit hash or tag.
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@YOUR_COMMIT_HASH
| best, best_free = 0, -1 | ||
| cur = torch.xpu.current_device() | ||
| for i in range(torch.xpu.device_count()): |
There was a problem hiding this comment.
The variable cur is assigned the current XPU device but is never used in the rest of the method. It can be safely removed to keep the code clean.
best, best_free = 0, -1\n for i in range(torch.xpu.device_count()):Addresses review on PR sgl-project#29935: torch.xpu.current_device() was assigned to 'cur' but never used (mem_get_info(i) reads a specific device without changing the current one). Drop it and correct the comment.
| # Use mem_get_info (real device free/total), NOT memory_allocated(): the | ||
| # latter is the allocator's bookkeeping and does not drop when the XPU | ||
| # memory saver releases physical pages via zeVirtualMemUnmap. | ||
| free, total = torch.xpu.mem_get_info() |
There was a problem hiding this comment.
is this api working ?
| # (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) |
There was a problem hiding this comment.
is this api working ?
| return _memory_saver.region(tag=tag, enable_cpu_backup=enable_cpu_backup) | ||
|
|
||
| def cuda_graph(self, **kwargs): | ||
| if is_xpu(): |
There was a problem hiding this comment.
why can't we support this ?
2 participants
Motivation
Enable release/resume_memory_occupation on Intel XPU using the upstream torch_memory_saver package (Level Zero VMM backend). Replaces in-repo implementation with pip-installable upstream dependency.
Features:
Modifications
Changes:
Installation (XPU builds from source; TMS_PLATFORM=xpu forces the XPU backend and --no-build-isolation lets the build match libsycl to the installed torch):
source /opt/intel/oneapi/setvars.sh
TMS_PLATFORM=xpu pip install --no-build-isolation
git+https://github.com/fzyzcjy/torch_memory_saver.git
Usage:
python -m sglang.launch_server --model-path --device xpu --enable-memory-saver
Depends on: fzyzcjy/torch_memory_saver#78
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-cicc @mingfeima @fzyzcjy @habaohaba @yueming-yuan @jianan-gu @hnyls2002 @merrymercy
CI States
Latest PR Test (Base): ❌ Run #28586566535
Latest PR Test (Extra): ❌ Run #28586566353