[XPU] Make checkpoint_engine worker device-agnostic (#32382)

.github/workflows/pr-test-xpu.yml

@@ -151,6 +151,11 @@ jobs:
           docker exec ci_sglang_xpu /opt/venv/bin/python3 -m pip install --no-cache-dir --no-deps xgrammar==0.1.33
           docker exec ci_sglang_xpu /bin/bash -c '/opt/venv/bin/hf auth login --token ${HF_TOKEN}'
 
+      - name: Install checkpoint-engine extra (optional; tests skip if unavailable)
+        continue-on-error: true
+        run: |
+          docker exec -w /sglang-checkout/python ci_sglang_xpu /opt/venv/bin/python3 -m pip install --no-cache-dir ".[checkpoint-engine]" --extra-index-url https://download.pytorch.org/whl/xpu
+
       - name: Run tests
         # Old stage-a (30) + stage-b (120) = 150; +20% headroom for --enable-retry.
         timeout-minutes: 180

docs/docs/advanced_features/checkpoint_engine.mdx

@@ -25,7 +25,7 @@ pip install 'checkpoint-engine[p2p]'
 
 The system consists of two main components:
 
-1. **SGLang Server**: Runs with `--wait-for-initial-weights` flag to wait for weights before becoming ready
+1. **SGLang Server**: Runs with `--checkpoint-engine-wait-weights-before-ready` flag to wait for weights before becoming ready
 2. **Checkpoint Engine Workers**: Separate processes (managed by torchrun) that load and distribute model weights
 
 The checkpoint engine uses a parameter server architecture with support for:
@@ -43,7 +43,7 @@ python -m sglang.launch_server \
     --model-path Qwen/Qwen3-8B \
     --tp 8 \
     --load-format dummy \
-    --wait-for-initial-weights
+    --checkpoint-engine-wait-weights-before-ready
 ```
 
 **Terminal 2 - Run Checkpoint Engine:**
@@ -75,7 +75,7 @@ python -m sglang.launch_server \
     --model-path Qwen/Qwen3-8B \
     --tp 8 \
     --load-format dummy \
-    --wait-for-initial-weights \
+    --checkpoint-engine-wait-weights-before-ready \
     --host [IP]
 ```
 
@@ -110,7 +110,7 @@ python -m sglang.launch_server \
     --model-path Qwen/Qwen3-8B \
     --tp 8 \
     --load-format dummy \
-    --wait-for-initial-weights \
+    --checkpoint-engine-wait-weights-before-ready \
     --host [IP]
 ```
 
@@ -147,7 +147,7 @@ python -m sglang.launch_server \
     --model-path Qwen/Qwen3-8B \
     --tp 8 \
     --load-format dummy \
-    --wait-for-initial-weights \
+    --checkpoint-engine-wait-weights-before-ready \
     --host [IP] \
     --dist-init-addr [IP]:9120 \
     --nnodes 2 \
@@ -185,7 +185,7 @@ python -m sglang.launch_server \
     --model-path Qwen/Qwen3-8B \
     --tp 8 \
     --load-format dummy \
-    --wait-for-initial-weights \
+    --checkpoint-engine-wait-weights-before-ready \
     --host [IP] \
     --dist-init-addr [IP]:9120 \
     --nnodes 2 \
@@ -220,7 +220,7 @@ torchrun --nproc-per-node 8 \
 ### SGLang Server Options
 
 - `--load-format dummy`: Use dummy format for initial loading (allows overlapping with other tasks)
-- `--wait-for-initial-weights`: Wait for checkpoint engine to provide weights before becoming ready
+- `--checkpoint-engine-wait-weights-before-ready`: Wait for checkpoint engine to provide weights before becoming ready
 - `--host`: Host address for multi-node setups
 - `--dist-init-addr`: Distributed initialization address for tensor parallelism
 

python/pyproject_xpu.toml

@@ -73,6 +73,7 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+checkpoint-engine = ["checkpoint-engine @ git+https://github.com/MoonshotAI/checkpoint-engine.git@f772d1858204fa82bdb0436fdfec137f5de2972c"]
 diffusion = [
   "addict==2.4.0",
   "av==16.1.0",

python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py

@@ -22,6 +22,9 @@ from typing import Callable, Dict, Optional
 import torch
 import zmq
 
+from sglang.srt.platforms import current_platform
+from sglang.srt.utils import get_device, get_device_module, is_npu
+
 try:
     from checkpoint_engine.worker import update_weights_from_ipc
 except ImportError:
@@ -100,17 +103,25 @@ class SGLangCheckpointEngineWorkerExtensionImpl(SGLangCheckpointEngineWorkerExte
         self.model_runner = model_runner
 
     def get_device_uuid(self) -> str:
-        """Get the UUID of current device."""
+        """Physical GPU id, matching checkpoint-engine's ParameterServer key.
-        # Get device UUID for current device
+
-        device_id = torch.cuda.current_device()
+        Must equal ps.py::_get_physical_gpu_id for the ZMQ handshake to resolve:
+        NPU uses an ``NPU-<uuid>`` key, every other accelerator ``GPU-<uuid>``.
+        NPU keeps its own branch because its key is derived from npu-smi rather
+        than device properties, and no NPU platform implements get_device_uuid."""
+        if is_npu():
+            from checkpoint_engine.device_utils import npu_generate_uuid
+
+            return f"NPU-{npu_generate_uuid()}"
+        device_id = get_device_module().current_device()
         try:
-            return f"GPU-{torch.cuda.get_device_properties(device_id).uuid!s}"
+            return f"GPU-{current_platform.get_device_uuid(device_id)}"
         except AssertionError as e:
             raise ValueError(f"Failed to get GPU UUID for device {device_id}") from e
 
     def get_device_id(self) -> int:
         """Get the device ID."""
-        return torch.cuda.current_device()
+        return get_device_module().current_device()
 
     def get_model_loader(self) -> Callable:
         """Get the model weight loader function."""
@@ -130,7 +141,7 @@ class SGLangCheckpointEngineWorkerExtensionImpl(SGLangCheckpointEngineWorkerExte
                     if quant_method is not None:
                         # Move parameters to device if needed for quantization processing
                         target_device = torch.device(
-                            "cuda", torch.cuda.current_device()
+                            get_device(), get_device_module().current_device()
                         )
                         with device_loading_context(module, target_device):
                             quant_method.process_weights_after_loading(module)

test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py

@@ -0,0 +1,141 @@
+"""Unit tests for srt/checkpoint_engine/checkpoint_engine_worker.py — no server, no model loading.
+
+Focus: device resolution so the ZMQ handshake key matches checkpoint-engine's
+ParameterServer (ps.py::_get_physical_gpu_id) on every backend -- ``GPU-<uuid>``
+for CUDA/XPU and ``NPU-<uuid>`` for NPU. These paths are pure namespace routing
+(``get_device`` / ``get_device_module`` / ``is_npu``) and are fully mockable on CPU.
+
+Skipped entirely unless the ``checkpoint-engine`` extra is installed, since the
+worker module refuses to import without it.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+# Unit tests may register CPU suites only (scripts/lint/check_registered_tests.py).
+# TestWorkerDeviceUuidOnXpu below still self-skips off XPU, so it stays runnable
+# by hand on an Intel GPU host.
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+import importlib.util
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+
+from sglang.srt.utils import is_xpu
+from sglang.test.test_utils import CustomTestCase
+
+# checkpoint-engine is an optional extra (sglang[checkpoint-engine]) that CI does
+# not install, and the worker module raises ImportError at import time without it.
+# Probe first so this file stays importable: an unguarded import would fail
+# collection and take the whole file down rather than skipping.
+_HAS_CHECKPOINT_ENGINE = importlib.util.find_spec("checkpoint_engine") is not None
+
+if _HAS_CHECKPOINT_ENGINE:
+    from sglang.srt.checkpoint_engine.checkpoint_engine_worker import (
+        SGLangCheckpointEngineWorkerExtensionImpl,
+    )
+
+_WORKER_MOD = "sglang.srt.checkpoint_engine.checkpoint_engine_worker"
+_NO_CKPT_ENGINE = "requires the checkpoint-engine optional dependency"
+
+
+@unittest.skipUnless(_HAS_CHECKPOINT_ENGINE, _NO_CKPT_ENGINE)
+class TestWorkerDeviceResolution(CustomTestCase):
+    """get_device_uuid / get_device_id must route through the active accelerator
+    namespace and emit the key the ParameterServer expects."""
+
+    def _make_worker(self):
+        # model_runner is unused by the device-resolution methods under test.
+        return SGLangCheckpointEngineWorkerExtensionImpl(model_runner=MagicMock())
+
+    def _fake_device_module(self, *, current=3):
+        mod = MagicMock()
+        mod.current_device.return_value = current
+        return mod
+
+    def _fake_platform(self, *, uuid="abcd-1234"):
+        # The uuid now comes from the platform layer (current_platform), which
+        # already returns str(get_device_properties(id).uuid) for cuda/xpu.
+        plat = MagicMock()
+        plat.get_device_uuid.return_value = uuid
+        return plat
+
+    def test_device_uuid_cuda(self):
+        worker = self._make_worker()
+        fake = self._fake_device_module(current=0)
+        with (
+            patch(f"{_WORKER_MOD}.is_npu", return_value=False),
+            patch(f"{_WORKER_MOD}.get_device_module", return_value=fake),
+            patch(
+                f"{_WORKER_MOD}.current_platform", self._fake_platform(uuid="cuda-uuid")
+            ),
+        ):
+            self.assertEqual(worker.get_device_uuid(), "GPU-cuda-uuid")
+            self.assertEqual(worker.get_device_id(), 0)
+
+    def test_device_uuid_xpu(self):
+        worker = self._make_worker()
+        fake = self._fake_device_module(current=2)
+        with (
+            patch(f"{_WORKER_MOD}.is_npu", return_value=False),
+            patch(f"{_WORKER_MOD}.get_device_module", return_value=fake),
+            patch(
+                f"{_WORKER_MOD}.current_platform", self._fake_platform(uuid="xpu-uuid")
+            ),
+        ):
+            # XPU shares CUDA's GPU-<uuid> format; only the namespace differs.
+            self.assertEqual(worker.get_device_uuid(), "GPU-xpu-uuid")
+            self.assertEqual(worker.get_device_id(), 2)
+
+    def test_device_uuid_npu_uses_npu_prefix(self):
+        # NPU must NOT be treated as CUDA: the ParameterServer keys it as
+        # NPU-<npu_generate_uuid()>, so a GPU-<uuid> key would never resolve.
+        worker = self._make_worker()
+        with (
+            patch(f"{_WORKER_MOD}.is_npu", return_value=True),
+            patch(
+                "checkpoint_engine.device_utils.npu_generate_uuid",
+                return_value="1.2.3.4-0",
+            ),
+        ):
+            self.assertEqual(worker.get_device_uuid(), "NPU-1.2.3.4-0")
+
+    def test_device_uuid_wraps_assertion_error(self):
+        worker = self._make_worker()
+        fake = self._fake_device_module(current=1)
+        plat = MagicMock()
+        plat.get_device_uuid.side_effect = AssertionError("no uuid")
+        with (
+            patch(f"{_WORKER_MOD}.is_npu", return_value=False),
+            patch(f"{_WORKER_MOD}.get_device_module", return_value=fake),
+            patch(f"{_WORKER_MOD}.current_platform", plat),
+            self.assertRaises(ValueError),
+        ):
+            worker.get_device_uuid()
+
+
+@unittest.skipUnless(_HAS_CHECKPOINT_ENGINE, _NO_CKPT_ENGINE)
+@unittest.skipUnless(is_xpu(), "requires an Intel XPU")
+class TestWorkerDeviceUuidOnXpu(CustomTestCase):
+    """Hardware-gated: the real XPU key must match what checkpoint-engine's
+    ParameterServer derives, or the ZMQ handshake silently fails on XPU."""
+
+    def test_real_uuid_matches_parameter_server(self):
+        from checkpoint_engine.device_utils import DeviceManager
+        from checkpoint_engine.ps import _get_physical_gpu_id
+
+        worker = SGLangCheckpointEngineWorkerExtensionImpl(model_runner=MagicMock())
+        key = worker.get_device_uuid()
+
+        self.assertTrue(key.startswith("GPU-"), key)
+        self.assertEqual(worker.get_device_id(), torch.xpu.current_device())
+
+        # Independently derived by the ParameterServer side; the two must agree.
+        dm = DeviceManager()
+        self.assertEqual(dm.device_type, "xpu")
+        self.assertEqual(key, _get_physical_gpu_id(dm, torch.xpu.current_device()))
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)