[None][fix] respect KV cache quotas and ignore constraint for non-GPU tiers

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h

@@ -58,6 +58,17 @@ enum class CacheTier : int
     DISK = 2,
 };
 
+inline char const* cacheTierName(CacheTier tier)
+{
+    switch (tier)
+    {
+    case CacheTier::GPU_MEM: return "GPU";
+    case CacheTier::HOST_MEM: return "host";
+    case CacheTier::DISK: return "disk";
+    }
+    return "unknown";
+}
+
 // PageIndexMode — how converted page indices relate to layers within a layer group.
 // Mirrors _common.py::PageIndexMode.
 enum class PageIndexMode : int

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h

@@ -17,10 +17,12 @@
 
 #pragma once
 
+#include "kv_cache_manager_v2/common.h"
 #include "kv_cache_manager_v2/utils/sharedPtr.h"
 
 #include "tensorrt_llm/common/logger.h"
 
+#include <cstddef>
 #include <cuda.h>
 #include <exception>
 #include <stdexcept>
@@ -111,6 +113,18 @@ class AssertionError : public std::logic_error
     }
 };
 
+// A configured tier quota cannot satisfy its minimum storage layout.
+class InsufficientQuotaError : public std::invalid_argument
+{
+public:
+    InsufficientQuotaError(CacheTier tier, size_t configuredQuota, size_t requiredQuota)
+        : std::invalid_argument(std::string(cacheTierName(tier)) + " cache tier quota "
+            + std::to_string(configuredQuota) + " is insufficient for the minimum storage layout (requires at least "
+            + std::to_string(requiredQuota) + ")")
+    {
+    }
+};
+
 // Wraps a CUDA driver API error (CUresult).
 class CuError : public std::runtime_error
 {

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp

@@ -1608,11 +1608,10 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional<size_t> ne
         : lvlStorage.totalQuota();
     auto const minSlots
         = level == kHotLevel ? mMinSlots : TypedVec<PoolGroupIndex, SlotCount>(numPoolGroups(level), SlotCount{1});
-    size_t minQuota = minQuotaForLevel(lvlStorage.slotSizeLists(), lvlStorage.poolSizeGranularity(), minSlots);
+    size_t const minQuota = minQuotaForLevel(lvlStorage.slotSizeLists(), lvlStorage.poolSizeGranularity(), minSlots);
     if (quota < minQuota)
     {
-        throw std::invalid_argument("Quota " + std::to_string(quota)
-            + " is insufficient for min_slots constraints (requires at least " + std::to_string(minQuota) + ")");
+        throw InsufficientQuotaError(lvlStorage.cacheTier(), quota, minQuota);
     }
     auto newNumSlots = lvlStorage.computeSlotCountList(ratioList, minSlots, quota);
 
@@ -1914,11 +1913,17 @@ TypedVec<PoolGroupIndex, SlotCount> StorageManager::computeSlotCountForLevel(Cac
     TypedVec<PoolGroupIndex, TypedVec<PoolIndex, size_t>> const& slotSizeLists,
     TypedVec<PoolGroupIndex, float> const& ratio, TypedVec<PoolGroupIndex, SlotCount> const& minSlots) const
 {
-    CacheTier tier = cacheTierOf(tierConfig);
-    size_t quota = cacheTierQuota(tierConfig);
-    size_t granularity = tier == CacheTier::GPU_MEM ? mGpuPhysMemAllocator->physMemSize()
-                                                    : CacheLevelManager::cacheTierGranularity(tier, quota);
-    quota = std::max(minQuotaForLevel(slotSizeLists, granularity, minSlots), roundUp(quota, granularity));
+    CacheTier const tier = cacheTierOf(tierConfig);
+    size_t const configuredQuota = cacheTierQuota(tierConfig);
+    size_t const granularity = tier == CacheTier::GPU_MEM
+        ? mGpuPhysMemAllocator->physMemSize()
+        : CacheLevelManager::cacheTierGranularity(tier, configuredQuota);
+    size_t const minQuota = minQuotaForLevel(slotSizeLists, granularity, minSlots);
+    size_t const quota = roundUp(configuredQuota, granularity);
+    if (quota < minQuota)
+    {
+        throw InsufficientQuotaError(tier, quota, minQuota);
+    }
     return CacheLevelStorage::ratioToSlotCountList(quota, slotSizeLists, ratio, granularity, minSlots);
 }
 

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp

@@ -792,6 +792,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
     static nb::object sResourceBusyError = nb::exception<kv::ResourceBusyError>(m, "ResourceBusyError");
     static nb::object sOutOfPagesError = nb::exception<kv::OutOfPagesError>(m, "OutOfPagesError");
     static nb::object sCuError = nb::exception<kv::CuError>(m, "CuError");
+    static nb::object sInsufficientQuotaError
+        = nb::exception<kv::InsufficientQuotaError>(m, "InsufficientQuotaError", PyExc_ValueError);
     // Default attribute so the class mirrors the pure-Python CuError surface.
     sCuError.attr("error_code") = nb::none();
 

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -41,6 +41,7 @@
                                                    supports_native_fp8_lora)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.mapping import CpType, Mapping
