[XPU] Make checkpoint_engine worker device-agnostic

[siju-samuel]

Important

Depends on MoonshotAI/checkpoint-engine#96
(adds Intel XPU support to checkpoint-engine). That PR must merge first
pyproject_xpu.toml tracks checkpoint-engine's upstream main, which only
resolves to XPU-capable code once #96 lands.

Motivation

SGLang's checkpoint-engine integration enables fast in-place weight updates
for a running server (the RL / online-serving use case), but the worker
hardcoded torch.cuda for device UUID/id resolution. As a result the ZMQ
handshake with checkpoint-engine's ParameterServer only worked on CUDA.

This makes the worker device-agnostic so the same path works on Intel XPU,
and along the way fixes a latent NPU bug: the ParameterServer keys NPU
devices as NPU-<uuid>, but the old CUDA fallback would have keyed them as
GPU-<uuid> — a key that never resolves.

Modifications

• srt/checkpoint_engine/checkpoint_engine_worker.py
  • _accelerator_type() now delegates to SGLang's device-agnostic
    get_device() instead of assuming CUDA; device-module lookups go through
    get_device_module() (removing the getattr(torch, ...) indirection).
  • get_device_uuid() mirrors checkpoint_engine.ps::_get_physical_gpu_id
    exactly: NPU → NPU-<uuid>, every other accelerator (CUDA, XPU) →
    GPU-<uuid>. This is the fix for the NPU key mismatch above.
• python/pyproject_xpu.toml
  • Add the checkpoint-engine optional dependency for the XPU build. XPU
    support is not yet in a released checkpoint-engine, so it tracks upstream
    source (git+https://github.com/MoonshotAI/checkpoint-engine.git) until a
    release ships it, mirroring how sgl-kernel-xpu is referenced in the same
    file.
• test/registered/unit/checkpoint_engine/test_checkpoint_engine_worker.py
  • CPU-mockable tests for cuda/xpu/npu key-format routing and the
    AssertionError→ValueError wrapping.
  • A hardware-gated XPU test asserting the worker's real device UUID equals
    the key the ParameterServer independently derives (the actual handshake
    contract).
  • Registered for both CPU (base-a-test-cpu) and XPU
    (stage-b-test-1-gpu-xpu) CI.

Accuracy Tests

Not applicable — this changes device-identity resolution for the weight-update
handshake, not model forward/kernels. Model outputs are unaffected.

Speed Tests and Profiling

Not applicable — no change to the inference hot path. get_device() /
get_device_module() are called only during the weight-update handshake, not
per token.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results (N/A — see above).
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #30145719133
Latest PR Test (Extra): ❌ Run #30145719088

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[polisettyvarma]

why not get_device, directly ?

[polisettyvarma]

after PR merge and any release let's use/pin it like CUDA