fix(vllm): report an exhausted KV pool as a scheduling miss, not a crash

kvcached/integration/vllm/autopatch.py

@@ -17,6 +17,7 @@
     GPUModelRunnerPatch,
     GPUWorkerPatch,
     KVCacheCoordinatorPatch,
+    KVCacheManagerAllocateSlotsPatch,
     KVCacheManagerPatch,
     TritonAttentionPatch,
 )
@@ -47,6 +48,7 @@ def _patch_vllm(_vllm: types.ModuleType) -> None:
             (GPUWorkerPatch(), VLLM_ALL_RANGE),
             (KVCacheCoordinatorPatch(), VLLM_V9_PLUS_RANGE),
             (KVCacheManagerPatch(), VLLM_V8_RANGE),
+            (KVCacheManagerAllocateSlotsPatch(), VLLM_ALL_RANGE),
             (TritonAttentionPatch(), VLLM_V9_PLUS_RANGE),
         ]
     )

kvcached/integration/vllm/patches.py

@@ -14,7 +14,7 @@
 
 from kvcached.integration.patch_base import BasePatch, enable_kvcached
 from kvcached.integration.version_utils import VersionAwarePatch, VersionRange, version_range
-from kvcached.utils import KVCachedConfigError, get_kvcached_logger
+from kvcached.utils import KVCachedConfigError, KVCachePoolExhausted, get_kvcached_logger
 
 if TYPE_CHECKING:
     # These types are imported from vLLM at runtime via getattr()
