feat(vllm): avoid CUDA context in EngineCore

examples/01_simple_two_models/README.md

@@ -22,6 +22,19 @@ vllm serve "${MODEL}" \
   --port "${PORT}"
 ```
 
+To avoid creating an additional CUDA context in the vLLM EngineCore process,
+enable the worker-owned physical-admission path explicitly:
+
+```bash
+export KVCACHED_ENGINECORE_NO_CUDA=true
+```
+
+This mode is disabled by default. When enabled, EngineCore manages logical KV
+capacity while the existing CUDA-owning TP workers perform final physical page
+mapping. It does not add a separate synchronous memory-info query to the
+scheduler path, and page preallocation is disabled for the control-only
+EngineCore. Unset the variable to restore the original behavior for A/B tests.
+
 For SGLang:
 
 ```bash

kvcached/integration/vllm/interfaces.py

@@ -27,7 +27,8 @@
 logger = get_kvcached_logger()
 
 _kvcached_initialized: bool = False
-_kvcached_device = None
+_kvcached_gpu_initialized: bool = False
+_kvcached_device: Optional[str] = None
 _async_sched = False
 _world_size: int = 1
 _pp_rank: int = 0
@@ -48,37 +49,71 @@ def get_world_size() -> int:
     return _world_size
 
 
+def _cuda_device_index(device: str) -> int:
+    normalized = normalize_gpu_device(device)
+    device_type, separator, index = normalized.partition(":")
+    if device_type.lower() != "cuda":
+        raise ValueError(f"Expected a CUDA device, got {device}")
+    if separator and index:
+        return int(index)
+    return int(torch.cuda.current_device())
+
+
 def init_kvcached(
     tp_rank: int = 0,
     world_size: int = 1,
     pp_rank: int = 0,
     is_worker: bool = False,
     device: Optional[str] = None,
     async_sched: bool = False,
+    control_only: bool = False,
 ) -> None:
-    global _kvcached_initialized, _kvcached_device, _world_size, _async_sched, _pp_rank, _is_worker
+    global _kvcached_initialized, _kvcached_gpu_initialized, _kvcached_device
+    global _world_size, _async_sched, _pp_rank, _is_worker
+    if control_only and is_worker:
+        raise ValueError("control_only mode is only valid in the engine process")
     if _kvcached_initialized:
         # EngineCore call init_kvcached(is_worker=False) first. When TP=1 GPUModelRunner
         # then calls init_kvcached(is_worker=True) in the same process; without this branch
         # the early return would leave _is_worker False, so KVCacheManager would try Unix IPC
         # (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)
+            if not _kvcached_gpu_initialized:
+                if device is None:
+                    device = f"cuda:{torch.cuda.current_device()}"
+                device = normalize_gpu_device(device)
+                _init_kvcached_impl(device, PAGE_SIZE, _contiguous_layout)
+                _kvcached_gpu_initialized = True
+                _kvcached_device = device
+            assert _kvcached_device is not None
+            start_worker_listener_thread(
+                tp_rank, pp_rank,
+                device_index=_cuda_device_index(_kvcached_device))
         if async_sched and not _async_sched:
             _async_sched = True
             logger.info("kvcached async scheduler enabled")
         _pp_rank = pp_rank
         _world_size = world_size
         return
 
+    if control_only:
+        _kvcached_initialized = True
+        _world_size = world_size
+        _pp_rank = pp_rank
+        _async_sched = async_sched
+        _is_worker = False
+        logger.info("kvcached initialized in control-only mode")
+        return
+
     if device is None:
         device = f"cuda:{torch.cuda.current_device()}"
-    device = normalize_gpu_device(device)
+    normalized_device = normalize_gpu_device(device)
 
-    _init_kvcached_impl(device, PAGE_SIZE, _contiguous_layout)
+    _init_kvcached_impl(normalized_device, PAGE_SIZE, _contiguous_layout)
     _kvcached_initialized = True
-    _kvcached_device = device
+    _kvcached_gpu_initialized = True
+    _kvcached_device = normalized_device
     _world_size = world_size
     _pp_rank = pp_rank
     _async_sched = async_sched
