[XPU] Support checkpoint-engine weight update on Intel XPU

README.md

@@ -75,6 +75,28 @@ Use the flexible P2P implementation, notice this will install `mooncake-transfer
 pip install 'checkpoint-engine[p2p]'
 ```
 
+### Intel XPU (build from source)
+
+Intel XPU is supported for the **broadcast** update path. The cross-process weight handoff uses a native SYCL `ipc_memory` extension (the XPU counterpart of CUDA IPC) that is JIT-compiled at runtime. Install from source since XPU support is not yet in the released package. P2P is not supported on XPU (Mooncake has no Level Zero backend for XPU device memory).
+
+Requirements:
+- Intel XPU build of PyTorch — `torch.xpu.is_available()` returns `True`, `torch>=2.9` (for the device `.uuid` property; the SYCL extension build also needs `torch>=2.7`). See the [PyTorch XPU install guide](https://pytorch.org/docs/stable/notes/get_start_xpu.html); this is not the default PyPI `torch`.
+- Intel oneAPI 2026.0+ providing the `icpx` compiler with SYCL IPC memory support. Needed at runtime (first weight update), not at `pip install` time.
+
+```Bash
+git clone https://github.com/MoonshotAI/checkpoint-engine.git
+cd checkpoint-engine
+pip install -e .    # no [p2p] extra on XPU
+```
+
+Make `icpx` discoverable in the runtime environment, either by sourcing oneAPI (`source /opt/intel/oneapi/setvars.sh`) or by setting `CMPLR_ROOT`. If neither is set, `icpx` is auto-detected under `/opt/intel/oneapi/compiler/*/bin` and then on `PATH`. The extension then builds automatically on first use; `ParameterServer` also prebuilds it at startup so the one-time compile stays out of the update window.
+
+Verify the build and the IPC path on the target machine:
+
+```Bash
+pytest tests/test_xpu_ipc.py    # hardware-gated; skipped without an Intel GPU + buildable extension
+```
+
 ## Getting Started
 
 Prepare an H800 or H20 machine with 8 GPUs with vLLM. Be sure to include [/collective_rpc API endpoint](https://github.com/vllm-project/vllm/commit/f7cf5b512ee41f36613deb2471a44de5f304f70d) commit (available in main branch) since checkpoint-engine will use this endpoint to update weights. vLLM version `v0.10.2` is fully tested and recommended.

checkpoint_engine/device_utils.py

@@ -1,4 +1,5 @@
 import ctypes
+import gc
 import os
 import re
 import socket
@@ -209,9 +210,20 @@ def _is_torch_npu_available(self) -> bool:
         except ImportError:
             return False
 
+    def _is_torch_xpu_available(self) -> bool:
+        try:
+            if hasattr(torch, "xpu") and callable(getattr(torch.xpu, "is_available", None)):
+                return torch.xpu.is_available()
+            else:
+                return False
+        except Exception:  # noqa: BLE001
+            return False
+
     def _detect_device_type(self) -> str:
         if self._is_torch_npu_available():
             return "npu"
+        elif self._is_torch_xpu_available():
+            return "xpu"
         elif torch.cuda.is_available():
             return "cuda"
         else:
@@ -222,6 +234,8 @@ def _setup_device_module(self):
             import torch_npu
 
             self.device_module = torch_npu.npu
+        elif self.device_type == "xpu":
+            self.device_module = torch.xpu
         elif self.device_type == "cuda":
             self.device_module = torch.cuda
         else:
@@ -231,6 +245,8 @@ def _setup_device_module(self):
     def backend(self) -> str:
         if self.device_type == "npu":
             return "hccl"
+        elif self.device_type == "xpu":
+            return "xccl"
         elif self.device_type == "cuda":
             return "nccl"
         else:
@@ -240,7 +256,7 @@ def backend(self) -> str:
     def transfer_engine_protocol(self) -> str:
         if self.device_type == "npu":
             return "ascend_direct"
-        elif self.device_type == "cuda":
+        elif self.device_type in ("cuda", "xpu"):
             if has_efa_pci():
                 return "efa"
             else:
@@ -255,3 +271,42 @@ def rdma_device(self, rank: int) -> str:
             return _get_my_rdma_device(rank, self.device_module.device_count(), _get_rdma_devices())
         else:
             raise TypeError("The current transfer engine protocol is not supported")
+
+    def ipc_collect(self) -> None:
+        """Reclaim memory held by stale IPC handles."""
+        if self.device_type in ("cuda", "npu"):
+            self.device_module.ipc_collect()
+        elif self.device_type == "xpu":
+            # SYCL ipc_memory frees on close_handle; there is no cache to collect.
+            pass
+        else:
+            raise TypeError("The current device type is not supported")
+
+    def supports_inplace_pin(self) -> bool:
+        """Whether in-place host-memory pinning (cudaHostRegister) is available -- CUDA only."""
+        return self.device_type == "cuda"
+
+    def supports_device_ipc(self) -> bool:
+        """Whether cross-process IPC of *device* tensors works for this backend.
+
+        CUDA/NPU use ``torch.multiprocessing.reductions``; XPU uses the native SYCL
+        ``ipc_memory`` extension, available when it can be built (oneAPI >= 2026.0).
+        """
+        if self.device_type in ("cuda", "npu"):
+            return True
+        if self.device_type == "xpu":
+            from checkpoint_engine import xpu_ipc
+
+            return xpu_ipc.is_available()
+        return False
+
+    def supports_device_p2p(self) -> bool:
+        """Whether P2P (Mooncake) transfer of *device* memory works for this backend (CUDA/NPU only)."""
+        return self.device_type in ("cuda", "npu")
+
+    def host_empty_cache(self) -> None:
+        """Release cached pinned host memory (``_host_emptyCache`` on CUDA; else ``gc.collect``)."""
+        if self.device_type == "cuda":
+            torch._C._host_emptyCache()
+        else:
+            gc.collect()

checkpoint_engine/distributed/base.py

@@ -229,7 +229,12 @@ def use_backend(backend: str | None):
         "vllm_hccl": ".vllm_hccl.DistributedHccl",
     }
     if backend not in mapping:
-        raise ValueError(f"Unsupported custom backend: {backend}")
+        # XPU has no custom backend; leave custom_dist unset to use the default xccl TorchBackend.
+        raise ValueError(
+            f"Unsupported custom backend: {backend}. "
+            f"Supported custom backends: {sorted(mapping)}. "
+            "XPU is not supported here; leave custom_dist unset to use the default xccl backend."
+        )
 
     module_path, class_name = mapping[backend].rsplit(".", 1)
     module = importlib.import_module(module_path, "checkpoint_engine.distributed")

checkpoint_engine/ipc_handler.py

@@ -0,0 +1,132 @@
+"""Pluggable device-buffer handoff between the ParameterServer and the worker.
+
+The broadcast path shares a device buffer with the colocated worker. CUDA/NPU use
+:class:`TorchIpcHandler` (``torch.multiprocessing`` CUDA IPC, wire-format
+unchanged); XPU uses :class:`XpuIpcHandler` (native SYCL ``ipc_memory``).
+The handle is always a picklable, self-contained value, so the producer's
+``export`` -> ZMQ ``send_pyobj`` -> consumer ``attach`` flow is identical for both;
+each side calls ``detach`` on cleanup.
+"""
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import torch
+from loguru import logger
+from torch.multiprocessing.reductions import reduce_tensor
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from typing_extensions import Self
+
+    from checkpoint_engine.device_utils import DeviceManager
+
+
+def _rebuild_ipc(handle: tuple["Callable", tuple], device_id: int | None = None) -> torch.Tensor:
+    func, args = handle
+    list_args = list(args)
+    if device_id is not None:
+        # the key is to change device id to the current device id
+        # in case two processes have different CUDA_VISIBLE_DEVICES
+        list_args[6] = device_id
+    return func(*list_args)
+
+
+class IpcHandler(ABC):
+    """Hands a device buffer from the producer (ps) to the consumer (worker)."""
+
+    @abstractmethod
+    def export(self, buffer: torch.Tensor) -> Any:
+        """Producer: return the picklable handle to send over ZMQ."""
+
+    @abstractmethod
+    def attach(self, handle: Any, device_id: int) -> torch.Tensor:
+        """Consumer: reconstruct the shared device buffer from ``handle``."""
+
+    def detach(self) -> None:
+        """Release IPC resources on either side. No-op by default."""
+
+    # Used as a context manager so the handle is always released, without the
+    # caller needing its own try/finally.
+    def __enter__(self) -> "Self":
+        return self
+
+    def __exit__(self, *exc_info: object) -> None:
+        self.detach()
+
+
+class TorchIpcHandler(IpcHandler):
+    """CUDA/NPU zero-copy handoff via torch.multiprocessing CUDA IPC (unchanged)."""
+
+    def export(self, buffer: torch.Tensor) -> Any:
+        return reduce_tensor(buffer)
+
+    def attach(self, handle: Any, device_id: int) -> torch.Tensor:
+        assert isinstance(handle, tuple), f"expected reduce_tensor tuple, got {type(handle)}"
+        buffer = _rebuild_ipc(handle, device_id)
+        assert buffer.dtype == torch.uint8
+        return buffer
+
+
+class XpuIpcHandler(IpcHandler):
+    """Intel XPU zero-copy handoff via native SYCL ``ipc_memory`` (portable byte blob)."""
+
+    kind = "xpu_sycl"  # wire tag: identifies this handle format to the consumer
+
+    def __init__(self) -> None:
+        self._opened_ptr: int | None = None  # consumer: the mapping to unmap
+        self._exported_ptr: int | None = None  # producer: the retained exporter handle
+
+    def export(self, buffer: torch.Tensor) -> Any:
+        from checkpoint_engine import xpu_ipc
+
+        ptr = buffer.data_ptr()
+        handle_bytes = xpu_ipc.get_handle(ptr)
+        # Release only in detach(): freeing before the consumer opens can drop the
+        # fd under the level-zero-v2 UR adapter.
+        self._exported_ptr = ptr
+        return {
+            "kind": self.kind,
+            "handle_bytes": handle_bytes,
+            "nbytes": buffer.nbytes,
+        }
+
+    def attach(self, handle: Any, device_id: int) -> torch.Tensor:
+        from checkpoint_engine import xpu_ipc
+
+        assert isinstance(handle, dict) and handle.get("kind") == self.kind, (
+            f"expected {self.kind} handle dict, got {type(handle)}"
+        )
+        ptr = xpu_ipc.open_handle(handle["handle_bytes"], device_id)
+        self._opened_ptr = ptr
+        buffer = xpu_ipc.wrap_tensor(ptr, handle["nbytes"], device_id)
+        assert buffer.dtype == torch.uint8
+        return buffer
+
+    def detach(self) -> None:
+        # Consumer unmaps its opened pointer; producer releases its exported handle.
+        # At most one of the two is set on any given instance.
+        from checkpoint_engine import xpu_ipc
+
+        if self._opened_ptr is not None:
+            try:
+                torch.xpu.synchronize()  # no in-flight reads before unmapping
+                xpu_ipc.close_handle(self._opened_ptr)
+            except Exception as e:  # noqa: BLE001
+                logger.debug(f"xpu ipc close_handle failed during detach: {e}")
+            self._opened_ptr = None
+        if self._exported_ptr is not None:
+            try:
+                xpu_ipc.release_handle(self._exported_ptr)
+            except Exception as e:  # noqa: BLE001
+                logger.debug(f"xpu ipc release_handle failed during detach: {e}")
+            self._exported_ptr = None
+
+
+def build_ipc_handler(device_manager: "DeviceManager") -> IpcHandler:
+    """Select the IPC handler for the current device backend."""
+    if device_manager.device_type == "xpu":
+        return XpuIpcHandler()
+    return TorchIpcHandler()