[XPU] Support checkpoint-engine weight update on Intel XPU - #96
Open
siju-samuel wants to merge 2 commits into
Open
[XPU] Support checkpoint-engine weight update on Intel XPU#96siju-samuel wants to merge 2 commits into
siju-samuel wants to merge 2 commits into
Conversation
5 tasks
siju-samuel
force-pushed
the
feat/xpu-support-phase1
branch
from
July 26, 2026 10:30
5c0d3f4 to
324468d
Compare
Add Intel XPU as a first-class device backend alongside CUDA and NPU. Device abstraction (device_utils.py): detect xpu, map to the xccl distributed backend, and gate capabilities per backend -- in-place host pinning (cudaHostRegister) and Mooncake device P2P stay CUDA/NPU-only, while cross-process device-tensor IPC is supported on XPU via a native extension. Weight transport (transport.py): the broadcast path shares a device buffer with the colocated inference worker. PyTorch has no XPU tensor IPC (reduce_tensor raises '_share_fd_: only available on CPU'), so a WeightTransport seam selects the CUDA/NPU reduce_tensor path or a native XPU implementation. XPU IPC (xpu_ipc/): implemented on SYCL ext::oneapi::experimental::ipc_memory, which torch's own libsycl (oneAPI >= 2026.0, libsycl.so.9) exports. A SYCL IPC handle is a self-contained, portable byte blob -- verified cross-process on Arc B60 -- so it needs no out-of-band dma-buf fd transfer and no sub-allocation offset handling: the whole handle rides the existing ZMQ channel exactly like CUDA's reduce_tensor tuple. The -fsycl extension coexists with torch's single SYCL runtime in-process. Only the broadcast update method is supported on XPU; P2P is rejected with a clear error (Mooncake has no Level Zero backend for XPU device memory). Robustness: the exported IPC handle is released in an outer try/finally so an early failure (export, ZMQ bind, or first send) can no longer leak it, while the collective barrier stays in the inner finally to avoid deadlocking peers on an asymmetric failure. The P2P store is skipped entirely on backends that do not support device P2P (e.g. XPU) instead of being eagerly initialized and only guarded against ImportError. icpx discovery falls back to PATH (via shutil.which) for oneAPI layouts outside /opt or a sourced setvars.sh that does not export CMPLR_ROOT. Tests: CPU-only coverage for device dispatch, the transport seam, the P2P guard, and the XPU parity paths (portable-handle contract, xccl backend selection, in-place-pin disable, detach-on-early-failure, icpx PATH fallback); hardware-gated tests for the SYCL IPC roundtrip, interior-pointer offset, and the full cross-process broadcast.
Author
|
@kip-cxj @specture724 @weixiao-huang @ppwwyyxx |
| ``reduce_tensor`` tuple -- see ``XpuIpcWeightTransport``. | ||
| """ | ||
|
|
||
| from __future__ import annotations |
Collaborator
There was a problem hiding this comment.
It is unnecessary to import this line. I thinks we can use double quote in type
| def ipc_collect(self) -> None: | ||
| """Reclaim memory held by stale IPC handles where the backend supports it (no-op otherwise).""" | ||
| fn = getattr(self.device_module, "ipc_collect", None) | ||
| if callable(fn): |
Collaborator
There was a problem hiding this comment.
If not callable, mayeb it should raise
| socket.recv() | ||
| gidx = 0 | ||
| ret_code = torch.zeros((), device=self.device_manager.device_type, dtype=torch.int64) | ||
| try: |
Collaborator
There was a problem hiding this comment.
Why there need additinal try and finally?
HubertZhang
reviewed
Jul 28, 2026
| @@ -0,0 +1,124 @@ | |||
| """Pluggable device-buffer handoff between the ParameterServer and the worker. | |||
Collaborator
There was a problem hiding this comment.
What about renaming WeightTransport to IPCHandler? This object transports nothing; it just hands an IPC handle for a device buffer across processes.
| #include <unordered_map> | ||
| #include <vector> | ||
|
|
||
| namespace ipc = sycl::ext::oneapi::experimental::ipc_memory; |
Collaborator
There was a problem hiding this comment.
It seems ipc_memory namespace is deprecated upstream and should migrate to ipc::memory?
Rework of the six review comments from @weixiao-huang and @HubertZhang. Drop `from __future__ import annotations` from the two new modules and quote the TYPE_CHECKING-only annotations instead; the rest of the repo already relies on native PEP 604 unions (requires-python >= 3.10). DeviceManager.ipc_collect no longer duck-types via getattr/callable: it dispatches on device_type like every other method in the class and raises TypeError on an unsupported backend. XPU is an explicit no-op (SYCL frees on close_handle; torch.xpu has no ipc_collect). Replace the extra try/finally in _update_per_bucket with a context manager: IpcHandler.__enter__/__exit__ release the exported handle, so the call site is a plain `with` and the early-failure leak stays fixed without added nesting. The collective barrier remains in the inner finally so an asymmetric early failure cannot deadlock peers. Rename the abstraction to match what it does -- it hands over an IPC handle rather than transporting bytes: WeightTransport -> IpcHandler, IpcWeightTransport -> TorchIpcHandler, XpuIpcWeightTransport -> XpuIpcHandler, build_transport -> build_ipc_handler, transport.py -> ipc_handler.py (and tests/test_transport.py -> tests/test_ipc_handler.py). Acronym casing follows the existing HcclCommConfig / DistributedNccl / VllmColocateWorkerExtension convention. The "xpu_sycl" wire tag is unchanged, so producer/consumer stay compatible across builds. Guard the SYCL namespace with __has_include: upstream intel/llvm split the API (functions -> ipc::memory, handle types -> the parent ipc namespace) and deprecated flat ipc_memory, but no oneAPI release ships that layout yet and the extension's feature-test macro is unversioned, so both spellings are supported. Verified: the shim compiles clean under -Werror=deprecated-declarations on oneAPI 2026.1. Build the extension with torch's with_sycl=True, which supplies the SYCL include paths and device link. This removes _find_sycl_include_dir entirely and the CC/CXX override block -- confirmed unnecessary, a clean build now succeeds with CXX=g++ exported. _find_icpx stays because torch shells out to a bare "icpx", so the located compiler is prepended to PATH for the build; -O2 is kept since without it the host object is compiled -O0 (6x larger .so).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This PR adds Intel XPU as a first-class device backend behind device abstraction, and fills the IPC gap with a native SYCL ipc_memory transport, which is the XPU counterpart of CUDA IPC
The CUDA and NPU paths are unchanged, XPU-specific code is reached only when the device type is xpu.
This unblocks checkpoint-engine's fast weight-update path on Intel GPUs, which is needed for SGLang's XPU for RL rollout weight sync.
Modifications
Device abstraction
device_utils.py: detectxpuand map it to thexcclcollective backend; addsupports_device_ipc()/supports_device_p2p()capability gates so XPU enables broadcast IPC but not P2P.distributed/base.py: reject custom distributed backends on XPU and fall back to the defaultxcclTorchBackend.**Weight transport **
transport.py: add aWeightTransportseam —IpcWeightTransport(CUDA/NPU, unchanged) andXpuIpcWeightTransport(XPU) — withbuild_transport()selecting by device type.xpu_ipc/sycl_ipc.cpp+xpu_ipc/__init__.py: native SYCLipc_memoryextension (get/open/close), JIT-compiled at runtime with-fsycl, providing the cross-process device-buffer IPC that PyTorch lacks on XPU.Parameter server & worker
ps.py: route the broadcast throughbuild_transport(), prebuild the SYCL extension at startup, and fail loudly with clear errors when device IPC/P2P is unavailable on XPU.worker.py: derive theGPU-<uuid>handshake key fromtorch.xpuon XPU, matching the ParameterServer's_get_physical_gpu_id.Packaging
pyproject.toml: ship the SYCL.cppsource as package data (JIT-compiled on XPU hosts; inert on CUDA).Tests
test_device_manager.py,test_transport.py,test_xpu_parity.py,test_p2p_guard.py— device dispatch, transport seam, portable-handle contract,xcclselection, P2P rejection.test_xpu_ipc.py— SYCL IPC roundtrip, interior-offset, full cross-process broadcast.Accuracy Test / Verified
Device-UUID handshake verified on Intel Arc Pro B60: the worker emits the exact
GPU-<uuid>key the ParameterServer expects. 5 CPU unit tests pass.Environment
Check list
pytest tests/ -m "not gpu")pytest tests/test_xpu_ipc.py)xpu)CC: @@kip-cxj @specture724 @weixiao-huang @ppwwyyxxc Could you please run the CI and review. TIA