+from tensorrt_llm.runtime.kv_cache_manager_v2 import InsufficientQuotaError
 
 from ..attention_backend import get_sparse_attn_kv_cache_manager
 from ..hostfunc import set_low_latency_dispatch
@@ -2004,6 +2005,22 @@ def build_managers(self,
                        resources: Dict,
                        estimating_kv_cache: bool = False) -> None:
         """Construct KV caches for model and draft model (if applicable)."""
+        try:
+            self._build_managers(resources, estimating_kv_cache)
+        except InsufficientQuotaError as error:
+            guidance = (
+                "For the GPU tier, increase kv_cache_config.max_gpu_total_bytes "
+                "(if set) or kv_cache_config.free_gpu_memory_fraction, or decrease "
+                "max_batch_size, max_seq_len, or max_num_tokens; for the host/CPU "
+                "tier, increase kv_cache_config.host_cache_size; for the disk tier, "
+                "increase kv_cache_config.disk_cache_size.")
+            raise ValueError(
+                f"Failed to create KV cache manager: {error}. {guidance}"
+            ) from error
+
+    def _build_managers(self,
+                        resources: Dict,
+                        estimating_kv_cache: bool = False) -> None:
         if self._skip_est:
             self.configure_kv_cache_capacity()
         original_max_seq_len = self._max_seq_len

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py

@@ -86,7 +86,12 @@
         KVCacheUpdatedData,
         UniqueToken,
     )
