fix: bind worker IPC listener to its CUDA device

kvcached/integration/sglang/interfaces.py

@@ -16,7 +16,7 @@
     clear_registered_kv_cache_pools,
     register_kv_cache_pool,
 )
-from kvcached.tp_ipc_util import start_worker_listener_thread
+from kvcached.tp_ipc_util import resolve_gpu_device_index, start_worker_listener_thread
 from kvcached.utils import CONTIGUOUS_LAYOUT, PAGE_SIZE, get_kvcached_logger, normalize_gpu_device
 from kvcached.vmm_ops import (
     create_kv_tensors,
@@ -58,7 +58,11 @@ def init_kvcached(
 
     if world_size > 1:
         # start the listener thread for tensor parallel kv cache management
-        start_worker_listener_thread(tp_rank, pp_rank)
+        start_worker_listener_thread(
+            tp_rank,
+            pp_rank,
+            device_index=resolve_gpu_device_index(device),
+        )
 
 
 def shutdown_kvcached() -> None:

kvcached/integration/vllm/interfaces.py

@@ -16,7 +16,7 @@
     clear_registered_kv_cache_pools,
     register_kv_cache_pool,
 )
-from kvcached.tp_ipc_util import start_worker_listener_thread
+from kvcached.tp_ipc_util import resolve_gpu_device_index, start_worker_listener_thread
 from kvcached.utils import CONTIGUOUS_LAYOUT, PAGE_SIZE, get_kvcached_logger, normalize_gpu_device
 from kvcached.vmm_ops import (
     create_kv_tensors,
@@ -64,7 +64,12 @@ def init_kvcached(
         # (broadcast_kv_tensors_created) and fail with ENOENT on the socket path.
         if is_worker and not _is_worker:
             _is_worker = True
-            start_worker_listener_thread(tp_rank, pp_rank)
+            listener_device = _kvcached_device or device
+            start_worker_listener_thread(
+                tp_rank,
+                pp_rank,
+                device_index=resolve_gpu_device_index(listener_device),
+            )
         if async_sched and not _async_sched:
             _async_sched = True
             logger.info("kvcached async scheduler enabled")
@@ -90,7 +95,11 @@ def init_kvcached(
     if is_worker:
         # start the listener thread for kv cache management regardless of TP size
         # because the vLLM EngineCore might need to reach this worker if PP > 1
-        start_worker_listener_thread(tp_rank, pp_rank)
+        start_worker_listener_thread(
+            tp_rank,
+            pp_rank,
+            device_index=resolve_gpu_device_index(device),
+        )
 
 
 def shutdown_kvcached() -> None:

kvcached/tp_ipc_util.py

@@ -7,9 +7,9 @@
 import socket
 import threading
 import uuid
-from typing import Any, Dict, cast
+from typing import Any, Dict, Optional, cast
 
-from kvcached.utils import DEFAULT_IPC_NAME
+from kvcached.utils import DEFAULT_IPC_NAME, normalize_gpu_device
 from kvcached.vmm_ops import kv_tensors_created, map_to_kv_tensors, unmap_from_kv_tensors
 
 
@@ -93,11 +93,29 @@ def recv_msg(sock: socket.socket) -> Message:
     return cast(Message, pickle.loads(data))
 
 
-def start_worker_listener_thread(rank: int, pp_rank: int = 0):
-    """
-    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.
+def resolve_gpu_device_index(device: Optional[str]) -> int:
+    """Resolve an integration device string to the CUDA runtime device index."""
+    import torch
+
+    if device is not None:
+        device_index = torch.device(normalize_gpu_device(device)).index
+        if device_index is not None:
+            return int(device_index)
+    return int(torch.cuda.current_device())
+
+
+def start_worker_listener_thread(
+    rank: int,
+    pp_rank: int = 0,
+    device_index: Optional[int] = None,
+):
+    """Start a thread that listens for messages on the worker socket.
+
+    ``pp_rank`` selects a PP-stage-specific socket directory so concurrent
+    stages do not bind the same path. When ``device_index`` is provided, the
+    listener restores that CUDA device inside the new thread before executing
+    CUDA-backed map or unmap operations because CUDA's current device is
+    thread-local.
     """
     socket_dir = os.path.join(SOCKET_DIR, f"pp{pp_rank}") if pp_rank > 0 else SOCKET_DIR
     os.makedirs(socket_dir, exist_ok=True)
@@ -114,6 +132,11 @@ def start_worker_listener_thread(rank: int, pp_rank: int = 0):
     server_sock.listen()
 
     def listen_loop():
+        if device_index is not None:
+            import torch
+
+            # CUDA's current device is thread-local, so restore the worker device.
+            torch.cuda.set_device(device_index)
         print(f"Worker {rank} IPC listener started at {socket_path}")
         while True:
             conn, _ = server_sock.accept()

tests/manifests/cpu.txt

@@ -13,6 +13,7 @@ tests/test_sglang_allocator_patch.py
 tests/test_shm_info_tracker.py
 tests/test_sleep_manager.py
 tests/test_test_classification.py
+tests/test_tp_listener_device.py
 tests/test_vllm_nixl_compat.py
 tests/test_vllm_pool_exhaustion.py
 tests/test_vllm_tp_world_size.py

tests/test_observability.py

@@ -269,6 +269,7 @@ def __init__(self, num_blocks, block_size, cell_size, num_layers, **kwargs):
     setattr(manager_module, "KVCacheManager", FakeKVCacheManager)
 
     tp_ipc_module = types.ModuleType("kvcached.tp_ipc_util")
+    setattr(tp_ipc_module, "resolve_gpu_device_index", lambda device: 0)
     setattr(tp_ipc_module, "start_worker_listener_thread", lambda *args: None)
 
     utils_module = types.ModuleType("kvcached.utils")
@@ -354,6 +355,7 @@ def __init__(
     setattr(manager_module, "KVCacheManager", FakeKVCacheManager)
 
     tp_ipc_module = types.ModuleType("kvcached.tp_ipc_util")
+    setattr(tp_ipc_module, "resolve_gpu_device_index", lambda device: 0)
     setattr(tp_ipc_module, "start_worker_listener_thread", lambda *args: None)
 
     utils_module = types.ModuleType("kvcached.utils")

tests/test_tp_listener_device.py

@@ -0,0 +1,116 @@
+# SPDX-FileCopyrightText: Copyright contributors to the kvcached project
+# SPDX-License-Identifier: Apache-2.0
+
+import importlib
+import sys
+from importlib.machinery import ModuleSpec
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+
+
+def _mock_torch(monkeypatch, current_device=0):
+    torch = mock.MagicMock()
+    torch.__version__ = "2.6.0"
+    torch.__spec__ = ModuleSpec("torch", loader=None)
+    torch.cuda.current_device.return_value = current_device
+
+    def parse_device(device):
+        value = str(device)
+        index = int(value.split(":", 1)[1]) if ":" in value else None
+        return SimpleNamespace(index=index)
+
+    torch.device.side_effect = parse_device
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    return torch
+
+
+def test_listener_thread_restores_cuda_device(monkeypatch, tmp_path):
+    torch = _mock_torch(monkeypatch)
+    monkeypatch.setitem(sys.modules, "kvcached.vmm_ops", mock.MagicMock())
+
+    import kvcached.tp_ipc_util as tp_ipc_util
+
+    thread_target = None
+
+    class FakeSocket:
+        def bind(self, path):
+            pass
+
+        def listen(self):
+            pass
+
+        def accept(self):
+            raise RuntimeError("stop listener")
+
+    class FakeThread:
+        def __init__(self, target, daemon):
+            nonlocal thread_target
+            thread_target = target
+
+        def start(self):
+            pass
+
+    monkeypatch.setattr(tp_ipc_util, "SOCKET_DIR", str(tmp_path))
+    monkeypatch.setattr(tp_ipc_util.socket, "AF_UNIX", 1, raising=False)
+    monkeypatch.setattr(
+        tp_ipc_util.socket, "socket", lambda *args, **kwargs: FakeSocket()
+    )
+    monkeypatch.setattr(tp_ipc_util.threading, "Thread", FakeThread)
+
+    tp_ipc_util.start_worker_listener_thread(2, 0, device_index=3)
+
+    assert thread_target is not None
+    with pytest.raises(RuntimeError, match="stop listener"):
+        thread_target()
+    torch.cuda.set_device.assert_called_once_with(3)
+
+
+@pytest.mark.parametrize("integration", ["vllm", "sglang"])
+@pytest.mark.parametrize("device", ["cuda:3", "hip:3"])
+def test_worker_listener_uses_explicit_device(monkeypatch, integration, device):
+    _mock_torch(monkeypatch, current_device=0)
+    monkeypatch.setitem(sys.modules, "kvcached.vmm_ops", mock.MagicMock())
+
+    module_name = f"kvcached.integration.{integration}.interfaces"
+    monkeypatch.delitem(sys.modules, module_name, raising=False)
+    interfaces = importlib.import_module(module_name)
+    listener = mock.Mock()
+    monkeypatch.setattr(interfaces, "start_worker_listener_thread", listener)
+
+    kwargs = {
+        "tp_rank": 2,
+        "world_size": 4,
+        "pp_rank": 1,
+        "device": device,
+    }
+    if integration == "vllm":
+        kwargs["is_worker"] = True
+    interfaces.init_kvcached(**kwargs)
+
+    listener.assert_called_once_with(2, 1, device_index=3)
+
+
+def test_vllm_reinit_listener_uses_initialized_device(monkeypatch):
+    _mock_torch(monkeypatch, current_device=0)
+    monkeypatch.setitem(sys.modules, "kvcached.vmm_ops", mock.MagicMock())
+
+    module_name = "kvcached.integration.vllm.interfaces"
+    monkeypatch.delitem(sys.modules, module_name, raising=False)
+    interfaces = importlib.import_module(module_name)
+    listener = mock.Mock()
+    monkeypatch.setattr(interfaces, "start_worker_listener_thread", listener)
+    monkeypatch.setattr(interfaces, "_kvcached_initialized", True)
+    monkeypatch.setattr(interfaces, "_kvcached_device", "cuda:3")
+    monkeypatch.setattr(interfaces, "_is_worker", False)
+
+    interfaces.init_kvcached(
+        tp_rank=2,
+        world_size=4,
+        pp_rank=1,
+        is_worker=True,
+        device="cuda:1",
+    )
+
+    listener.assert_called_once_with(2, 1, device_index=3)