diff --git a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py index c7ca1891b642..4ef6569aff41 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py @@ -14,16 +14,13 @@ # limitations under the License. from collections import defaultdict +from dataclasses import replace from typing import Dict, List, Optional, Tuple import torch from tensorrt_llm._torch.pyexecutor import llm_request -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( - GPU_LEVEL, - BlockReusePolicy, - KVCacheManagerV2, -) +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import GPU_LEVEL, KVCacheManagerV2 from tensorrt_llm._utils import ( TensorWrapper, convert_to_torch_tensor, @@ -42,13 +39,10 @@ BatchDesc, BufferConfig, DataRole, - GpuCacheTierConfig, - HostCacheTierConfig, KVCacheDesc, LayerId, PageIndexMode, ScratchDesc, - SwaScratchReuseConfig, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX @@ -254,9 +248,10 @@ def __init__( self._init_indexer_dtype(sparse_attn_config) - # _build_cache_config() needs them to build constraints + # _build_cache_config() needs these values to build the typical step. self._max_input_len = max_input_len self._max_num_tokens = max_num_tokens + self._avg_seq_len = kv_cache_config.avg_seq_len # General initialization super().__init__( @@ -768,16 +763,9 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: return float("inf") return self._max_num_tokens + (quota - context_limit_quota) / generation_size_per_token - def _build_cache_config( - self, - kv_cache_config: KvCacheConfig, - *, - tokens_per_block: int, - vocab_size: int | None, - cache_tiers: List[GpuCacheTierConfig | HostCacheTierConfig], - ) -> KVCacheManagerConfigPy: + def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """ - Create the cache manager config for DeepSeek-V4. + Add DeepSeek-V4 layers and a typical step to the cache config. """ layers: List[AttentionLayerConfig] = [] layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {} @@ -864,19 +852,14 @@ def _add_layer( # number of layers in the KVCacheManagerPy self._num_manager_layers = len(layers) - # Build constraints and typical_step for better pool ratio. + # Build a typical step for the DeepSeek-V4 pool ratio. max_batch_size = self.max_batch_size max_seq_len = self.max_seq_len max_num_tokens = self._max_num_tokens max_draft_len = self._max_draft_len typical_step = None - constraints = [] - if kv_cache_config.pool_ratio is None: - typical_seq_len = ( - kv_cache_config.avg_seq_len - if kv_cache_config.avg_seq_len is not None - else max_seq_len - ) + if config.initial_pool_ratio is None: + typical_seq_len = self._avg_seq_len if self._avg_seq_len is not None else max_seq_len if typical_seq_len > max_seq_len: raise ValueError( f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or " @@ -903,50 +886,10 @@ def _add_layer( * (max_batch_size - 1), ) - # Constraint 1: cuda graph generation warmup — one decode request that has - # accumulated to the tail of max_seq_len. Using history_length=max_seq_len-1 - # (instead of 0) lets SWA / SSM pools collapse to their windowed working set, - # while full-cache pools still need max_seq_len/tokens_per_block blocks - # because they don't age. - constraints.append( - BatchDesc([KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1)]) - ) - - # Constraint 2: general / chunked-prefill warmup — one fresh context request - # at max_num_tokens (the per-iteration token budget). - if max_num_tokens is not None: - constraints.append( - BatchDesc( - [ - KVCacheDesc( - capacity=max_num_tokens + self.num_extra_kv_tokens, history_length=0 - ) - ] - ) - ) - - scratch_reuse_config = None - if self.enable_swa_scratch_reuse: - # Context requests will allocate num_extra_kv_tokens tokens for spec decoding. - # Cache manager should not take them into account when calculating scratch range. - # Therefore set max_rewind_len to num_extra_kv_tokens. - scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens) - - return KVCacheManagerConfigPy( - tokens_per_block=tokens_per_block, - cache_tiers=cache_tiers, - max_util_for_resume=kv_cache_config.max_util_for_resume, - enable_partial_reuse=kv_cache_config.enable_partial_reuse, - swa_scratch_reuse=scratch_reuse_config, - commit_min_snapshot=( - kv_cache_config.enable_block_reuse - and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE - ), + return replace( + config, layers=layers, typical_step=typical_step, - constraints=constraints, - enable_stats=self.enable_stats, - initial_pool_ratio=kv_cache_config.pool_ratio, ) def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None: diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py index b9badca6ea35..a46f812d7a63 100644 --- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py @@ -168,7 +168,7 @@ def __init__( disable_index_value_layer_ids = list(sparse_layer_ids) # Must be set BEFORE super().__init__ — the base - # ``_build_cache_config`` invokes ``_extra_buffers_per_layer`` + # ``_build_base_config`` invokes ``_extra_buffers_per_layer`` # which reads these attributes. self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids) self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids) @@ -198,7 +198,7 @@ def _extra_buffers_per_layer(self, *, tokens_per_block): ``size`` is bytes per **block**: ``1 * sparse_index_dim * elem_bytes * tokens_per_block``. Keyed by **local** layer id — - the base ``_build_cache_config`` iterates local ids, so keying + the base ``_build_base_config`` iterates local ids, so keying by global ids would silently skip registration on non-trivial PP ranks. """ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 46d3c55a0771..80b2e03f38bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -17,6 +17,7 @@ import os import sys from collections import OrderedDict, defaultdict +from dataclasses import replace from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union import numpy as np @@ -45,6 +46,7 @@ GPU_LEVEL, AttentionLayerConfig, AttnLifeCycle, + BatchDesc, BufferConfig, CacheLevel, CacheTierConfig, @@ -53,6 +55,7 @@ DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, + KVCacheDesc, KVCacheEventManager, KVCacheIterationStatsDelta, LayerId, @@ -121,8 +124,8 @@ class Role: # Sparse-attention per-layer index-K cache (MiniMax-M3 and similar # sparse-block-selection backends). Registered as a native V2 # BufferConfig on sparse layers via the extra_buffers_per_layer hook on - # _build_cache_config, so allocation, free, slot reuse, and prefix - # reuse share the lifecycle of the main K/V buffers for the same layer. + # _build_base_config, so allocation, free, slot reuse, and prefix reuse + # share the lifecycle of the main K/V buffers for the same layer. INDEX_KEY = DataRole("index_key") ALL = DataRole("all") @@ -797,6 +800,7 @@ def append_to_kv_heads_per_layer( self.is_vswa = len(set(self.max_attention_window_vec)) > 1 + max_util_for_resume = kv_cache_config.max_util_for_resume quota = sys.maxsize if ( kv_cache_config.max_gpu_total_bytes is not None @@ -808,7 +812,7 @@ def append_to_kv_heads_per_layer( quota_from_max_tokens = int( math.ceil( self._get_quota_from_max_tokens(kv_cache_config.max_tokens) - / kv_cache_config.max_util_for_resume + / max_util_for_resume ) ) quota = min(quota, quota_from_max_tokens) @@ -822,13 +826,13 @@ def append_to_kv_heads_per_layer( "Quota not set. Check kv_cache_config.max_tokens or kv_cache_config.max_gpu_total_bytes" ) - # Sync KV cache token capacity across ranks so all ranks allocate - # the same number of tokens and the scheduler produces identical - # batches. Normalize to token count before the allreduce because - # bytes_per_token varies across PP ranks (different local layers). + # Sync resumable token capacity across ranks so the scheduler produces + # identical batches. Normalize to tokens because cache costs vary + # across PP ranks, including fixed per-rank costs. if mapping.world_size > 1: dist = Distributed.get(mapping) - max_tokens = self._get_max_tokens_from_quota(quota) + resumable_quota = int(quota * max_util_for_resume) + max_tokens = self._get_max_tokens_from_quota(resumable_quota) max_tokens = dist.allreduce(max_tokens, op=ReduceOp.MIN) # inf max_tokens means all layers are SWA and every rank quota can # fit all SWA fixed cache. @@ -837,7 +841,10 @@ def append_to_kv_heads_per_layer( # token↔quota round-trip is not identity when SWA layers # dominate (full_attn_size_per_token==0), so clamp to guard # against a bogus inflation (nvbugs/6418103). - quota = min(quota, self._get_quota_from_max_tokens(max_tokens)) + synced_quota = int( + math.ceil(self._get_quota_from_max_tokens(max_tokens) / max_util_for_resume) + ) + quota = min(quota, synced_quota) logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") @@ -891,12 +898,12 @@ def append_to_kv_heads_per_layer( self.vocab_size = vocab_size - config = self._build_cache_config( + config = self._build_base_config( kv_cache_config, tokens_per_block=tokens_per_block, - vocab_size=vocab_size, cache_tiers=cache_tiers, ) + config = self._build_cache_config(config) self.kv_cache_manager_py_config = config @@ -910,12 +917,7 @@ def append_to_kv_heads_per_layer( "Retrying without host cache tier." ) cache_tiers_gpu_only = [t for t in cache_tiers if isinstance(t, GpuCacheTierConfig)] - config = self._build_cache_config( - kv_cache_config, - tokens_per_block=tokens_per_block, - vocab_size=vocab_size, - cache_tiers=cache_tiers_gpu_only, - ) + config = replace(config, cache_tiers=cache_tiers_gpu_only) cache_tiers = cache_tiers_gpu_only self.kv_cache_manager_py_config = config self.impl = KVCacheManagerPy(config, event_manager=self.event_manager) @@ -1489,17 +1491,52 @@ def _copy_batch_block_offsets_per_layer( non_blocking=True, ) - def _build_cache_config( + def _build_base_config( self, kv_cache_config: KvCacheConfig, *, tokens_per_block: int, - vocab_size: int | None, cache_tiers: List[CacheTierConfig], ) -> KVCacheManagerConfigPy: - # Kept in the virtual method contract for cache-manager subclasses. - # The generic C++ config no longer stores the vocabulary size. - del vocab_size + """Build the general cache configuration used by most models. + + Models that need a custom configuration should subclass + :class:`KVCacheManagerV2` and override :meth:`_build_cache_config` to + update the necessary fields. + """ + scratch_reuse_config = None + if self.enable_swa_scratch_reuse: + # Context requests allocate num_extra_kv_tokens for spec decoding. + # They should not count toward the scratch range. + scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens) + + constraints = [] + if kv_cache_config.pool_ratio is None: + # CUDA graph generation warmup uses one request at max_seq_len and + # enough minimal decode requests to fill max_batch_size. + min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens + constraints.append( + BatchDesc( + [KVCacheDesc(capacity=self.max_seq_len, history_length=self.max_seq_len - 1)] + + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)] + * (self.max_batch_size - 1) + ) + ) + + # General and chunked-prefill warmup uses one fresh context request + # at the per-iteration token budget. + if self.max_num_tokens is not None: + constraints.append( + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_num_tokens + self.num_extra_kv_tokens, + history_length=0, + ) + ] + ) + ) + buffer_type = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: buffer_type.append(Role.VALUE) @@ -1512,12 +1549,6 @@ def _build_cache_config( if self.kv_cache_type != CacheTypeCpp.SELFKONLY: buffer_type.append(Role.VALUE_BLOCK_SCALE) - scratch_reuse_config = None - if self.enable_swa_scratch_reuse: - # Context requests allocate num_extra_kv_tokens for spec decoding. - # They should not count toward the scratch range. - scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens) - # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional # per-layer BufferConfig entries — for example a sparse index-K # buffer — without overriding the K/V/NVFP4 scale wiring above. @@ -1559,6 +1590,9 @@ def _build_cache_config( return KVCacheManagerConfigPy( tokens_per_block=tokens_per_block, cache_tiers=cache_tiers, + layers=layer_configs, + typical_step=None, + constraints=constraints, max_util_for_resume=kv_cache_config.max_util_for_resume, enable_partial_reuse=kv_cache_config.enable_partial_reuse, enable_stats=self.enable_stats, @@ -1568,9 +1602,12 @@ def _build_cache_config( and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE ), initial_pool_ratio=kv_cache_config.pool_ratio, - layers=layer_configs, ) + def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + """Customize the general cache config for a specialized cache manager.""" + return config + def _extra_buffers_per_layer( self, *, tokens_per_block: int ) -> Optional[dict[int, List[BufferConfig]]]: @@ -1582,7 +1619,7 @@ def _extra_buffers_per_layer( a sparse index-K buffer for each sparse local layer. Each ``BufferConfig.size`` is interpreted as bytes per block (i.e., ``bytes_per_token * tokens_per_block``), matching the standard - buffers built in :meth:`_build_cache_config`. The block storage + buffers built in :meth:`_build_base_config`. The block storage groups buffers by lifecycle and size with an opaque role key, so new roles do not require C++ changes. """ diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index a7a07926a4e1..2e60f6eca8e8 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3528,7 +3528,7 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): max_util_for_resume: float = Field( default=0.95, - ge=0, + gt=0, le=1, status="prototype", description= @@ -3677,14 +3677,6 @@ def validate_max_attention_window(cls, v: Optional[List[int]]): ) return v - @field_validator('max_util_for_resume') - @classmethod - def validate_max_util_for_resume(cls, v: float): - if not 0 <= v <= 1: - raise ValueError( - "kv_cache_config.max_util_for_resume must be between 0 and 1") - return v - @field_validator('pool_ratio') @classmethod def validate_pool_ratio(cls, v: Optional[List[float]]): diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index 460fdde0ab82..83633d3eaf25 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -202,6 +202,7 @@ def _build_deepseek_v4_cache_config_for_test( cache_manager._swa_window_size = 128 cache_manager._max_draft_len = max_draft_len cache_manager._max_num_tokens = max_num_tokens + cache_manager._avg_seq_len = kv_cache_config.avg_seq_len cache_manager.compressed_block_sizes = [128, 32, 1] cache_manager.index_head_dim = 128 cache_manager.head_dim = 512 @@ -216,12 +217,21 @@ def _build_deepseek_v4_cache_config_for_test( cache_manager.block_reuse_policy = BlockReusePolicy(kv_cache_config.block_reuse_policy) cache_manager.is_draft = is_draft - return cache_manager._build_cache_config( - kv_cache_config, + config = KVCacheManagerConfig( tokens_per_block=128, - vocab_size=129280, cache_tiers=[GpuCacheTierConfig(quota=1 << 30)], + layers=[], + typical_step=None, + max_util_for_resume=kv_cache_config.max_util_for_resume, + enable_partial_reuse=kv_cache_config.enable_partial_reuse, + enable_stats=cache_manager.enable_stats, + commit_min_snapshot=( + kv_cache_config.enable_block_reuse + and cache_manager.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE + ), + initial_pool_ratio=kv_cache_config.pool_ratio, ) + return cache_manager._build_cache_config(config) def test_deepseek_v4_pool_ratio_overrides_typical_step_and_constraints(): @@ -249,8 +259,6 @@ def test_deepseek_v4_avg_seq_len_updates_typical_step(): assert config.typical_step.kv_caches[0].history_length == 0 assert [kv.capacity for kv in config.typical_step.kv_caches[1:]] == [256, 256] assert [kv.history_length for kv in config.typical_step.kv_caches[1:]] == [253, 253] - assert config.constraints[0].kv_caches[0].capacity == 1024 - assert config.constraints[0].kv_caches[0].history_length == 1023 def test_deepseek_v4_avg_seq_len_must_not_exceed_max_seq_len(): 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 be1f2936abf8..abd5e5656f37 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -2,13 +2,21 @@ # SPDX-License-Identifier: Apache-2.0 from types import SimpleNamespace +from unittest.mock import Mock import pytest +import torch +import tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 as kv_cache_manager_v2 from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 -from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.bindings import DataType +from tensorrt_llm.bindings.internal.batch_manager import CacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig -from tensorrt_llm.runtime.kv_cache_manager_v2 import GpuCacheTierConfig, KVCacheManagerConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + GpuCacheTierConfig, + KVCacheDesc, + KVCacheManagerConfig, +) class _FakeKVCache: @@ -25,11 +33,12 @@ def stop_committing(self) -> None: self.stopped_committing = True -def _build_cache_config_for_test( +def _make_cache_config_for_test( kv_cache_config: KvCacheConfig, *, is_draft: bool = False ) -> KVCacheManagerConfig: cache_manager = object.__new__(KVCacheManagerV2) - cache_manager.kv_cache_type = CacheTypeCpp.SELFKONLY + cache_manager.kv_cache_type = CacheType.SELFKONLY + cache_manager.dtype = DataType.HALF cache_manager.head_dim_per_layer = [128] cache_manager.enable_swa_scratch_reuse = False cache_manager.num_extra_kv_tokens = 0 @@ -39,12 +48,15 @@ def _build_cache_config_for_test( cache_manager.num_local_layers = 1 cache_manager.pp_layers = [0] cache_manager.max_attention_window_vec = [None] + cache_manager.max_seq_len = 1024 + cache_manager.max_batch_size = 1 + cache_manager.max_num_tokens = None + cache_manager.max_draft_len = 0 cache_manager.get_layer_bytes_per_token = lambda **_: 128 - return cache_manager._build_cache_config( + return cache_manager._build_base_config( kv_cache_config, tokens_per_block=128, - vocab_size=129280, cache_tiers=[GpuCacheTierConfig(quota=1 << 30)], ) @@ -64,7 +76,7 @@ def test_commit_min_snapshot_follows_block_reuse_policy( is_draft: bool, commit_min_snapshot: bool, ) -> None: - config = _build_cache_config_for_test( + config = _make_cache_config_for_test( KvCacheConfig( enable_block_reuse=enable_block_reuse, block_reuse_policy=block_reuse_policy, @@ -79,11 +91,93 @@ def test_commit_min_snapshot_follows_block_reuse_policy( @pytest.mark.parametrize("enable_partial_reuse", [False, True]) def test_propagates_partial_reuse_config(enable_partial_reuse: bool) -> None: - config = _build_cache_config_for_test(KvCacheConfig(enable_partial_reuse=enable_partial_reuse)) + config = _make_cache_config_for_test(KvCacheConfig(enable_partial_reuse=enable_partial_reuse)) assert config.enable_partial_reuse is enable_partial_reuse +def test_config_preserves_resume_headroom(monkeypatch): + """Keep resume headroom for distributed quotas and constrained batches.""" + max_util_for_resume = 0.95 + mapping = SimpleNamespace( + cp_config={}, + tp_size=1, + pp_size=1, + rank=1, + world_size=2, + enable_attention_dp=False, + ) + dist = Mock() + dist.allreduce.return_value = 85 + impl = Mock(layer_grouping=[[0]]) + impl.get_layer_group_id.return_value = 0 + + def tokens_from_quota(quota): + return (quota - 40) / 2 + + def quota_from_tokens(max_tokens): + return 2 * max_tokens + 40 + + monkeypatch.setattr(kv_cache_manager_v2, "get_pp_layers", lambda *args, **kwargs: ([0], 1)) + monkeypatch.setattr(kv_cache_manager_v2.Distributed, "get", lambda mapping: dist) + monkeypatch.setattr(kv_cache_manager_v2, "KVCacheManagerPy", lambda *args, **kwargs: impl) + monkeypatch.setattr(kv_cache_manager_v2, "IndexMapper", lambda *args, **kwargs: Mock()) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: Mock()) + monkeypatch.setattr( + KVCacheManagerV2, + "_get_max_tokens_from_quota", + lambda self, quota: tokens_from_quota(quota), + ) + monkeypatch.setattr( + KVCacheManagerV2, + "_get_quota_from_max_tokens", + lambda self, max_tokens: quota_from_tokens(max_tokens), + ) + monkeypatch.setattr( + KVCacheManagerV2, + "_build_cache_config", + lambda self, config: config, + ) + monkeypatch.setattr(KVCacheManagerV2, "get_num_available_tokens", lambda self, **kwargs: 1) + monkeypatch.setattr(KVCacheManagerV2, "_prepare_page_table_tensor", lambda *args: None) + monkeypatch.setattr(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping", lambda self: None) + + manager = KVCacheManagerV2( + KvCacheConfig( + max_tokens=None, + max_gpu_total_bytes=300, + host_cache_size=0, + max_util_for_resume=max_util_for_resume, + ), + CacheType.SELF, + num_layers=1, + num_kv_heads=1, + head_dim=1, + tokens_per_block=1, + max_seq_len=1024, + max_batch_size=4, + max_num_tokens=8, + mapping=mapping, + vocab_size=1, + ) + + device_quota = manager.kv_cache_manager_py_config.cache_tiers[0].quota + required_quota = quota_from_tokens(dist.allreduce.return_value) + assert dist.allreduce.call_args.args[0] == 122.5 + assert device_quota == 222 + assert device_quota * max_util_for_resume >= required_quota + + constraints = manager.kv_cache_manager_py_config.constraints + warmup = constraints[0].kv_caches + assert len(warmup) == 4 + assert warmup[0].capacity == 1024 + assert warmup[0].history_length == 1023 + assert all(kv.capacity == 1 for kv in warmup[1:]) + assert all(kv.history_length == 0 for kv in warmup[1:]) + assert constraints[1].kv_caches == [KVCacheDesc(capacity=8, history_length=0)] + manager.shutdown() + + def test_try_commit_blocks_commits_partial_block_at_context_end() -> None: request = SimpleNamespace( py_request_id=1, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index c77fa26d7fd9..a31d41ffd205 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -409,6 +409,8 @@ def test_KvCacheConfig_declaration(): use_kv_cache_manager_v2=False).use_kv_cache_manager_v2 is False with pytest.raises(ValidationError, match="use_kv_cache_manager_v2"): KvCacheConfig(use_kv_cache_manager_v2="invalid") + with pytest.raises(ValidationError, match="max_util_for_resume"): + KvCacheConfig(max_util_for_resume=0) config = KvCacheConfig(enable_block_reuse=True, max_tokens=1024,