-    from ._exceptions import CuError, OutOfMemoryError, OutOfPagesError  # noqa: F401
+    from ._exceptions import (  # noqa: F401
+        CuError,
+        InsufficientQuotaError,
+        OutOfMemoryError,
+        OutOfPagesError,
+    )
     from ._life_cycle_registry import AttnLifeCycle, LayerGroupId, LifeCycleId  # noqa: F401
     from ._stats import (  # noqa: F401
         _KV_CACHE_ITERATION_STATS_DELTA_FIELDS,
@@ -219,6 +224,7 @@ class _KVCacheManagerConfigFieldSpec:
     _KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple(KVCacheIterationStatsDelta._field_names)
     PlannedDropHandle = _cpp.PlannedDropHandle
     CuError = _cpp.CuError
+    InsufficientQuotaError = _cpp.InsufficientQuotaError
 
     # Symbols added on main that are not yet ported to the C++ backend.
     # TODO(kvCacheManagerV2-cpp): port these and replace the fallbacks.
@@ -311,6 +317,7 @@ def typed_range(*args: int) -> range:
     "GpuCacheTierConfig",
     "HalfOpenRange",
     "HostCacheTierConfig",
+    "InsufficientQuotaError",
     "KVCacheDesc",
     "KVCacheCreatedData",
     "KVCacheEvent",

tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py

@@ -15,6 +15,8 @@
 
 import cuda.bindings.driver as drv
 
+from ._common import CacheTier
+
 
 class OutOfMemoryError(Exception):
     pass
@@ -41,6 +43,26 @@ def __init__(self, message: str) -> None:
         super().__init__(message)
 
 
+class InsufficientQuotaError(ValueError):
+    """A configured cache-tier quota cannot satisfy its minimum layout."""
+
+    pass
+
+
+def _make_insufficient_quota_error(
+    cache_tier: CacheTier, quota: int, min_quota: int
+) -> InsufficientQuotaError:
+    tier_name = {
+        CacheTier.GPU_MEM: "GPU",
+        CacheTier.HOST_MEM: "host",
+        CacheTier.DISK: "disk",
+    }[cache_tier]
+    return InsufficientQuotaError(
+        f"{tier_name} cache tier quota {quota} is insufficient for the minimum storage layout "
+        f"(requires at least {min_quota})"
+    )
+
+
 class CuError(Exception):
     error_code: drv.CUresult
 

tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py

@@ -44,7 +44,7 @@
 from ._copy_engine import CopyTask, batched_copy
 from ._event_manager import KVCacheEventDiff
 from ._eviction_controller import EvictablePage, PerLevelEvictionController
-from ._exceptions import LogicError, OutOfPagesError
+from ._exceptions import LogicError, OutOfPagesError, _make_insufficient_quota_error
 from ._life_cycle_registry import (
     AttnLifeCycle,
     LifeCycleId,
@@ -124,7 +124,9 @@ def cache_tier_granularity(tier: CacheTier, quota: int) -> int:
         match tier:
             case CacheTier.GPU_MEM:
                 page_size = 2 << 20
-                return page_size << min(4, max(0, int(math.log(quota / (page_size * 512), 2))))
+                ratio = quota // (page_size * 512)
+                exponent = 0 if ratio <= 0 else min(4, int(math.log2(ratio)))
+                return page_size << exponent
             case CacheTier.HOST_MEM:
                 return HostCacheLevelStorage.POOL_SIZE_GRANULARITY
             case CacheTier.DISK:
@@ -848,10 +850,7 @@ def adjust_cache_level(
             lvl_storage.slot_size_lists, lvl_storage.pool_size_granularity, min_slots
         )
         if new_quota < min_quota:
-            raise ValueError(
-                f"Quota {new_quota} is insufficient for min_slots constraints "
-                f"(requires at least {min_quota})"
-            )
+            raise _make_insufficient_quota_error(lvl_storage.cache_tier, new_quota, min_quota)
         new_num_slots = lvl_storage.compute_slot_count_list(new_ratio_list, min_slots, new_quota)
         if level != num_cache_levels - 1:
             assert persistent_pages is None, (
@@ -1091,14 +1090,15 @@ def _compute_slot_count_for_level(
     ) -> TypedIndexList[PoolGroupIndex, int]:
         """Compute slot counts for a cache level from its tier config and ratio.
 
-        Applies hot constraint floors or a one-slot structural floor on colder levels.
+        Workload constraints and resume headroom apply only to the hot GPU level.
+        Colder levels use a one-slot structural floor per pool group.
         """
         granularity = CacheLevelManager.cache_tier_granularity(tier_config.tier, tier_config.quota)
         min_slots = self._min_slots_for_level(level)
-        quota = max(
-            self._min_quota_for_level(slot_size_lists, granularity, min_slots),
-            round_up(tier_config.quota, granularity),
-        )
+        min_quota = self._min_quota_for_level(slot_size_lists, granularity, min_slots)
+        quota = round_up(tier_config.quota, granularity)
+        if quota < min_quota:
+            raise _make_insufficient_quota_error(tier_config.tier, quota, min_quota)
         return CacheLevelStorage.ratio_to_slot_count_list(
             quota, slot_size_lists, ratio, granularity, min_slots
         )

tests/unittest/_torch/executor/test_kv_cache_budget_split.py

@@ -23,6 +23,7 @@
 from tensorrt_llm._torch.pyexecutor.config_utils import uses_vswa_kv_cache_layout
 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig
+from tensorrt_llm.runtime.kv_cache_manager_v2 import InsufficientQuotaError
 
 pytestmark = pytest.mark.cpu_only
 
@@ -77,6 +78,25 @@ def _make_creator(
     return c
 
 
+def test_build_managers_reports_quota_guidance() -> None:
+    creator = object.__new__(KvCacheCreator)
+    quota_error = InsufficientQuotaError("Cache tier quota is insufficient")
+    creator._build_managers = Mock(side_effect=quota_error)
+
+    with pytest.raises(ValueError) as raised:
+        creator.build_managers({})
+
+    message = str(raised.value)
+    assert "max_gpu_total_bytes" in message
+    assert "free_gpu_memory_fraction" in message
+    assert "max_batch_size" in message
+    assert "max_seq_len" in message
+    assert "max_num_tokens" in message
+    assert "host_cache_size" in message
+    assert "disk_cache_size" in message
+    assert raised.value.__cause__ is quota_error
+
+
 class TestSplitGpuBudgetForDraft:
     @pytest.mark.parametrize(
         "is_external_drafter",

tests/unittest/_torch/executor/test_kv_cache_manager_v2.py

@@ -31,6 +31,7 @@
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     DEFAULT_BEAM_INDEX,
     BatchDesc,
+    CacheLevel,
     DiskCacheTierConfig,
     GpuCacheTierConfig,
     HostCacheTierConfig,
@@ -368,6 +369,74 @@ def test_host_init_fallback_drops_only_host_tier(tmp_path) -> None:
     ]
 
 
+@pytest.mark.parametrize("max_util_for_resume", [1.0, 0.5])
+def test_host_quota_follows_host_cache_size(max_util_for_resume: float) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+    init_cuda_once()
+    host_cache_size = 1 << 20
+    manager = KVCacheManagerV2(
+        KvCacheConfig(
+            avg_seq_len=512,
+            enable_block_reuse=False,
+            host_cache_size=host_cache_size,
+            max_gpu_total_bytes=16 << 20,
+            max_util_for_resume=max_util_for_resume,
+        ),
+        CacheType.SELF,
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=128,
+        tokens_per_block=32,
+        max_seq_len=1024,
+        max_batch_size=64,
+        max_num_tokens=128,
+        mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
+        dtype=DataType.HALF,
+        vocab_size=4096,
+        enable_stats=False,
+    )
+    try:
+        assert manager.impl.get_quota(CacheLevel(1)) == host_cache_size
+    finally:
+        manager.shutdown()
+
+
+@pytest.mark.parametrize("max_util_for_resume", [1.0, 0.5])
+def test_disk_quota_follows_disk_cache_size(tmp_path, max_util_for_resume: float) -> None:
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+    init_cuda_once()
+    disk_cache_size = 16 << 20
+    manager = KVCacheManagerV2(
+        KvCacheConfig(
+            avg_seq_len=512,
+            disk_cache_path=str(tmp_path),
+            disk_cache_size=disk_cache_size,
+            enable_block_reuse=False,
+            host_cache_size=0,
+            max_gpu_total_bytes=16 << 20,
+            max_util_for_resume=max_util_for_resume,
+        ),
+        CacheType.SELF,
+        num_layers=1,
+        num_kv_heads=1,
+        head_dim=128,
+        tokens_per_block=32,
+        max_seq_len=1024,
+        max_batch_size=64,
+        max_num_tokens=128,
+        mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1),
+        dtype=DataType.HALF,
+        vocab_size=4096,
+        enable_stats=False,
+    )
+    try:
+        assert manager.impl.get_quota(CacheLevel(1)) == disk_cache_size
+    finally:
+        manager.shutdown()
+
+
 def test_extra_tokens_are_in_context_capacity() -> None:
     config = _make_cache_config_for_test(
         KvCacheConfig(avg_seq_len=264),