Skip to content

Add Intel XPU (Level Zero) backend - #78

Open
siju-samuel wants to merge 8 commits into
fzyzcjy:masterfrom
siju-samuel:feat/xpu-support
Open

Add Intel XPU (Level Zero) backend#78
siju-samuel wants to merge 8 commits into
fzyzcjy:masterfrom
siju-samuel:feat/xpu-support

Conversation

@siju-samuel

@siju-samuel siju-samuel commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

Adds Intel GPU support alongside the existing CUDA and ROCm backends, providing the same pause/resume (release/re-commit) behavior built natively on Level Zero virtual memory:

malloc -> zeVirtualMemReserve + zePhysicalMemCreate + zeVirtualMemMap
pause  -> zeVirtualMemUnmap + zePhysicalMemDestroy   (VA stays reserved)
resume -> zePhysicalMemCreate + zeVirtualMemMap      (same VA)

Wired into PyTorch via torch.xpu.memory.XPUPluggableAllocator + torch.xpu.MemPool, so only hook_mode="torch" is supported (LD_PRELOAD preload is CUDA/HIP-specific). Pauseable graph capture, disable(), disk backup and memory_margin_bytes are not supported on XPU and are rejected explicitly rather than silently ignored.

C++

  • macro.h / utils.h: USE_XPU platform branch (types, device helpers, host-copy shims so the shared disk-backup path compiles unforked).
  • core.h: AllocationMetadata XPU fields (XPUAllocExtra).
  • hardware_xpu_support.{h,cpp}: Level Zero backend — malloc/free/pause/resume, optional CPU backup, per-device SYCL context prewarm, and byte-accounting (committed/leaked/tracked/affected_devices). Uses the SYCL platform default context so VA mappings are visible to torch kernels.
  • Transactional error handling: every pause/resume/free step either completes or leaves a recoverable state — a failed zePhysicalMemDestroy retains the handle (surfaced as leaked bytes) so resume re-maps that exact handle; resume restores the CPU backup before committing; free keeps its ownership record until Level Zero confirms release. tms_pause/tms_resume return non-zero on partial failure and Python raises.
  • core.cpp / entrypoint.cpp: dispatch to XPUImplementation under USE_XPU. CUDA/ROCm paths unchanged (XPU is an #elif before the CUDA #else).

Python

  • hooks/mode_torch.py: use XPUPluggableAllocator on XPU.
  • entrypoint.py: torch.xpu MemPool/use_mem_pool, prewarm, atexit pool cleanup. MemPool key includes the current device so multiple devices in one process are each captured. pause/resume drain every device the backend will unmap (queried from tms_xpu_affected_devices), not just the current one — a missed device risks a DEVICE_LOST hang.
  • utils.py / binary_wrapper.py: XPU binary resolution + tms_xpu_* signatures.

Tests

  • test/test_examples.py is device-aware: on XPU it runs only hook_mode='torch' and skips CUDA/HIP-only examples (cuda_graph, disk_backup, multi_device, rl_example, training_engine, nested_region) with clear reasons.
  • test/examples/multi_device_torch_mode.py: torch-mode multi-GPU pause/resume (pins each device at alloc time), runs on both CUDA and XPU. Preload-mode multi-device coverage stays in the existing multi_device.py.
  • test/examples/xpu_scenarios.py: XPU-only scenarios, each in its own spawned subprocess — fault injection for the resume/cleanup/free transactional paths (asserted via the committed/leaked/tracked counters), multi-device drain, and the disable()/memory_margin_bytes rejections.
  • On Intel XPU: 9 passed, 7 skipped.

Build & docs

  • setup.py: detect XPU via icpx; build with -fsycl, link sycl + ze_loader; build only the torch hook variant. Picks an icpx whose libsycl major matches the installed torch+xpu so the .so links a compatible ABI.
  • make build-xpu: TMS_PLATFORM=xpu pip install --no-build-isolation .
  • Examples made device-portable via testing_utils.get_device()/empty_cache().
  • README: Platform Support table + Intel XPU section.

Validated on Intel GPUs: pause/resume release real physical memory and restore usable tensors at the same virtual address. Verified via the driver-independent tms_xpu_committed_bytes (sysman free-bytes is frozen on newer Intel drivers), with oneAPI 2025.3 and 2026.0.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for the Intel XPU backend using Level Zero virtual memory management. It introduces build scripts, updates setup configurations for the icpx compiler, and integrates XPU-specific allocation, pause, and resume hooks. The review feedback highlights three key issues: a critical bug in hardware_xpu_support.cpp where a failed CPU backup still proceeds to destroy physical memory (risking data loss), an incorrect directory traversal depth in setup.py when resolving the oneAPI root path, and missing return-value checks on secondary Sysman API calls in xpu_device_free_bytes.

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.

Comment thread csrc/hardware_xpu_support.cpp
Comment thread setup.py Outdated
Comment thread csrc/hardware_xpu_support.cpp Outdated
siju-samuel added a commit to siju-samuel/torch_memory_saver that referenced this pull request Jun 25, 2026
- Fix path resolution in setup.py: go up 4 levels to reach oneAPI root
  (was resolving to /opt/intel/oneapi/compiler/compiler/latest/include)
- Add CPU backup failure handling in xpu_pause: skip pause if malloc
  or memcpy fails to prevent silent data loss
- Add missing return value checks in xpu_device_free_bytes for all
  sysman API calls (zesDriverGet, zesDeviceGet, zesDeviceEnumMemoryModules)
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jun 25, 2026
Replace in-repo XPU memory saver with upstream torch_memory_saver package.
Removes ~961 lines of in-repo build infrastructure.

Changes:
- Deleted python/sglang/srt/utils/xpu_memory_saver/ (entire package)
- Deleted _TorchMemorySaverAdapterXpu class
- Updated adapter to use torch_memory_saver for all platforms (CUDA/ROCm/XPU)
- Set hook_mode='torch' for XPU at module import
- configure_subprocess() and cuda_graph() are no-ops on XPU
- Fixed test_xpu_memory_saver.py stability (device probe, tearDown cleanup)
- Updated Dockerfile to pip install upstream package
- Updated docs with installation instructions

Depends on: fzyzcjy/torch_memory_saver#78

Test: All 10 XPU memory saver tests pass (3.028s)
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jun 25, 2026
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 (tested on 2+ XPUs)
- 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 pip install upstream package
- Add comprehensive test suite (10 tests, all passing)

Installation:
  pip install 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
Test: 10/10 tests pass in 3.028s
@fzyzcjy

fzyzcjy commented Jun 25, 2026

Copy link
Copy Markdown
Owner

Looks great! Feel free to ping me when it is ready; if I do not see the comment, also feel free to ping me in SGLang slack

@siju-samuel
siju-samuel marked this pull request as ready for review June 26, 2026 09:14
@siju-samuel

Copy link
Copy Markdown
Contributor Author

@fzyzcjy please review and merge. Thanks in advance
Verified with Intel OneAPI version 2025.3 and 2026.0

@siju-samuel siju-samuel changed the title [DRAFT] Add Intel XPU (Level Zero) backend Add Intel XPU (Level Zero) backend Jun 26, 2026
@siju-samuel

Copy link
Copy Markdown
Contributor Author

@fzyzcjy Somehow i was not able to find you in SGLang slack.
Could you please have a look at this PR when you get a chance?

siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jul 2, 2026
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 (tested on 2+ XPUs)
- 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 pip install upstream package
- Add comprehensive test suite (10 tests, all passing)

Installation:
  pip install 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
Test: 10/10 tests pass in 3.028s
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jul 2, 2026
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 (tested on 2+ XPUs)
- 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 pip install upstream package
- Add comprehensive test suite (10 tests, all passing)

Installation:
  pip install 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
Test: 10/10 tests pass in 3.028s
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jul 2, 2026
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
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jul 2, 2026
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
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jul 2, 2026
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
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jul 2, 2026
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
siju-samuel added a commit to siju-samuel/sglang that referenced this pull request Jul 2, 2026
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
@fzyzcjy

fzyzcjy commented Jul 7, 2026

Copy link
Copy Markdown
Owner

Let me review now. SGLang slack is here:

image

@siju-samuel

Copy link
Copy Markdown
Contributor Author

Thanks @fzyzcjy
image
Somehow SGL wont show up your name.
Hope you got this message in slack
image

Comment thread csrc/entrypoint.cpp Outdated
Comment thread csrc/entrypoint.cpp
Comment thread test/examples/xpu_multi_device.py Outdated
Comment thread torch_memory_saver/hooks/mode_torch.py Outdated
Comment thread torch_memory_saver/entrypoint.py Outdated
Comment thread torch_memory_saver/entrypoint.py Outdated
Comment thread torch_memory_saver/entrypoint.py Outdated
Comment thread README.md
Comment thread README.md Outdated
Comment thread setup.py Outdated
@fzyzcjy

fzyzcjy commented Jul 7, 2026

Copy link
Copy Markdown
Owner

@siju-samuel just @Tom in sgl slack; I do not recv @ for that message indeed

@siju-samuel
siju-samuel force-pushed the feat/xpu-support branch 2 times, most recently from 9b1f4d2 to 43df4d4 Compare July 8, 2026 15:06

@fzyzcjy fzyzcjy left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

generally LGTM, will have a deeper look (and w/ agents) after prereq pr is merged

@fzyzcjy

fzyzcjy commented Jul 17, 2026

Copy link
Copy Markdown
Owner

hi, could you plz resolve merge conflicts (and I will have a final review ASAP having time)

@siju-samuel

Copy link
Copy Markdown
Contributor Author

hi, could you plz resolve merge conflicts (and I will have a final review ASAP having time)

Thanks. I have rebased and resolved the merge conflicts.

@fzyzcjy

fzyzcjy commented Jul 22, 2026

Copy link
Copy Markdown
Owner

looks great, busy recently and will review ASAP when having time

@fzyzcjy fzyzcjy left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 Agent-authored inline review. This review was generated and submitted autonomously by a coding agent on behalf of the repository owner. The findings below were independently checked across semantic, lifecycle, architecture, performance, and adversarial passes. Please point out anything inaccurate.

Requesting changes for the blocking findings attached to the exact diff lines. This was a source-level review; no Intel XPU hardware was available locally.

Comment thread torch_memory_saver/entrypoint.py Outdated
Comment thread torch_memory_saver/entrypoint.py Outdated
Comment thread csrc/hardware_xpu_support.cpp Outdated
pdesc.stype = ZE_STRUCTURE_TYPE_PHYSICAL_MEM_DESC;
pdesc.size = aligned;
ze_physical_mem_handle_t phys = {};
ze_result_t rc = zePhysicalMemCreate(ze_ctx, ze_dev, &pdesc, &phys);

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] A create/map failure is silently treated as a successful resume(). Both failure branches only log and continue, while the C entrypoint and Python API return no status. A tag can therefore be half-resumed, with some tensors still unmapped; the next access can device-fault even though resume() returned normally. Please make the transition transactional, roll back partial work, propagate an error to Python, and add failure-injection coverage for create/map failures.

