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 913dbbd8689c..6eb9c366ff7c 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, @@ -39,16 +36,11 @@ from tensorrt_llm.runtime import ModelConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, - 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,7 +246,6 @@ def __init__( self._init_indexer_dtype(sparse_attn_config) - # _build_cache_config() needs them to build constraints self._max_input_len = max_input_len self._max_num_tokens = max_num_tokens @@ -271,6 +262,7 @@ def __init__( vocab_size=vocab_size, mapping=mapping, dtype=dtype, + max_num_tokens=max_num_tokens, **kwargs, ) self.is_vswa = True # DeepSeek-V4 must has VSWA @@ -768,16 +760,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 to the cache config. """ layers: List[AttentionLayerConfig] = [] layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {} @@ -864,91 +849,9 @@ def _add_layer( # number of layers in the KVCacheManagerPy self._num_manager_layers = len(layers) - # Build constraints and typical_step for better 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 typical_seq_len > max_seq_len: - raise ValueError( - f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or " - f"equal to max_seq_len ({max_seq_len})" - ) - - # For aggregated serving in large batch size: - # Use 1 context request + (max_batch_size - 1) generation requests as - # the typical step. An all-generation typical_step over-provisions the - # compressed-cache pool at the expense of the SWA pool, starving the - # SWA pool and artificially capping the achievable batch size. - ctx_capacity = ( - max_num_tokens if max_num_tokens is not None else typical_seq_len - ) + self.num_extra_kv_tokens - generation_history_length = max(0, typical_seq_len - max_draft_len - 1) - typical_step = BatchDesc( - kv_caches=[ - KVCacheDesc(capacity=ctx_capacity, history_length=0), - ] - + [ - KVCacheDesc( - capacity=typical_seq_len, - history_length=generation_history_length, - ) - ] - * (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: @@ -1163,9 +1066,9 @@ def get_layer_bytes_per_token( local_layer_idx: int, data_role: DataRole, ) -> int: - raise NotImplementedError( - "DeepSeek-V4 doesn't support get_layer_bytes_per_token, use _get_attn_bytes_per_block" - ) + # The generic layers in the base config are replaced by + # _build_cache_config, so their buffer sizes are only placeholders. + return 1 def get_indexer_k_cache_buffers(self, layer_idx: int) -> torch.Tensor: """ 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 84d20ce8caa5..b612c9c3ed93 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -17,7 +17,7 @@ import os import sys from collections import OrderedDict, defaultdict -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union import numpy as np @@ -46,6 +46,7 @@ GPU_LEVEL, AttentionLayerConfig, AttnLifeCycle, + BatchDesc, BufferConfig, CacheLevel, CacheTierConfig, @@ -54,6 +55,7 @@ DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, + KVCacheDesc, KVCacheEventManager, KVCacheIterationStatsDelta, LayerId, @@ -123,8 +125,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") @@ -885,6 +887,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 @@ -896,7 +899,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) @@ -910,13 +913,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. @@ -925,7 +928,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") @@ -979,12 +985,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 @@ -998,12 +1004,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) @@ -1583,17 +1584,81 @@ 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) + + 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 self.max_seq_len + ) + if typical_seq_len > self.max_seq_len: + raise ValueError( + f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or " + f"equal to max_seq_len ({self.max_seq_len})" + ) + + # Model one context request and enough generation requests to fill + # max_batch_size without over-provisioning windowed cache pools. + context_capacity = ( + self.max_num_tokens if self.max_num_tokens is not None else typical_seq_len + ) + self.num_extra_kv_tokens + generation_history_length = max(0, typical_seq_len - self.max_draft_len - 1) + typical_step = BatchDesc( + [KVCacheDesc(capacity=context_capacity, history_length=0)] + + [ + KVCacheDesc( + capacity=typical_seq_len, + history_length=generation_history_length, + ) + ] + * (self.max_batch_size - 1) + ) + + # 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) @@ -1606,12 +1671,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. @@ -1653,6 +1712,9 @@ def _build_cache_config( return KVCacheManagerConfigPy( tokens_per_block=tokens_per_block, cache_tiers=cache_tiers, + layers=layer_configs, + typical_step=typical_step, + 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, @@ -1662,9 +1724,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]]]: @@ -1676,7 +1741,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 6cdd95827ccf..4d0b9d317fb5 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3625,7 +3625,7 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): max_util_for_resume: float = Field( default=0.95, - ge=0, + gt=0, le=1, status="prototype", description= @@ -3781,14 +3781,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 3fcfc8bafb78..bcba3baf1b24 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 @@ -34,7 +34,6 @@ ) from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.pyexecutor._util import CacheCost -from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import binding_to_torch_dtype @@ -42,11 +41,7 @@ from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.llmapi.llm_args import DeepSeekV4SparseAttentionConfig, KvCacheConfig from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import ( - GpuCacheTierConfig, - KVCacheManagerConfig, - PageIndexMode, -) +from tensorrt_llm.runtime.kv_cache_manager_v2 import BatchDesc, KVCacheDesc, PageIndexMode from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX _RequestCache = Dict[ @@ -84,35 +79,6 @@ def get_num_attention_layers(self) -> int: assert cost.intercept == 0 -def test_mtp_extra_tokens_are_in_context_capacity(): - cache_manager = object.__new__(DeepseekV4CacheManager) - cache_manager.pp_layers = [0] - cache_manager._compress_ratios = [1] - cache_manager._get_attn_bytes_per_block = lambda _attn_type, _layer_idx: 1 - cache_manager._get_window_size = lambda _compress_ratio, _attn_type: 128 - cache_manager.max_batch_size = 1 - cache_manager.max_seq_len = 264 - cache_manager._max_num_tokens = 256 - cache_manager._max_draft_len = 3 - cache_manager.num_extra_kv_tokens = 2 - cache_manager.enable_stats = False - cache_manager.enable_swa_scratch_reuse = False - cache_manager.block_reuse_policy = BlockReusePolicy.ALL_REUSABLE - - config = cache_manager._build_cache_config( - KvCacheConfig(), - tokens_per_block=128, - vocab_size=129280, - cache_tiers=[GpuCacheTierConfig(quota=1024)], - ) - - assert config.typical_step is not None - assert config.typical_step.kv_caches[0].capacity == 258 - assert config.typical_step.kv_caches[0].history_length == 0 - assert config.constraints[1].kv_caches[0].capacity == 258 - assert config.constraints[1].kv_caches[0].history_length == 0 - - def test_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(DeepseekV4CacheManager) manager.pp_layers = [0, 1] @@ -216,117 +182,6 @@ def _view_fp8_as_uint8(buffer: torch.Tensor) -> torch.Tensor: return buffer -def _build_deepseek_v4_cache_config_for_test( - kv_cache_config: KvCacheConfig, - *, - max_batch_size: int = 4, - max_seq_len: int = 1024, - max_num_tokens: int | None = 2048, - max_draft_len: int = 0, - is_draft: bool = False, -) -> KVCacheManagerConfig: - cache_manager = object.__new__(DeepseekV4CacheManager) - cache_manager.pp_layers = [0, 1, 2] - cache_manager._compress_ratios = [1, 4, 128] - cache_manager._swa_window_size = 128 - cache_manager._max_draft_len = max_draft_len - cache_manager._max_num_tokens = max_num_tokens - cache_manager.compressed_block_sizes = [128, 32, 1] - cache_manager.index_head_dim = 128 - cache_manager.head_dim = 512 - cache_manager.tokens_per_block = 128 - cache_manager.dtype = DataType.BF16 - cache_manager._indexer_k_dtype = "fp8" - cache_manager.max_batch_size = max_batch_size - cache_manager.max_seq_len = max_seq_len - cache_manager.enable_stats = False - cache_manager.enable_swa_scratch_reuse = False - cache_manager.num_extra_kv_tokens = 0 - 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, - tokens_per_block=128, - vocab_size=129280, - cache_tiers=[GpuCacheTierConfig(quota=1 << 30)], - ) - - -def test_deepseek_v4_pool_ratio_overrides_typical_step_and_constraints(): - config = _build_deepseek_v4_cache_config_for_test( - KvCacheConfig(pool_ratio=[0.2, 0.3, 0.5], avg_seq_len=256) - ) - - assert config.initial_pool_ratio == pytest.approx([0.2, 0.3, 0.5]) - assert config.typical_step is None - assert config.constraints == [] - - -def test_deepseek_v4_avg_seq_len_updates_typical_step(): - config = _build_deepseek_v4_cache_config_for_test( - KvCacheConfig(avg_seq_len=256), - max_batch_size=3, - max_seq_len=1024, - max_num_tokens=2048, - max_draft_len=2, - ) - - assert config.initial_pool_ratio is None - assert config.typical_step is not None - assert config.typical_step.kv_caches[0].capacity == 2048 - 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(): - with pytest.raises(ValueError, match="avg_seq_len"): - _build_deepseek_v4_cache_config_for_test( - KvCacheConfig(avg_seq_len=2048), - max_seq_len=1024, - ) - - -@pytest.mark.parametrize( - ("enable_block_reuse", "block_reuse_policy", "is_draft", "commit_min_snapshot"), - [ - (True, "all_reusable", False, False), - (True, "per_request", False, True), - (False, "per_request", False, False), - (True, "per_request", True, True), - ], -) -def test_deepseek_v4_commit_min_snapshot_follows_block_reuse_policy( - enable_block_reuse: bool, - block_reuse_policy: str, - is_draft: bool, - commit_min_snapshot: bool, -) -> None: - config = _build_deepseek_v4_cache_config_for_test( - KvCacheConfig( - enable_block_reuse=enable_block_reuse, - block_reuse_policy=block_reuse_policy, - enable_partial_reuse=True, - ), - is_draft=is_draft, - ) - - assert config.commit_min_snapshot is commit_min_snapshot - assert config.enable_partial_reuse - - -@pytest.mark.parametrize("enable_partial_reuse", [False, True]) -def test_deepseek_v4_propagates_partial_reuse_config(enable_partial_reuse: bool) -> None: - config = _build_deepseek_v4_cache_config_for_test( - KvCacheConfig(enable_partial_reuse=enable_partial_reuse) - ) - - assert config.enable_partial_reuse is enable_partial_reuse - - @pytest.fixture(params=[False, True], ids=["scratch_reuse_disabled", "scratch_reuse_enabled"]) def scratch_reuse_enabled(request) -> bool: return request.param @@ -1128,6 +983,28 @@ def _assert_cache_equal( msg=f"Mismatch for layer {layer_idx}, attention type {attn_type.name} (scales)", ) + def test_max_num_tokens_is_used_by_base_config(self): + max_batch_size = 2 + max_seq_len = 1024 + max_input_len = 127 + max_num_tokens = max_batch_size * (max_input_len + 1) + cache_manager, _ = self._create_deepseek_v4_cache_manager( + tokens_per_block=self.tokens_per_block, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + max_input_len=max_input_len, + compress_ratios=[1, 4], + dtype=DataType.BF16, + compressor_dtype=DataType.FLOAT, + ) + + assert cache_manager.kv_cache_manager_py_config.typical_step == BatchDesc( + [ + KVCacheDesc(capacity=max_num_tokens, history_length=0), + KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1), + ] + ) + def test_indexer_cache_layout_default(self): """DeepSeek-V4 defaults to FP4 indexer K cache on Blackwell+.""" cache_manager, _ = self._create_deepseek_v4_cache_manager( 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 307588e0d2a0..6211bc2ffb6a 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -17,7 +17,9 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( DEFAULT_BEAM_INDEX, + BatchDesc, GpuCacheTierConfig, + KVCacheDesc, KVCacheManagerConfig, ) from tensorrt_llm.runtime.kv_cache_manager_v2._utils import init_cuda_once @@ -40,26 +42,37 @@ def stop_committing(self) -> None: self.stopped_committing = True -def _build_cache_config_for_test( - kv_cache_config: KvCacheConfig, *, is_draft: bool = False +def _make_cache_config_for_test( + kv_cache_config: KvCacheConfig, + *, + is_draft: bool = False, + max_batch_size: int = 1, + max_seq_len: int = 1024, + max_num_tokens: int | None = None, + max_draft_len: int = 0, + num_extra_kv_tokens: int = 0, ) -> KVCacheManagerConfig: cache_manager = object.__new__(KVCacheManagerV2) 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 + cache_manager.num_extra_kv_tokens = num_extra_kv_tokens cache_manager.enable_stats = False cache_manager.block_reuse_policy = BlockReusePolicy(kv_cache_config.block_reuse_policy) cache_manager.is_draft = is_draft cache_manager.num_local_layers = 1 cache_manager.pp_layers = [0] cache_manager.max_attention_window_vec = [None] + cache_manager.max_seq_len = max_seq_len + cache_manager.max_batch_size = max_batch_size + cache_manager.max_num_tokens = max_num_tokens + cache_manager.max_draft_len = max_draft_len 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)], ) @@ -79,7 +92,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, @@ -94,11 +107,86 @@ 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_pool_ratio_overrides_constraints() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(pool_ratio=[1.0], avg_seq_len=256, host_cache_size=0), + max_batch_size=3, + max_num_tokens=2048, + ) + + assert config.initial_pool_ratio == pytest.approx([1.0]) + assert config.typical_step is None + assert config.constraints == [] + + +def test_builds_warmup_constraints() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(host_cache_size=0), + max_batch_size=3, + max_seq_len=1024, + max_num_tokens=2048, + max_draft_len=2, + ) + + assert config.initial_pool_ratio is None + assert config.typical_step == BatchDesc( + [KVCacheDesc(capacity=2048, history_length=0)] + + [KVCacheDesc(capacity=1024, history_length=1021)] * 2 + ) + assert config.constraints == [ + BatchDesc( + [ + KVCacheDesc(capacity=1024, history_length=1023), + KVCacheDesc(capacity=3, history_length=0), + KVCacheDesc(capacity=3, history_length=0), + ] + ), + BatchDesc([KVCacheDesc(capacity=2048, history_length=0)]), + ] + + +def test_avg_seq_len_updates_typical_step() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(avg_seq_len=256), + max_batch_size=3, + max_seq_len=1024, + max_num_tokens=2048, + max_draft_len=2, + ) + + assert config.typical_step == BatchDesc( + [KVCacheDesc(capacity=2048, history_length=0)] + + [KVCacheDesc(capacity=256, history_length=253)] * 2 + ) + + +def test_avg_seq_len_must_not_exceed_max_seq_len() -> None: + with pytest.raises(ValueError, match="avg_seq_len"): + _make_cache_config_for_test( + KvCacheConfig(avg_seq_len=2048), + max_seq_len=1024, + ) + + +def test_extra_tokens_are_in_context_capacity() -> None: + config = _make_cache_config_for_test( + KvCacheConfig(), + max_batch_size=1, + max_seq_len=264, + max_num_tokens=256, + max_draft_len=3, + num_extra_kv_tokens=2, + ) + + assert config.typical_step == BatchDesc([KVCacheDesc(capacity=258, history_length=0)]) + assert config.constraints[1] == BatchDesc([KVCacheDesc(capacity=258, history_length=0)]) + + 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 4f17a87123c3..6f8603f860ca 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -549,6 +549,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,