@@ -802,7 +802,12 @@ def get_new_blocks(
                         break
 
                 if block_ids is None:
-                    raise ValueError(
+                    # Transient, not a defect: a colocated engine took the last
+                    # physical pages. KVCacheManagerAllocateSlotsPatch turns
+                    # this into the scheduler's own "cannot allocate now"
+                    # signal, so keep it a distinct type the patch can catch
+                    # without also swallowing real contract violations.
+                    raise KVCachePoolExhausted(
                         "Unable to allocate KV cache blocks from physical pool; "
                         f"requested={num_blocks}, available={self.kv_cache_manager.available_size()}"
                     )
@@ -1893,6 +1898,79 @@ def _patched_init_device(self, *args: Any, **kwargs: Any):  # type: ignore[no-se
         return True
 
 
+class KVCacheManagerAllocateSlotsPatch(VersionAwarePatch, BasePatch):
+    """Report an exhausted physical KV pool the way vLLM's scheduler expects.
+
+    vLLM's own block pool can raise from `get_new_blocks()` because its free
+    count is process-local and authoritative: if the count says the blocks are
+    there, the allocation cannot fail, so the raise is an invariant guard that
+    never fires. Under kvcached the same count reads device-wide state shared
+    with colocated engines, so it is a snapshot rather than a reservation, and
+    a peer can take the last pages before they are claimed. The guard becomes
+    reachable.
+
+    The scheduler already handles this exact situation -- `allocate_slots()`
+    returning None means "not now", and it preempts a running request and
+    retries on the next step, which incidentally releases physical pages back
+    to the shared pool. What it does not handle is an exception: `schedule()`
+    catches nothing and EngineCore's own handler wraps only `execute_model`,
+    so the exception terminates the engine and every in-flight request with
+    it.
+
+    Translate only `KVCachePoolExhausted`. A plain ValueError from the pool
+    (asking for more blocks than were just reported free) is a contract
+    violation and must stay fail-loud.
+    """
+
+    library = "vllm"
+    target_module = "vllm.v1.core.kv_cache_manager"
+    target_class = "KVCacheManager"
+    patch_name = "allocate_slots"
+
+    def apply(self, kvcache_manager_mod: types.ModuleType) -> bool:
+        if not self.initialize_version_info():
+            return False
+        return self.patch_allocate_slots(kvcache_manager_mod)
+
+    @version_range(VLLM_ALL_RANGE)
+    def patch_allocate_slots(self, kvcache_manager_mod: types.ModuleType) -> bool:
+        KVCacheManager = self._get_target_class(kvcache_manager_mod)
+        if KVCacheManager is None:
+            return False
+
+        original_allocate_slots = getattr(KVCacheManager, "allocate_slots", None)
+        if original_allocate_slots is None:
+            self.logger.warning(
+                "KVCacheManager.allocate_slots was not found; an exhausted "
+                "physical KV pool will terminate EngineCore")
+            return False
+        if self._is_already_patched(original_allocate_slots, "allocate_slots"):
+            self.logger.debug("KVCacheManager.allocate_slots already patched")
+            return True
+
+        logger = self.logger
+
+        def _patched_allocate_slots(self, *args: Any, **kwargs: Any) -> Any:
+            if not enable_kvcached():
+                return original_allocate_slots(self, *args, **kwargs)
+            try:
+                return original_allocate_slots(self, *args, **kwargs)
+            except KVCachePoolExhausted as exhausted:
+                # None is the scheduler's own "cannot schedule this request
+                # now" path. Partially allocated blocks are released when the
+                # scheduler preempts or frees the request, so returning here
+                # does not strand them.
+                logger.warning(
+                    "Shared physical KV pool is exhausted; reporting a "
+                    "scheduling miss so the engine can preempt and retry: %s",
+                    exhausted)
+                return None
+
+        self._mark_as_patched(_patched_allocate_slots, "allocate_slots")
+        KVCacheManager.allocate_slots = _patched_allocate_slots  # type: ignore[assignment]
+        return True
+
+
 class TritonAttentionPatch(VersionAwarePatch, BasePatch):
     """Build the per-token-head scale views from the KV tensor, not raw storage.
 

kvcached/utils.py

@@ -12,6 +12,22 @@ class KVCachedConfigError(RuntimeError):
     abort startup instead of silently falling back to non-kvcached behavior."""
 
 
+class KVCachePoolExhausted(ValueError):
+    """Raised when the shared physical KV pool cannot back an allocation.
+
+    This is a transient condition, not a defect: colocated engines share one
+    physical pool, so a peer can take the last pages between the moment
+    availability is observed and the moment the pages are claimed. Serving
+    engines already know how to respond -- free something and try again -- so
+    integrations translate this into whatever "cannot allocate right now"
+    signal their engine understands, rather than letting it terminate the
+    process.
+
+    It subclasses ValueError so callers written against the pre-existing
+    behavior keep working.
+    """
+
+
 def _sanitize_segment(segment: str) -> str:
     """Sanitize a segment to safe characters for SHM names."""
     allowed = []

tests/manifests/cpu.txt

@@ -14,4 +14,5 @@ tests/test_shm_info_tracker.py
 tests/test_sleep_manager.py
 tests/test_test_classification.py
 tests/test_vllm_nixl_compat.py
+tests/test_vllm_pool_exhaustion.py
 tests/test_vllm_tp_world_size.py

tests/test_prefix_cache.py

@@ -156,8 +156,9 @@ def pool_factory(monkeypatch):
     Returns (pool, manager) so tests can inspect both.
     """
 
-    def _make(num_blocks: int = 100, enable_caching: bool = True):
-        manager = MockKVCacheManager(num_blocks)
+    def _make(num_blocks: int = 100, enable_caching: bool = True, manager=None):
+        if manager is None:
+            manager = MockKVCacheManager(num_blocks)
 
         kv_cache_utils = types.ModuleType("vllm.v1.core.kv_cache_utils")
         setattr(
@@ -724,3 +725,32 @@ def test_get_usage(self, pool_factory):
         pool.get_new_blocks(50)
         # 50+1(null) allocated, 0 evictable -> 49 free from kvcached
         assert pool.get_usage() == pytest.approx(0.51)
+
+
+class DrainedPoolManager(MockKVCacheManager):
+    """Leave the pool in the state a colocated peer produces.
+
+    ``available_size()`` reads device-wide free memory, so it is a snapshot of
+    state shared with every colocated engine, not a reservation: a peer can
+    take the last pages between the pool reading it and the pages being
+    claimed.
+    """
+
+    def available_size(self) -> int:
+        return 1000
+
+    def alloc(self, n: int):
+        return None
+
+
+def test_exhaustion_raises_the_type_the_integration_translates(pool_factory):
+    """The other half of this fix lives in KVCacheManagerAllocateSlotsPatch,
+    which turns exactly this exception into a scheduling miss. Widening it back
+    to a plain ValueError would silently restore the EngineCore crash."""
+    from kvcached.utils import KVCachePoolExhausted
+
+    pool, _ = pool_factory(manager=DrainedPoolManager(100))
+
+    with pytest.raises(KVCachePoolExhausted,
+                       match="Unable to allocate KV cache blocks"):
+        pool.get_new_blocks(4)

tests/test_vllm_pool_exhaustion.py

@@ -0,0 +1,93 @@
+# SPDX-FileCopyrightText: Copyright contributors to the kvcached project
+# SPDX-License-Identifier: Apache-2.0
+"""An exhausted shared KV pool must reach vLLM as a scheduling miss.
+
+vLLM's scheduler has one channel for "this request cannot be served right
+now": ``allocate_slots()`` returns None, and it preempts a running request and
+retries on the next step. It has no channel for an exception -- ``schedule()``
+contains no exception handler, and EngineCore's own handler wraps only
+``execute_model`` -- so an exception from the block pool terminates the engine
+along with every in-flight request.
+
+Under kvcached the pool can legitimately fail to back an allocation: colocated
+engines share one physical pool, and a peer can take the last pages between the
+moment availability is observed and the moment they are claimed. These tests
+pin the translation, and pin that it stays narrow.
+"""
+from __future__ import annotations
+
+import importlib
+import sys
+import types
+from typing import Any
+from unittest import mock
+
+import pytest
+
+
+@pytest.fixture
+def vllm_patches(monkeypatch):
+    torch = mock.MagicMock()
+    torch.__version__ = "2.6.0"
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.setitem(sys.modules, "torch.cuda", torch.cuda)
+    monkeypatch.setitem(sys.modules, "torch.utils", torch.utils)
+    monkeypatch.setitem(sys.modules, "torch.utils.cpp_extension",
+                        torch.utils.cpp_extension)
+    monkeypatch.setitem(sys.modules, "posix_ipc", mock.MagicMock())
+    monkeypatch.setitem(sys.modules, "kvcached.vmm_ops", mock.MagicMock())
+    monkeypatch.delitem(sys.modules, "kvcached.integration.vllm.patches",
+                        raising=False)
+    return importlib.import_module("kvcached.integration.vllm.patches")
+
+
+def _module_with_manager(raises: BaseException | None):
+    """Build a stand-in ``vllm.v1.core.kv_cache_manager`` module.
+
+    ``allocate_slots`` either raises what the block pool would raise, or
+    returns a sentinel so the success path stays observable.
+    """
+
+    class KVCacheManager:
+        def allocate_slots(self, *args: Any, **kwargs: Any) -> Any:
+            if raises is not None:
+                raise raises
+            return ("blocks", args, kwargs)
+
+    module = types.ModuleType("vllm.v1.core.kv_cache_manager")
+    module.KVCacheManager = KVCacheManager  # type: ignore[attr-defined]
+    return module
+
+
+def _apply(vllm_patches, monkeypatch, module, *, kvcached_enabled=True):
+    patch = vllm_patches.KVCacheManagerAllocateSlotsPatch()
+    monkeypatch.setattr(patch, "initialize_version_info", lambda: True)
+    monkeypatch.setattr(vllm_patches, "enable_kvcached",
+                        lambda: kvcached_enabled)
+    assert patch.patch_allocate_slots(module) is True
+    return module.KVCacheManager()
+
+
+def test_pool_exhaustion_becomes_a_scheduling_miss(vllm_patches, monkeypatch):
+    from kvcached.utils import KVCachePoolExhausted
+
+    manager = _apply(
+        vllm_patches, monkeypatch,
+        _module_with_manager(KVCachePoolExhausted("physical pool empty")))
+
+    assert manager.allocate_slots("request", 8) is None
+
+
+def test_contract_violations_still_terminate(vllm_patches, monkeypatch):
+    """Only exhaustion is transient.
+
+    Asking the pool for more blocks than it just reported free is a defect in
+    the caller. Downgrading it to a scheduling miss would turn a loud bug into
+    a request that is quietly never scheduled.
+    """
+    manager = _apply(
+        vllm_patches, monkeypatch,
+        _module_with_manager(ValueError("Cannot get 999 free blocks")))
+
+    with pytest.raises(ValueError, match="Cannot get 999 free blocks"):
+        manager.allocate_slots("request", 8)