Comment thread csrc/hardware_xpu_support.cpp
Comment thread csrc/hardware_xpu_support.cpp Outdated
if (it == allocation_metadata.end())
return cudaErrorInvalidDevicePointer;
metadata = it->second;
allocation_metadata.erase(it);

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] xpu_free() drops the only ownership record before any Level Zero cleanup succeeds. The subsequent unmap/destroy/free return values are all ignored, so any failure permanently leaks physical memory or VA and cannot be retried; an unmap failure is followed by destroying a still-mapped physical handle, despite xpu_pause() explicitly avoiding that unsafe sequence. Please retain explicit ownership/state until every cleanup step succeeds and propagate cleanup failures to the allocator.

Comment thread csrc/hardware_xpu_support.cpp
Comment thread csrc/core.cpp
Comment thread setup.py Outdated

@fzyzcjy fzyzcjy left a comment

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM w/ nits (below are human reviews; above are agent reviews as serve as references not mandatory)

Comment thread csrc/entrypoint.cpp Outdated
Comment thread test/examples/multi_device_torch_mode.py
Comment thread torch_memory_saver/entrypoint.py Outdated
Comment thread torch_memory_saver/entrypoint.py Outdated
@siju-samuel
siju-samuel force-pushed the feat/xpu-support branch 6 times, most recently from 812b64e to ab5a5ce Compare July 26, 2026 14:17
@fzyzcjy

