diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h index 675a6888f23a..dfb4dcdad333 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h +++ b/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 diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h index 428362e61af3..121f39cceb11 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h +++ b/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 #include #include #include @@ -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 { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index c99c0b2b164e..77111e9e8574 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -1608,11 +1608,10 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional ne : lvlStorage.totalQuota(); auto const minSlots = level == kHotLevel ? mMinSlots : TypedVec(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 StorageManager::computeSlotCountForLevel(Cac TypedVec> const& slotSizeLists, TypedVec const& ratio, TypedVec 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); } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index a18f67cff29a..5c650d938ab1 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -792,6 +792,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) static nb::object sResourceBusyError = nb::exception(m, "ResourceBusyError"); static nb::object sOutOfPagesError = nb::exception(m, "OutOfPagesError"); static nb::object sCuError = nb::exception(m, "CuError"); + static nb::object sInsufficientQuotaError + = nb::exception(m, "InsufficientQuotaError", PyExc_ValueError); // Default attribute so the class mirrors the pure-Python CuError surface. sCuError.attr("error_code") = nb::none(); diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 5bedb58e526f..adbd0e64ef29 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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 diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index ac0c604d02c6..a2102e85a21a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/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", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py index 7c037cb36d3d..c841a8ea5b04 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_exceptions.py +++ b/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 diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index b228728ca662..1cb7708a276e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/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 ) diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py index 477e459b3e2c..803c7c7771a7 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py +++ b/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", diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index 3b6df54cb3bf..1886f8c7e2eb 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/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), diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 77a9c1cffaae..d842a7349503 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -44,6 +44,7 @@ DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, + InsufficientQuotaError, KVCacheDesc, KVCacheManager, KVCacheManagerConfig, @@ -98,6 +99,7 @@ DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, + InsufficientQuotaError, KVCacheDesc, KVCacheManager, KVCacheManagerConfig, @@ -3348,11 +3350,13 @@ def test_constraint_reserves_resume_headroom(self): num_requests = 32 constraint = BatchDesc(kv_caches=[KVCacheDesc(capacity=1, history_length=0)] * num_requests) granularity = 2 << 20 - gpu_quota = round_up(num_requests * self.PG0_SLOT_SIZE, granularity) + round_up( - num_requests * self.PG1_SLOT_SIZE, granularity + max_util_for_resume = 0.95 + min_slots = math.ceil(num_requests / max_util_for_resume) + gpu_quota = round_up(min_slots * self.PG0_SLOT_SIZE, granularity) + round_up( + min_slots * self.PG1_SLOT_SIZE, granularity ) cfg = self._make_config(gpu_quota=gpu_quota, constraints=[constraint]) - cfg.max_util_for_resume = 0.95 + cfg.max_util_for_resume = max_util_for_resume manager = KVCacheManager(cfg) stream_holder = CachedCudaStream() stream = cast(CudaStream, stream_holder.handle) @@ -3368,6 +3372,35 @@ def test_constraint_reserves_resume_headroom(self): kv_cache.close() manager.shutdown() + def test_gpu_quota_below_constraint_minimum_raises(self): + """A configured GPU quota below its constraint floor is rejected.""" + num_requests = 32 + constraint = BatchDesc(kv_caches=[KVCacheDesc(capacity=1, history_length=0)] * num_requests) + granularity = 2 << 20 + gpu_quota = round_up(num_requests * self.PG0_SLOT_SIZE, granularity) + round_up( + num_requests * self.PG1_SLOT_SIZE, granularity + ) + cfg = self._make_config(gpu_quota=gpu_quota, constraints=[constraint]) + cfg.max_util_for_resume = 0.95 + + with self.assertRaisesRegex( + InsufficientQuotaError, + rf"^GPU cache tier quota {gpu_quota} is insufficient " + r"for the minimum storage layout \(requires at least \d+\)$", + ): + KVCacheManager(cfg) + + def test_zero_gpu_quota_raises(self): + """A zero configured GPU quota is rejected.""" + cfg = self._make_config(gpu_quota=0) + + with self.assertRaisesRegex( + InsufficientQuotaError, + r"^GPU cache tier quota 0 is insufficient " + r"for the minimum storage layout \(requires at least \d+\)$", + ): + KVCacheManager(cfg) + def test_constraint_floor_overrides_infeasible_initial_pool_ratio(self): """A constraint's feasibility floor overrides an infeasible initial_pool_ratio.