Replace hardcoded torch.cuda calls in the Encoder/Prefill/Decode path with
torch.get_device_module equivalents, so staging-buffer KV transfer works on
Intel XPU. Two of the fixes are not XPU-specific: the encoder /send
synchronize was gated on is_cuda and so was skipped on every other
accelerator, and staging init hardcoded a "cuda:{gpu_id}" device string.
Breaking change for out-of-tree backends: init_staging_buffers,
init_staging_allocator and init_dcp_pack_buffers take a required keyword-only
device_type. The "cuda" default was dropped so no backend can silently
allocate staging on CUDA.
The DCP pack kernel is Triton-only, so off CUDA init_dcp_pack_buffers returns
no buffers and DCP transfers fall back to per-token RDMA.
Adds create_device_event and current_device_stream to srt/utils/common.py.
The decode-side scatter now waits on an event recorded on schedule_stream when
the room is registered, right after pre-alloc writes its req_to_token row,
instead of having no ordering edge. Waiting on the stream itself would also
queue each scatter behind any forward running there without overlap.
The XPU PD tests run on Mooncake in CI (nightly-xpu-2-gpu and
nightly-xpu-4-gpu). The NIXL XPU tests are manual tests under
test/manual/kv_transfer/ (test_disaggregation_nixl_xpu.py and
test_disaggregation_staging_nixl_xpu.py): they mirror the registered scenarios,
skip when nixl is not installed.
docker/xpu.Dockerfile builds the Mooncake transfer
engine from a pinned upstream commit with icpx (-DUSE_XPU=ON -DUSE_TENT=ON
-DWITH_STORE=OFF) and installs sglang-router. The Mooncake classes skip when
mooncake.engine is not importable, so the nightly stays green on an image
built before that Dockerfile step.
Mooncake's XPU support lives in its TENT engine, so those classes set
MC_USE_TENT=1 and MOONCAKE_PROTOCOL=tcp and need no RDMA NIC. The staging
decode server (TP 2) runs with TORCHINDUCTOR_USE_STATIC_CUDA_LAUNCHER=0. On
torch 2.13.0+xpu, inductor's static Triton launcher segfaults in load_kernel on
the compiled vocab-parallel embedding mask, which only runs at TP > 1. This is
not PD-specific: a plain Qwen2.5-1.5B or Llama-3.1-8B server at TP 2 crashes
the same way, with a cold inductor cache too. Turning off only the static
launcher keeps torch.compile on, unlike TORCHDYNAMO_DISABLE=1. Likely
pytorch/pytorch#199511 (XPU static launcher ABI mismatch); drop the override
once the XPU torch pin includes it.
Verified on 3x Intel Arc Pro B60: xpu/test_disaggregation_staging_xpu.py
(prefill TP 1, decode TP 2, so KV goes through the staging path), 4 passed;
test/registered/unit/disaggregation, 487 passed. With Mooncake over TCP:
disaggregation/test_disaggregation_xpu.py 2 passed, staging 4 passed. The same
Mooncake tests pass against a transfer-engine-only install of upstream Mooncake
built with the Dockerfile steps.
Motivation
PD disaggregation with staging-buffer KV transfer assumes CUDA in several places along the Encoder/Prefill/Decode path. On Intel XPU these assumptions either fail immediately (
torch.cuda.*calls on a build without CUDA) or silently skip required synchronization because it is gated onis_cuda.This PR makes the staging-buffer path device-agnostic and enables it end to end on Intel XPU, verified on 3x Intel Arc Pro B60.
Two of the fixes are not XPU-specific:
/sendpath synchronized only whenmm_embedding.is_cuda. The sync is required on every accelerator becausetransfer_syncdoes not order against the device stream; without it, the transport can read a buffer before the device writes land.device = f"cuda:{gpu_id}", so a non-CUDA backend could silently allocate staging buffers on CUDA.NIXL with staging was also broken on main for every device:
prefill.pyanddecode.pypassslot_layer_ids=toset_kv_buffer_tensors, which NIXL did not accept, so manager init failed withunexpected keyword argument 'slot_layer_ids'.Modifications
Device helpers (
srt/utils/common.py)create_device_eventandcurrent_device_stream, which resolve the device module from the given device instead of hard-codingtorch.cuda.Staging buffers (
disaggregation/common/staging_buffer.py)SGLANG_DISAGG_STAGING_BUFFER.gpu_id(with a warning on mismatch), so a caller passing"xpu:3"cannot allocate on a different device.data_ptr()values can set bit 63, whichtorch.tensor(..., dtype=torch.int64)rejects, and there is no Triton NPU backend; other supported accelerators take the torch gather/scatter path.ptr_type=XPU-kernel-space/XPU-user-spaceon XPU.Stream ordering
StagingBuffer.producer_stream()), resolved per call, instead oftorch.cuda.default_stream. The current stream is thread-local, and transfer workers run on their own threads.schedule_streamwhen the room is registered, right after pre-alloc writes itsreq_to_tokenrow. Previously there was no ordering edge; XPU has no CUDA-style default-stream synchronization, and a dedicated CUDA stream does not implicitly sync with unrelated work either. Waiting on the stream itself would also queue each scatter behind any forward running there without overlap.release_roomdrains it only when the room still has an in-flight scatter. The stream is shared by all rooms, so draining unconditionally would stall the scheduler thread behind unrelated rooms on the success path. The prefill gather stream stays lazy, since the decode-side allocator's buffer never gathers.Staging handler and transports (
common/staging_handler.py,mooncake/conn.py,nixl/conn.py)DecodeStagingHandlerderives its device from the KV buffers and requires an indexed device.device_typefromserver_args.deviceinto staging initialization.np.uint64for staged transfer descriptors so high device pointers are not truncated or read as signed.set_kv_buffer_tensorsacceptsslot_layer_ids(see Motivation); NIXL staging ispp_size == 1only, so the value is stored but not read.DCP pack (
disaggregation/common/dcp_pack.py)copy_mla_rows_into_packis Triton-only, so off CUDAinit_dcp_pack_buffersreturns no buffers and DCP transfers fall back to per-token RDMA.torch.cuda.default_stream, and enter it through the device-agnostic stream context.Encoder (
disaggregation/encoder/server.py,encoder/receiver.py)TensorWrappercopies from any non-CPU device, not only CUDA./sendsynchronizes for every non-CPU accelerator instead of gating onis_cuda, and the embedding-cache paths synchronize throughcurrent_device_stream.EmbeddingPoolallocate on the runtime device type instead ofcuda:{gpu_id}.Metadata buffers and QK normalization (
disaggregation/utils.py,models/utils.py)INTRA_NODE_NVLINKmetadata-buffer placement now also requiresis_cuda(). Theuse_mem_poolcontext stays ontorch.cuda, sincecustom_mem_poolis atorch.cuda.MemPool.apply_qk_normentersalt_streamthrough the device-agnostic stream context.Docker (
docker/xpu.Dockerfile)-DUSE_XPU=ON -DUSE_TENT=ON -DWITH_STORE=OFF) and installsglang-router, which the PD tests launch between prefill and decode.Breaking changes for out-of-tree backends
init_staging_buffers,init_staging_allocatorandinit_dcp_pack_bufferstake a required keyword-onlydevice_type. The old implicit"cuda"is intentionally gone, so omittingdevice_typefails loudly instead of allocating staging on the wrong accelerator.gather_all_layers_to_stagingno longer takesgpu_id; the device comes from the KV buffers.EmbeddingPooltakes keyword-onlydevice_type,gpu_idandsize_bytes.Total: 19 files, +891/-123.
Accuracy Tests
Unit tests (CPU, no accelerator required)
New tests run under
register_cpu_ciwith stubs/mocks:test_mooncake_transfer_batching.py(TestMooncakeStagingDeviceType): prefill staging buffers and the decode staging allocator receive the runtime device instead of an implicit"cuda". These cover the device-selection regression only, not end-to-end Mooncake staging transfer.test_dcp_pack.py(TestInitDcpPackBuffers): a non-Triton device gets no DCP pack buffers and never registers memory; CUDA buffers use the bare device type andgpu_id.test_encode_server.py: updated for the keyword-onlyEmbeddingPoolsignature.test/registered/unit/disaggregation/: 487 passed.XPU end-to-end
The registered XPU PD tests run on Mooncake. Mooncake's XPU support lives in its TENT engine, so the test classes set
MC_USE_TENT=1andMOONCAKE_PROTOCOL=tcpand need no RDMA NIC. They skip whenmooncake.engineis not importable, so the nightly stays green on an image built before the Dockerfile step.Verified on 3x Intel Arc Pro B60:
test/registered/xpu/test_disaggregation_staging_xpu.py(prefill TP 1, decode TP 2, so KV goes through the staging path): 4 passed, onnightly-xpu-4-gpu.test/registered/disaggregation/test_disaggregation_xpu.py: 2 passed, onnightly-xpu-2-gpu.The concurrent-vs-serial case matters most: it exercises
enable_overlaptogether with a radix-cache prefix hit, and compares against the serial answers at temperature 0, so KV corruption shows up as a token mismatch. It needs the dependency PR above.The staging e2e test lives under the
xpuvendor tree because the registered-test taxonomy rejects new paths outside the<kind>/<subsystem>layout, while vendor directories are exempt.The NIXL variants are manual tests under
test/manual/kv_transfer/:test_disaggregation_nixl_xpu.py(basic completion) andtest_disaggregation_staging_nixl_xpu.py(staging). They mirror the registered scenarios and skip when nixl is not installed.The staging decode server (TP 2) runs with
TORCHINDUCTOR_USE_STATIC_CUDA_LAUNCHER=0. On torch 2.13.0+xpu, inductor's static Triton launcher segfaults inload_kernelon the compiled vocab-parallel embedding mask, which only runs at TP > 1. This is not PD-specific: a plain Qwen2.5-1.5B or Llama-3.1-8B server at TP 2 crashes the same way. Turning off only the static launcher keepstorch.compileon, unlikeTORCHDYNAMO_DISABLE=1. Likely pytorch/pytorch#199511; the override can go once the XPU torch pin includes the fix.Performance
No CUDA-path performance regression is intended. The changes are mostly device selection and stream/event handling on initialization or correctness-critical paths.
release_roomnow skips the scatter-stream drain for rooms with no in-flight scatter, which removes a scheduler-thread stall on the success path.Checklist
CI States
Latest PR Test (Base): ⏳ Run #38372317850
Latest PR Test (Extra): ❌ Run #38372317378
Latest PR Test (AMD ROCm 10): ⏳ Run #38372317893