From 2a06b91c98465cd18627eabaeb037c6d3fa8e4f8 Mon Sep 17 00:00:00 2001 From: Siju Samuel Date: Fri, 24 Jul 2026 04:16:45 +0000 Subject: [PATCH] [XPU] Support checkpoint_engine weight updates on XPU The checkpoint_engine worker hardcoded torch.cuda for device UUID/id resolution, so the ZMQ handshake with checkpoint-engine's ParameterServer only worked on CUDA. Make it device-agnostic: - _accelerator_type() now delegates to SGLang's get_device() instead of assuming CUDA, and device-module lookups go through get_device_module(). - get_device_uuid() mirrors ps.py::_get_physical_gpu_id exactly: NPU keys as NPU-, every other accelerator (CUDA, XPU) as GPU-. This also fixes NPU, which the old cuda-fallback would have keyed as GPU- and never matched the ParameterServer. Wire the checkpoint-engine optional dependency into pyproject_xpu.toml. XPU support is not yet in a released checkpoint-engine, so it tracks the upstream source (git+https://github.com/MoonshotAI/checkpoint-engine.git) until a release ships it, mirroring how sgl-kernel-xpu is referenced. Tests (test/registered/unit/checkpoint_engine/): CPU-mockable routing for cuda/xpu/npu key formats and the AssertionError->ValueError wrapping, plus a hardware-gated XPU case asserting the worker's real UUID matches the ParameterServer's independently derived key. Registered for CPU and XPU CI. --- python/pyproject_xpu.toml | 1 + .../checkpoint_engine_worker.py | 29 ++++- .../test_checkpoint_engine_worker.py | 122 ++++++++++++++++++ 3 files changed, 146 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py diff --git a/python/pyproject_xpu.toml b/python/pyproject_xpu.toml index bcd3d1a72abc..ff6b31b6650d 100644 --- a/python/pyproject_xpu.toml +++ b/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", diff --git a/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py b/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py index 6f11c7872540..dd0c5f069c42 100644 --- a/python/sglang/srt/checkpoint_engine/checkpoint_engine_worker.py +++ b/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-`` key, every other accelerator ``GPU-``.""" + 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) diff --git a/test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py b/test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py new file mode 100644 index 000000000000..5f3d188e1fb1 --- /dev/null +++ b/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-`` +for CUDA/XPU and ``NPU-`` 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- 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-, so a GPU- 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)