fzyzcjy commented Jul 26, 2026

Copy link
Copy Markdown
Owner

(feel free to ping me when ready)

@siju-samuel

Copy link
Copy Markdown
Contributor Author

(feel free to ping me when ready)

There were merge conflicts due to #80. I'm validating the changes to make sure everything still works correctly after resolving them.

@siju-samuel

Copy link
Copy Markdown
Contributor Author
image

@fzyzcjy I've completed the changes from my side and re-verified everything after resolving the merge conflicts. The PR is ready for review again. When you have a chance, could you please take another look and help merge it? Thanks!

Comment thread csrc/entrypoint.cpp Outdated
}

#if defined(USE_XPU)
// Pre-warm per-device SYCL contexts from Python init: creating a sycl::context

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: only have absolutely necessary comments, i.e. remove most of them in the pr

Comment thread torch_memory_saver/entrypoint.py Outdated
# work before user kernels touch the remapped addresses.
self._sync_affected_devices(tag)

def _sync_affected_devices(self, tag: Optional[str]):

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm why add some functions that are not guarded under xpu

Add Intel XPU support to torch_memory_saver using Level Zero Virtual Memory
Management (VMM) API. Enables pause/resume of GPU memory occupation on Intel
discrete GPUs (Arc Pro B60, PVC with Xe Link) for the same use cases as the
existing CUDA/ROCm backends.

