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

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."""
-        # 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>``.
+        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)