fix: bind worker IPC listener to its CUDA device

[shipiyouniao]

Summary

Bind each TP IPC listener thread to the CUDA device selected by its worker before the listener performs CUDA-backed map/unmap operations.

CUDA current-device state is thread-local. A newly-created listener thread otherwise starts on device 0 even when the worker owns another TP device. The integration now captures the worker's current device and the listener restores it at thread entry.

Scope

• pass the worker CUDA device from the vLLM integration
• pass the worker CUDA device from the SGLang integration
• restore that device in the IPC listener thread
• keep the listener API backward compatible when no device is supplied

Real GPU reproduction

Validated on two isolated RTX 4090 GPUs using vllm/vllm-openai:v0.22.1, PyTorch 2.11.0+cu130, CUDA 13.0, and MPS active thread percentage 100.

The test initializes the main worker thread on logical CUDA device 1, starts the real Unix-socket listener, and sends a real kv_tensors_created IPC request. The VMM callback is replaced only with a recorder for torch.cuda.current_device() so the test isolates the listener thread's CUDA binding without allocating model KV pages.

Before this change (55a534c):

PASS main_device=1 listener_device=0 expected=0

With this change (0c9386e):

PASS main_device=1 listener_device=1 expected=1

This reproduces the wrong-device behavior on real CUDA and confirms that the listener executes on the worker-selected device after the fix.

Tests

• real two-GPU CUDA + Unix-socket before/after reproduction described above
• python -m pytest -q tests/test_tp_listener_device.py (3 passed)
• python -m ruff check ... (passed)
• git diff --check (passed)

[gemini-code-assist]

Caution

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

[Copilot]

Pull request overview

This PR ensures the TP worker IPC listener thread restores the worker’s selected CUDA device before performing CUDA-backed KV cache map/unmap operations, addressing the fact that CUDA “current device” is thread-local and defaults to device 0 in new threads.

Changes:

• Extend start_worker_listener_thread to accept an optional device_index and restore it at listener thread entry.
• Update vLLM and SGLang integrations to pass the worker’s CUDA device index into the listener thread.
• Add tests that mock CUDA state to verify the listener restores the expected device and integrations pass the correct index.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
tests/test_tp_listener_device.py	Adds unit tests validating listener-thread device restoration and integration call behavior.
kvcached/tp_ipc_util.py	Adds optional device_index parameter and restores CUDA device inside the listener thread.
kvcached/integration/vllm/interfaces.py	Passes a CUDA device index into the worker listener thread.
kvcached/integration/sglang/interfaces.py	Passes a CUDA device index into the worker listener thread.

Comments suppressed due to low confidence (2)

kvcached/integration/vllm/interfaces.py:82

• Same issue here: device_index is taken from torch.cuda.current_device(), but this may not match the device string used to initialize kvcached (especially if the caller passes an explicit device without changing the current device). Derive the index from the normalized device string when possible.

        start_worker_listener_thread(
            tp_rank,
            pp_rank,
            device_index=int(torch.cuda.current_device()),
        )

kvcached/tp_ipc_util.py:105

• The docstring doesn't mention the new device_index parameter or that it changes thread-local CUDA state, which is important for callers (and for understanding why this argument exists). Document the new behavior briefly in the function docstring.

    """
    Start a thread that listens for messages on the worker socket.
    pp_rank is used to create a PP-stage-specific subdirectory so that
    concurrent SGLang PP stages do not bind the same socket path.
    """

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[shipiyouniao]

Moving this back to draft until the listener-thread wrong-device failure is reproduced and validated on real multi-GPU hardware. The mocked CUDA regression test is useful, but it is not sufficient evidence by itself for this low-level runtime change.

[shipiyouniao]

Added a real two-GPU before/after reproduction to the PR description. The old code runs the Unix-socket listener callback on logical CUDA 0 while the worker main thread owns CUDA 1; this change restores CUDA 1 in the listener. The test uses real CUDA and real UDS transport, with only the VMM action replaced by a current-device recorder.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.