Architecture:
- VMM: zeVirtualMemReserve, zePhysicalMemCreate, zeVirtualMemMap/Unmap
- Hook mode: torch.xpu.memory.XPUPluggableAllocator (in-process, not LD_PRELOAD)
- Platform context: Uses SYCL platform-default context (same as torch streams)
- Memory verification: Sysman (zesMemoryGetState) for accurate free-memory reading

Features:
- Pause/resume physical memory while preserving virtual addresses
- Optional CPU backup for content preservation across pause/resume
- Multi-device support (per-device SYCL contexts)
- Passthrough allocation for non-region memory
- Build-from-source (links libsycl.so.<N>, must match torch+xpu runtime)

Changes:
- Add csrc/hardware_xpu_support.{h,cpp} - Level Zero backend implementation
- Add XPU branches to macro.h, utils.h, core.{h,cpp}, entrypoint.cpp
- Update setup.py to auto-detect XPU from icpx on PATH
- Update Python layer (hooks/mode_torch.py, entrypoint.py, testing_utils.py)
- Add scripts/build_xpu.sh with SONAME verification
- Add test/examples/xpu_multi_device.py for multi-GPU validation
- Update README with XPU installation instructions

Installation:
  # Requires Intel oneAPI (icpx) on PATH
  pip install git+https://github.com/fzyzcjy/torch_memory_saver.git

Tests: All XPU examples pass (simple, cpu_backup, multi_device)
…e filter

- setup.py: select an icpx whose libsycl major matches the installed
  torch+xpu and derive include/lib from that compiler, so the built .so
  links the same libsycl the runtime loads (avoids SYCL-runtime corruption
  on a mismatched oneAPI).
