[XPU] Make checkpoint_engine worker device-agnostic

python/pyproject_xpu.toml

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

python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py

@@ -22,6 +22,8 @@
 import torch
 import zmq
 
+from sglang.srt.utils import get_device, get_device_module, is_npu
+
 try:
     from checkpoint_engine.worker import update_weights_from_ipc
 except ImportError:
@@ -33,6 +35,14 @@
 logger = logging.getLogger(__name__)
 
 
+def _accelerator_type() -> str:
+    """Active accelerator device type -- ``"cuda"``, ``"xpu"``, ``"npu"``, etc.
+
+    Resolved via SGLang's device-agnostic ``get_device()`` rather than assuming
+    CUDA, so each backend is labeled correctly (e.g. NPU is not mistaken for CUDA)."""
+    return get_device()
+
+
 class SGLangCheckpointEngineWorkerExtension:
     """
     Worker extension for SGLang to support checkpoint-engine IPC weight updates.
@@ -100,17 +110,24 @@ def __init__(self, model_runner):
         self.model_runner = model_runner
 
     def get_device_uuid(self) -> str:
-        """Get the UUID of current device."""
-        # Get device UUID for current device
-        device_id = torch.cuda.current_device()
+        """Physical GPU id, matching checkpoint-engine's ParameterServer key.
+
+        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>``."""
+        if is_npu():
+            from checkpoint_engine.device_utils import npu_generate_uuid
+
+            return f"NPU-{npu_generate_uuid()}"
+        device_module = get_device_module()
+        device_id = device_module.current_device()
         try:
-            return f"GPU-{torch.cuda.get_device_properties(device_id).uuid!s}"
+            return f"GPU-{device_module.get_device_properties(device_id).uuid!s}"
         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 +147,7 @@ def post_hook():
                     if quant_method is not None:
                         # Move parameters to device if needed for quantization processing
                         target_device = torch.device(
-                            "cuda", torch.cuda.current_device()
+                            _accelerator_type(), 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,122 @@
+"""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.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci, register_xpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+register_xpu_ci(est_time=30, suite="stage-b-test-1-gpu-xpu")
+
+import unittest
+from unittest.mock import MagicMock, patch
+
+import torch
+from sglang.srt.checkpoint_engine.checkpoint_engine_worker import (
+    SGLangCheckpointEngineWorkerExtensionImpl,
+    _accelerator_type,
+)
+from sglang.srt.utils import is_xpu
+from sglang.test.test_utils import CustomTestCase
+
+_WORKER_MOD = "sglang.srt.checkpoint_engine.checkpoint_engine_worker"
+
+
+class TestAcceleratorType(CustomTestCase):
+    def test_delegates_to_get_device(self):
+        # _accelerator_type must not hardcode cuda; it reports whatever get_device
+        # resolves, so NPU is labeled "npu" rather than mistaken for "cuda".
+        for dev in ("cuda", "xpu", "npu"):
+            with patch(f"{_WORKER_MOD}.get_device", return_value=dev):
+                self.assertEqual(_accelerator_type(), dev)
+
+
+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, uuid="abcd-1234"):
+        mod = MagicMock()
+        mod.current_device.return_value = current
+        props = MagicMock()
+        props.uuid = uuid
+        mod.get_device_properties.return_value = props
+        return mod
+
+    def test_device_uuid_cuda(self):
+        worker = self._make_worker()
+        fake = self._fake_device_module(current=0, uuid="cuda-uuid")
+        with (
+            patch(f"{_WORKER_MOD}.is_npu", return_value=False),
+            patch(f"{_WORKER_MOD}.get_device_module", return_value=fake),
+        ):
+            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, uuid="xpu-uuid")
+        with (
+            patch(f"{_WORKER_MOD}.is_npu", return_value=False),
+            patch(f"{_WORKER_MOD}.get_device_module", return_value=fake),
+        ):
+            # 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 = MagicMock()
+        fake.current_device.return_value = 1
+        fake.get_device_properties.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),
+            self.assertRaises(ValueError),
+        ):
+            worker.get_device_uuid()
+
+
+@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)