@@ -90,20 +125,28 @@ 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=_cuda_device_index(normalized_device))
 
 
 def shutdown_kvcached() -> None:
-    global _kvcached_initialized, _kvcached_device, _async_sched
+    global _kvcached_initialized, _kvcached_gpu_initialized, _kvcached_device
+    global _async_sched, _world_size, _pp_rank, _is_worker
     if not _kvcached_initialized:
         clear_registered_kv_cache_pools(integration="vllm")
         return
 
-    _shutdown_kvcached_impl()
+    if _kvcached_gpu_initialized:
+        _shutdown_kvcached_impl()
     clear_registered_kv_cache_pools(integration="vllm")
     _kvcached_initialized = False
+    _kvcached_gpu_initialized = False
     _kvcached_device = None
     _async_sched = False
+    _world_size = 1
+    _pp_rank = 0
+    _is_worker = False
 
 
 def build_kv_views(
@@ -592,6 +635,7 @@ def get_kv_cache_manager(
         group_id=group_id,
         reserve_null_block=True,
         pool_name=pool_name,
+        worker_physical_admission=not _kvcached_gpu_initialized,
     )
     register_kv_cache_pool(
         manager,

kvcached/integration/vllm/patches.py

@@ -977,11 +977,14 @@ def _patched_engine_init(self, vllm_config, *args: Any, **kwargs: Any):
                 # within a single PP stage's TP group (w0.sock … w(tp-1).sock).
                 # Each PP stage manages its own KV memory independently, so
                 # cross-stage IPC is neither needed nor correct.
+                from kvcached.utils import ENGINECORE_NO_CUDA
+
                 init_kvcached(
                     tp_rank=0,
                     world_size=vllm_config.parallel_config.tensor_parallel_size,
                     is_worker=False,
                     async_sched=_should_enable_async_sched(vllm_config),
+                    control_only=ENGINECORE_NO_CUDA,
                 )
             return original_init(self, vllm_config, *args, **kwargs)
 
@@ -1074,11 +1077,14 @@ def _setup_kvcached_coordinator(self) -> None:
             # Use tp_size (not TP*PP global world size) for the KVCacheManager world_size.
             # Each PP stage manages its own KV tensors independently. The IPC sockets
             # are registered per TP rank within each stage (w0.sock … w(tp_size-1).sock).
+            from kvcached.utils import ENGINECORE_NO_CUDA
+
             kvi.init_kvcached(
                 tp_rank=0,
                 world_size=tp_size,
                 is_worker=False,
                 async_sched=_should_enable_async_sched(getattr(self, "vllm_config", None)),
+                control_only=ENGINECORE_NO_CUDA,
             )
 
             # Import ElasticBlockPool from the patched module

kvcached/kv_cache_manager.py

@@ -69,6 +69,7 @@ def __init__(
         num_kv_buffers: int = 2,
         group_id: int = 0,
         pool_name: Optional[str] = None,
+        worker_physical_admission: bool = False,
     ):
         """
         Args:
@@ -87,6 +88,9 @@ def __init__(
                 Different groups have independent FTensors and page spaces.
             pool_name: Stable, low-cardinality name assigned by the engine
                 integration when this pool is created.
+            worker_physical_admission: Report logical capacity to the scheduler
+                and let the CUDA-owning workers perform final physical
+                admission while mapping pages.
         """
         self.num_blocks = num_blocks
         self.block_mem_size = block_size * cell_size
@@ -95,6 +99,7 @@ def __init__(
         self.reserve_null_block = reserve_null_block
         self.group_id = group_id
         self._pool_name = pool_name
+        self.worker_physical_admission = worker_physical_admission
 
         # The physical page size used by kvcached page allocator.
         self.page_size = PAGE_SIZE
@@ -129,7 +134,10 @@ def __init__(
             pp_rank=self.pp_rank,
             async_sched=async_sched,
             contiguous_layout=CONTIGUOUS_LAYOUT,
-            enable_page_prealloc=PAGE_PREALLOC_ENABLED,
+            # A control-only process must not let the C++ preallocation thread
+            # enter Python to reach remote workers.
+            enable_page_prealloc=(PAGE_PREALLOC_ENABLED and
+                                  not self.worker_physical_admission),
             num_kv_buffers=self.num_kv_buffers,
             group_id=self.group_id,
             ipc_name=DEFAULT_IPC_NAME,
@@ -657,9 +665,16 @@ def available_size(self) -> int:
             blocks_from_free_pages = 0
         else:
             virtual_free_pages = self.page_allocator.get_num_free_pages()
-            physical_free_pages = self.page_allocator.get_avail_physical_pages(
-            ) + self.page_allocator.get_num_reserved_pages()
-            free_pages = min(virtual_free_pages, physical_free_pages)
+            if not getattr(self, "worker_physical_admission", False):
+                physical_free_pages = (
+                    self.page_allocator.get_avail_physical_pages() +
+                    self.page_allocator.get_num_reserved_pages())
+                free_pages = min(virtual_free_pages, physical_free_pages)
+            else:
+                # EngineCore owns only logical KV state in control-only mode.
+                # The workers that own the CUDA contexts perform the final
+                # physical admission as part of transactional page mapping.
+                free_pages = virtual_free_pages
             blocks_from_free_pages = free_pages * InternalPage.get_num_blocks(
                 self.page_size, self.block_mem_size)
         return avail_blocks + blocks_from_free_pages

kvcached/tp_ipc_util.py

@@ -7,7 +7,7 @@
 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.vmm_ops import kv_tensors_created, map_to_kv_tensors, unmap_from_kv_tensors
@@ -93,7 +93,19 @@ def recv_msg(sock: socket.socket) -> Message:
     return cast(Message, pickle.loads(data))
 
 
-def start_worker_listener_thread(rank: int, pp_rank: int = 0):
+def _set_listener_device(device_index: Optional[int]) -> None:
+    if device_index is None:
+        return
+
+    import torch
+
+    # CUDA's current device is thread-local. Restore the worker's initialized
+    # device before serving CUDA-backed map operations.
+    torch.cuda.set_device(device_index)
+
+
+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 is used to create a PP-stage-specific subdirectory so that
@@ -114,6 +126,7 @@ def start_worker_listener_thread(rank: int, pp_rank: int = 0):
     server_sock.listen()
 
     def listen_loop():
+        _set_listener_device(device_index)
         print(f"Worker {rank} IPC listener started at {socket_path}")
         while True:
             conn, _ = server_sock.accept()
@@ -137,7 +150,10 @@ def listen_loop():
                     })
             except Exception as e:
                 print(f"Worker {rank} error processing message: {e}")
-                send_msg(conn, {"status": "error", "message": str(e)})
+                try:
+                    send_msg(conn, {"status": "error", "message": str(e)})
+                except (BrokenPipeError, ConnectionError, OSError):
+                    pass
             finally:
                 conn.close()
 

kvcached/utils.py

@@ -136,6 +136,8 @@ def _get_page_size() -> int:
 GPU_UTILIZATION = float(os.getenv("KVCACHED_GPU_UTILIZATION", "0.95"))
 PAGE_PREALLOC_ENABLED = os.getenv("KVCACHED_PAGE_PREALLOC_ENABLED",
                                   "true").lower() == "true"
+ENGINECORE_NO_CUDA = os.getenv("KVCACHED_ENGINECORE_NO_CUDA",
+                               "false").lower() == "true"
 MIN_RESERVED_PAGES = int(os.getenv("KVCACHED_MIN_RESERVED_PAGES", "5"))
 MAX_RESERVED_PAGES = int(os.getenv("KVCACHED_MAX_RESERVED_PAGES", "10"))
 MAX_CACHED_BLOCKS = int(os.getenv("KVCACHED_MAX_CACHED_BLOCKS", "1000"))

tests/manifests/cpu.txt

@@ -1,5 +1,6 @@
 tests/test_alloc_rollback.py
 tests/test_bestfit_page_selection.py
+tests/test_enginecore_no_cuda.py
 tests/test_get_max_cached_blocks.py
 tests/test_ipc_name.py
 tests/test_ipc_timeout.py