- hardware_xpu_support: filter the SYCL device list to the Level-Zero
  backend before indexing by a torch device ordinal (fixes 'Backends
  mismatch' when GPUs are exposed via both L0 and OpenCL).
- Add tms_xpu_committed_bytes: driver-independent physical-commit query
  for pause/resume tests (sysman free-bytes is frozen on newer drivers).
- build_xpu.sh: ABI-aware build helper; scan both .so locations.
- README/docs: install with --no-build-isolation; clarify libsycl matching.
- Remove the never-reached non-region passthrough branches in tms_torch_malloc
  / tms_torch_free, plus the now-orphaned xpu_passthrough_malloc/free,
  xpu_is_managed, and TorchMemorySaver::is_managed. The torch MemPool allocator
  is pool-scoped, so the hook only fires for in-region (managed) allocations.
- Remove scripts/build_xpu.sh; inline its command into 'make build-xpu'
  (TMS_PLATFORM=xpu pip install --no-build-isolation .).
- README: document TMS_PLATFORM=xpu + --no-build-isolation for the install.
- pause()/resume(): sync the current device before unmap/after remap (required:
  unmapping Level-Zero pages with kernels in flight triggers DEVICE_LOST).
  Cross-device sync is the caller's responsibility, matching the CUDA backend;
  the one-process-per-rank (one device) deployment satisfies this implicitly.
- xpu_pause: on zeVirtualMemUnmap failure, keep the allocation ACTIVE instead of
  destroying the physical handle and marking PAUSED (which would orphan the VA
  and make the next resume fail). Lets it retry/free cleanly.
- entrypoint.cpp: clarify why the in-region assert is dropped on XPU only.
- setup.py: pass -std=c++17 explicitly (backend uses structured bindings).
- testing_utils: note _xpu_used_memory returns saver-committed vs whole-device
  bytes on different paths (fine for delta asserts, not an absolute figure).
- README: generalize the libsycl-major example (oneAPI 2025.x -> libsycl.so.8).
- ze_alloc_granularity now passes the requested size to
  zeVirtualMemQueryPageSize (was hard-coded 1), so the driver returns the page
  size recommended for that allocation instead of the smallest-allocation
  granularity. Strictly more correct for large allocations; 2 MiB fallback
  unchanged.
- setup.py: note _find_platform_home's xpu branch is kept only for symmetry
  (the XPU build derives include/lib from the resolved icpx, not this value).
- Deduplicate _is_xpu into utils._is_xpu_torch (was defined 4x).
- Collapse per-platform MemPool/use_mem_pool/current_device branching via
  torch.get_device_module().
- Densify comments (prewarm, atexit, pool-key) per review.
- setup.py: group the XPU helper functions under a '# For Intel XPU' banner at
  the end of the file; trim their comments.
- README: drop bold, remove the redundant usage example and 'Why no XPU wheel'
  section, keep the essentials.
- Make the multi-device test cross-platform (multi_device_torch_mode.py: pins
  each device at alloc time, runs on CUDA and XPU) replacing the XPU-only one.
Rework of the Intel XPU (Level Zero VMM) memory-saver backend in response to
upstream PR review feedback.

- xpu_free: make it transactional. Retain the ownership record and check
  every unmap/destroy/free-VA return code, advancing state per step and
  erasing the entry only after all confirm; a failed step no longer leaks
  physical memory/VA silently and the error propagates to the allocator.
- xpu_pause: on zePhysicalMemDestroy failure keep the physical handle
  (marked leaked) instead of clearing it and reporting the tag paused, so the
  next resume re-maps the same handle. Surface it via xpu_leaked_bytes
  (independent of the ACTIVE-only committed counter) and xpu_tracked_bytes
  (any state, incl. the free-VA-failure state where committed and leaked
  both read 0).
- memory_margin_bytes: reject on XPU at both entry points. The OOM guard needs
  a device-wide free-bytes reading the Intel driver reports frozen; the Python
  setter raises NotImplementedError and the C-ABI set_memory_margin_bytes
  refuses to store a non-zero margin (warns) rather than silently ignoring it.
- setup.py: append the SYCL flags to the extension's existing compile/link
  args instead of overwriting them, so the -O3/-std=c++17 seeded in
  _create_ext_modules() survives (icpx honors the last -O flag).
- get_binary_path_from_package: prefer the package dir and fall back to its
  parent only when empty, instead of globbing both; a stray inplace-build .so
  in the repo root no longer trips "expected exactly one .so ... found 2".
- test_examples: import example modules by absolute path (sys.path) rather
  than `from examples import ...`, avoiding a name collision with Intel's
  pyzes `examples` package (no test/examples/__init__.py needed).
- Add failure-injection fault hooks (TMS_XPU_FAULT_*) and a test-only
  tms_xpu_free entrypoint returning the error code instead of aborting, so
  the transactional paths are validated by injected failures, not only
  self-reported counters.
- Tests: cover free/pause/resume failure injection, margin rejection (incl.
  fd-level capture of the C-ABI warning), multi-device sync, and disable()
  rejection. The six XPU-only scenarios live in one module
  (test/examples/xpu_scenarios.py), each still run in its own spawned
  subprocess so a fault-injection abort can't affect the others.
- Keep the comments added by this rework dense and short (2-3 lines each),
  preserving the load-bearing invariant each documents.
@siju-samuel
siju-samuel force-pushed the feat/xpu-support branch 2 times, most recently from b6c0035 to 5299faa Compare July 30, 2026 03:51
Comments: reduce volume across the XPU backend per review feedback, keeping
only load-bearing rationale (driver quirks, transactional invariants, why a
path is rejected) and dropping restatements of the code, step narration and
section banners.

Scope the multi-device drain helpers to XPU: rename _sync_affected_devices /
_affected_devices to _xpu_* (they are only reachable under _is_xpu) and drop
the _region_devices mirror, which every platform populated but only XPU read.
Its sole reader was a hasattr() fallback for a symbol that always exists in an
XPU build, so both go away together.

Drop the #if !defined(USE_XPU) guards from disk_backend.cpp so the disk-backup
path is no longer platform-forked: macro.h now supplies the XPU host-copy
shims (cudaMallocHost/cudaMemcpy) the same way the ROCm branch maps them to
hip*, returning an error since there is no Level Zero pinned-host transfer.
The feature stays rejected up front (Python raises; TorchMemorySaver::malloc
asserts), and xpu_malloc now initializes enable_disk_backup so the gate that
keeps DiskBackend unreachable reads an initialized value.

Also adapt the XPU path to the AllocationMetadata::size -> raw_size and
malloc(size -> raw_size) renames from the CUDA VMM granularity change; XPU
still aligns internally via zeVirtualMemQueryPageSize, so it takes raw_size.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants