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 5773059bea9a..010932d58846 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 @@ -310,14 +310,10 @@ def _format_kv_cache_pool_lifecycle_entry(self, layer_id: LayerId, role: DataRol return super()._format_kv_cache_pool_lifecycle_entry(layer_id, role) model_layer_idx, attn_type = layer_semantics - attr = self.impl._storage.get_buffer_attr(layer_id, role) - pool_group_id = self.impl._storage.get_pool_group_index(attr.life_cycle_id) - lifecycle = self.impl._life_cycles.get_life_cycle(attr.life_cycle_id) return ( f"deepseek_role={attn_type.name}, " f"compress_ratio={self._compress_ratios[model_layer_idx]}, " - f"pool_group_id={int(pool_group_id)}, " - f"lifecycle_id={int(attr.life_cycle_id)}, lifecycle={lifecycle}" + f"{super()._format_kv_cache_pool_lifecycle_entry(layer_id, role)}" ) def get_buffers(self, layer_idx: int, attn_type: DeepseekV4AttentionType) -> torch.Tensor: @@ -934,7 +930,6 @@ def _add_layer( return KVCacheManagerConfigPy( tokens_per_block=tokens_per_block, - vocab_size=vocab_size, cache_tiers=cache_tiers, max_util_for_resume=kv_cache_config.max_util_for_resume, swa_scratch_reuse=scratch_reuse_config, diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index c124d678de5d..92f0a9c2f4ad 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -288,136 +288,109 @@ def _compute_global_layer_ids(manager, lg_idx: int) -> List[int]: def _build_page_table_v2(manager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManagerV2. - Uses the V2 storage layer APIs (pool.slot_address, pool.slot_size, - pool.num_slots) for accurate pool metadata, and stamps each PoolView - with the manager's native role-name strings (``pool_role``) plus the - closed-set ``mapper_kind`` discriminator used by ``build_kv_mapper``. - - Important: iterates over life cycles (layer groups), not storage pool - groups. Multiple life cycles with different sliding-window sizes may - share the same underlying storage pool group when their buffer sizes - are identical. The page table must reflect life cycles so that - per-window transfer logic works correctly. + Uses KVCacheManagerV2's public pool_group_descs layout API. A physical + pool group may be shared by several layer groups; layer_groups remains + indexed by layer_group_id while pool_group_idx points at the shared + physical pool group entry. + + Each PoolView is stamped with the manager's native role-name strings + (``pool_role``) plus the closed-set ``mapper_kind`` discriminator used + by ``build_kv_mapper``. """ - from collections import defaultdict - - from tensorrt_llm.runtime.kv_cache_manager_v2 import CacheTier - - storage = manager.impl._storage - config = manager.impl._init_config - - # Find GPU level - gpu_level = 0 - for level_idx, cache_tier_config in enumerate(config.cache_tiers): - if cache_tier_config.tier == CacheTier.GPU_MEM: - gpu_level = level_idx - break - - # Collect buffer entries keyed by (life_cycle_id, pool_idx). - # Also collect the set of native role-name strings per pool — used as - # ``PoolView.pool_role``, the manager-supplied equivalence label that - # disagg uses to match pools across peers without enumerating roles. - buffer_by_lc_pool: Dict[tuple, list] = defaultdict(list) - native_roles_by_pool: Dict[tuple, set] = defaultdict(set) - - for buffer_id, attr in storage._buffer_attr.items(): - layer_id, role = buffer_id - lc_id = attr.life_cycle_id - pool_idx = attr.pool_index - pool_key = (int(lc_id), pool_idx) - buffer_by_lc_pool[pool_key].append((layer_id, attr.offset, attr.size)) - native_roles_by_pool[pool_key].add(str(role)) - - # Iterate over life cycles (layer groups), not storage pool groups. - # Multiple layer_groups can share the same storage pool_group when their - # slot_size_list (coalesced buffer sizes) are identical. In that case, - # different layer_groups draw slots from the same physical pool, but a - # slot is exclusively allocated to one layer_group at a time (managed by - # SlotAllocator). Within a slot, each layer_group's buffer offsets start - # from 0 independently — the memory is reused, not concatenated. - # Therefore, slot_bytes / num_layers_for_this_layer_group correctly gives - # the per-layer size, and buffer offsets within a slot are contiguous for - # each layer_group. - num_life_cycles = storage.num_life_cycles - pool_group_storage = storage._levels[gpu_level].storage._pool_groups + config = manager.impl.init_config + pool_group_descs = manager.impl.pool_group_descs - pool_groups: List[PhysicalPoolGroup] = [] - storage_pg_to_list_idx: Dict[int, int] = {} - layer_groups: List[LayerGroup] = [] - - for lc_idx in range(num_life_cycles): - # Resolve the storage pool group for this life cycle. - # storage_pg_idx may be the same for multiple lc_idx values. - storage_pg_idx = storage.get_pool_group_index(lc_idx) - pool_group = pool_group_storage[storage_pg_idx] - num_pools = pool_group.num_pools - - # Build PhysicalPoolGroup once per unique storage pool group. - if storage_pg_idx not in storage_pg_to_list_idx: - storage_pg_to_list_idx[storage_pg_idx] = len(pool_groups) - pool_groups.append( - PhysicalPoolGroup( - pools=[ - PhysicalPool( - base_address=int(pool_group._pools[pi].slot_address(0)), - slot_bytes=int(pool_group._pools[pi].slot_size), - num_slots=int(pool_group._pools[pi].num_slots), - ) - for pi in range(num_pools) - ] - ) - ) + def _window_size_for_layer(internal_layer_id: int): + if internal_layer_id < len(config.layers): + return getattr(config.layers[internal_layer_id], "window_size", None) - # Compute group-level global layer IDs and internal layer IDs - all_internal_layer_ids = list(manager.impl.layer_grouping[lc_idx]) - all_global_layer_ids = _compute_global_layer_ids(manager, lc_idx) + if hasattr(manager, "_layer_attn_to_layer_id"): + for (model_layer, _attn_type), layer_id in manager._layer_attn_to_layer_id.items(): + if layer_id != internal_layer_id: + continue + local_layer = manager.layer_offsets.get(model_layer) + if local_layer is not None and local_layer < len(config.layers): + return getattr(config.layers[local_layer], "window_size", None) + if model_layer < len(config.layers): + return getattr(config.layers[model_layer], "window_size", None) - local_layers = [ - LocalLayer(local_layer_id=int(iid), global_layer_id=int(gid)) - for iid, gid in zip(all_internal_layer_ids, all_global_layer_ids) - ] + raise ValueError(f"Cannot resolve layer config for internal layer {internal_layer_id}") - pool_views = [] - for pool_idx in range(num_pools): - pool_key = (lc_idx, pool_idx) - buffers_info = buffer_by_lc_pool.get(pool_key, []) - - # Skip pools that have no buffers for this layer group. - # Multiple life cycles may share the same storage pool group; - # only include pools that actually belong to this life cycle. - if not buffers_info: - continue - - pool_views.append( - PoolView( - pool_idx=pool_idx, - buffer_entries=np.array(buffers_info, dtype=BUFFER_ENTRY_DTYPE), - pool_role=frozenset(native_roles_by_pool[pool_key]), - mapper_kind=MapperKind.INDEXED, - ) + pool_groups: List[PhysicalPoolGroup] = [] + storage_pg_to_list_idx: Dict[int, int] = {} + layer_groups_by_id: List[LayerGroup | None] = [None] * len(manager.impl.layer_grouping) + + for pg_desc in pool_group_descs: + storage_pg_idx = int(pg_desc.pool_group_index) + storage_pg_to_list_idx[storage_pg_idx] = len(pool_groups) + pool_groups.append( + PhysicalPoolGroup( + pools=[ + PhysicalPool( + base_address=int(pool.base_address), + slot_bytes=int(pool.slot_bytes), + num_slots=int(pg_desc.num_slots), + ) + for pool in pg_desc.pools + ] ) + ) - # Determine layer group metadata. - # For managers with virtual layers, internal layer_ids - # may exceed the length of num_kv_heads_per_layer. Use index 0 as all - # layers within a pool group share the same kv_heads count. - first_local_layer = all_internal_layer_ids[0] - if first_local_layer < len(manager.num_kv_heads_per_layer): - num_kv_heads = manager.num_kv_heads_per_layer[first_local_layer] - else: - num_kv_heads = manager.num_kv_heads_per_layer[0] - life_cycle = manager.impl._life_cycles[lc_idx] - sliding_window_size = life_cycle.window_size + for variant in pg_desc.slot_desc.variants: + layer_group_id = int(variant.layer_group_id) + all_internal_layer_ids = list(manager.impl.layer_grouping[layer_group_id]) + all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id) + + local_layers = [ + LocalLayer(local_layer_id=int(iid), global_layer_id=int(gid)) + for iid, gid in zip(all_internal_layer_ids, all_global_layer_ids) + ] + + pool_views = [] + for pool_idx, coalesced_buffer in enumerate(variant.coalesced_buffers): + entries = [] + # Native role-name strings for this pool — used as + # ``PoolView.pool_role``, the manager-supplied equivalence + # label that disagg uses to match pools across peers without + # enumerating roles. + native_roles: set = set() + offset = 0 + single_buffer_size = int(coalesced_buffer.single_buffer_size) + for buffer_id in coalesced_buffer.buffer_ids: + entries.append((int(buffer_id.layer_id), offset, single_buffer_size)) + native_roles.add(str(buffer_id.role)) + offset += single_buffer_size + + if entries: + pool_views.append( + PoolView( + pool_idx=pool_idx, + buffer_entries=np.array(entries, dtype=BUFFER_ENTRY_DTYPE), + pool_role=frozenset(native_roles), + mapper_kind=MapperKind.INDEXED, + ) + ) - layer_groups.append( - AttentionLayerGroup( + first_local_layer = all_internal_layer_ids[0] + if first_local_layer < len(manager.num_kv_heads_per_layer): + num_kv_heads = manager.num_kv_heads_per_layer[first_local_layer] + else: + num_kv_heads = manager.num_kv_heads_per_layer[0] + sliding_window_size = _window_size_for_layer(first_local_layer) + + layer_groups_by_id[layer_group_id] = AttentionLayerGroup( pool_group_idx=storage_pg_to_list_idx[storage_pg_idx], kv_head_num_per_rank=num_kv_heads, sliding_window_size=sliding_window_size, local_layers=local_layers, pool_views=pool_views, ) - ) + + layer_groups: List[LayerGroup] = [] + for layer_group_id, layer_group in enumerate(layer_groups_by_id): + if layer_group is None: + raise ValueError(f"Missing V2 layer group descriptor for layer group {layer_group_id}") + layer_groups.append(layer_group) if isinstance(manager, MambaHybridCacheManager): mamba_layer_group_idx = len(pool_groups) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index ee03e2171a22..21c4d98bd2aa 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -47,15 +47,29 @@ def get_pool_view_global_layer_ids( pool_view: PoolView, layer_group: AttentionLayerGroup ) -> List[int]: """ - Global layer IDs for the layers that appear in *pool_view*, ordered as in - *layer_group.local_layers*. - """ - local_ids_in_pool = get_unique_layers(pool_view) - return [ - ll.global_layer_id - for ll in layer_group.local_layers - if ll.local_layer_id in local_ids_in_pool - ] + Global layer IDs for the layers that appear in *pool_view*, ordered by their + physical offset within the coalesced buffer (ascending). + + The order is derived from the buffer entries' physical offsets rather than + from ``layer_group.local_layers`` order on purpose: the KV transfer maps + layers positionally (a layer's position in this list times the per-layer + slot size gives its byte offset), so the position must reflect the physical + slot layout. Deriving it from offsets keeps the transceiver decoupled from + the KV-cache manager's layer-grouping order (which is an implementation + detail, not an API contract). This mirrors ``get_aggregated_pages``, which + likewise sorts buffers by their offset inside the coalesced buffer. + """ + local_to_global = {ll.local_layer_id: ll.global_layer_id for ll in layer_group.local_layers} + # A layer may contribute several buffer entries (e.g. KEY and VALUE); use the + # smallest offset as that layer's position within the slot. + min_offset: dict[int, int] = {} + for entry in pool_view.buffer_entries: + local_layer_id = int(entry["local_layer_id"]) + offset = int(entry["offset"]) + if local_layer_id not in min_offset or offset < min_offset[local_layer_id]: + min_offset[local_layer_id] = offset + ordered_local_ids = sorted(min_offset, key=lambda lid: min_offset[lid]) + return [local_to_global[lid] for lid in ordered_local_ids] # ------------------------------------------------------------------------- diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 2f93f0fb028f..0a9d8bd8a6d1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -15,8 +15,8 @@ import hashlib import math import os +import sys from collections import OrderedDict, defaultdict -from dataclasses import fields from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union import numpy as np @@ -38,40 +38,38 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.runtime.kv_cache_hash import get_effective_kv_cache_event_hash_algo from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, + BAD_PAGE_INDEX, + CACHE_LEVEL1, DEFAULT_BEAM_INDEX, + GPU_LEVEL, AttentionLayerConfig, + AttnLifeCycle, BufferConfig, + CacheLevel, CacheTierConfig, + CuError, + DataRole, DiskCacheTierConfig, GpuCacheTierConfig, HostCacheTierConfig, + KVCacheEventManager, KVCacheIterationStatsDelta, LayerId, + LifeCycleId, PageIndexMode, + PoolGroupPeakBlockStats, ReuseScope, SwaScratchReuseConfig, TokenIdExt, _KVCache, + exact_div, + gen_multimodal_cache_key_tokens, + typed_range, ) from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as KVCacheManagerPy from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy -from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( - gen_multimodal_cache_key_tokens, -) -from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( - BAD_PAGE_INDEX, - CACHE_LEVEL1, - GPU_LEVEL, - CacheLevel, -) -from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole -from tensorrt_llm.runtime.kv_cache_manager_v2._event_manager import KVCacheEventManager -from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import CuError -from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import ( - OutOfMemoryError as KVCacheOutOfMemoryError, -) -from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import AttnLifeCycle, LifeCycleId -from tensorrt_llm.runtime.kv_cache_manager_v2._utils import exact_div, typed_range +from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfMemoryError as KVCacheOutOfMemoryError from tensorrt_llm.sampling_params import SamplingParams from ..._utils import binding_to_torch_dtype, mpi_rank, nvtx_range, str_dtype_to_torch @@ -101,9 +99,7 @@ if TYPE_CHECKING: from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata -KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple( - field.name for field in fields(KVCacheIterationStatsDelta) -) +KV_CACHE_ITERATION_STATS_DELTA_FIELDS = _KV_CACHE_ITERATION_STATS_DELTA_FIELDS KV_CACHE_ITERATION_STATS_REUSE_FIELDS = ( "iter_reused_blocks", "iter_full_reused_blocks", @@ -788,7 +784,7 @@ def append_to_kv_heads_per_layer( self.is_vswa = len(set(self.max_attention_window_vec)) > 1 - quota = float("inf") + quota = sys.maxsize if ( kv_cache_config.max_gpu_total_bytes is not None and kv_cache_config.max_gpu_total_bytes > 0 @@ -809,7 +805,7 @@ def append_to_kv_heads_per_layer( f"New quota is {quota / (1 << 30)}GiB" ) - assert quota != float("inf"), ( + assert quota < sys.maxsize, ( "Quota not set. Check kv_cache_config.max_tokens or kv_cache_config.max_gpu_total_bytes" ) @@ -832,7 +828,7 @@ def append_to_kv_heads_per_layer( logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB") - cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=quota)] + cache_tiers: List[CacheTierConfig] = [GpuCacheTierConfig(quota=int(quota))] if kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size >= 0: host_quota = kv_cache_config.host_cache_size else: @@ -865,7 +861,7 @@ def append_to_kv_heads_per_layer( if host_quota <= 0: host_quota = quota if host_quota > 0: - cache_tiers.append(HostCacheTierConfig(quota=host_quota)) + cache_tiers.append(HostCacheTierConfig(quota=int(host_quota))) logger.info( f"KV cache manager v2 host cache quota set to {host_quota / (1 << 30):.2f}GiB" ) @@ -873,7 +869,9 @@ def append_to_kv_heads_per_layer( if disk_cache_size is not None and disk_cache_size > 0: disk_cache_path = kv_cache_config.disk_cache_path assert disk_cache_path is not None - cache_tiers.append(DiskCacheTierConfig(quota=disk_cache_size, path=disk_cache_path)) + cache_tiers.append( + DiskCacheTierConfig(quota=int(disk_cache_size), path=disk_cache_path) + ) logger.info( f"KV cache manager v2 disk cache quota set to {disk_cache_size / (1 << 30):.2f}GiB at {disk_cache_path}" ) @@ -1483,9 +1481,12 @@ def _build_cache_config( kv_cache_config: KvCacheConfig, *, tokens_per_block: int, - vocab_size: Optional[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 buffer_type = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: buffer_type.append(Role.VALUE) @@ -1544,7 +1545,6 @@ def _build_cache_config( return KVCacheManagerConfigPy( tokens_per_block=tokens_per_block, - vocab_size=vocab_size, cache_tiers=cache_tiers, max_util_for_resume=kv_cache_config.max_util_for_resume, enable_stats=self.enable_stats, @@ -2220,15 +2220,48 @@ def _stats_life_cycle_window_size(self, life_cycle) -> Optional[int]: return None return self._stats_window_size(life_cycle.window_size) + def _get_storage_statistics(self, cache_level: CacheLevel): + return self.impl._storage.get_statistics(cache_level) + + def _stats_life_cycle_metadata(self) -> dict[int, tuple[int, Optional[int], str]]: + pool_groups_by_life_cycle = [ + self.impl._storage.get_pool_group_index(LifeCycleId(life_cycle_id)) + for life_cycle_id in range(len(self.impl.layer_grouping)) + ] + + metadata: dict[int, tuple[int, Optional[int], str]] = {} + for life_cycle_id, layer_ids in enumerate(self.impl.layer_grouping): + if not layer_ids: + continue + layer = self.kv_cache_manager_py_config.layers[int(layer_ids[0])] + is_attention = isinstance(layer, AttentionLayerConfig) + metadata[life_cycle_id] = ( + int(pool_groups_by_life_cycle[life_cycle_id]), + self._stats_window_size(layer.sliding_window_size) if is_attention else None, + "attention" if is_attention else "ssm", + ) + return metadata + def _storage_pool_groups_by_window(self) -> dict[int, set[int]]: pool_groups_by_window: dict[int, set[int]] = defaultdict(set) - for life_cycle_id, life_cycle in self.impl._life_cycles.attention_life_cycles(): - pool_group_id = self.impl._storage.get_pool_group_index(life_cycle_id) - pool_groups_by_window[self._stats_window_size(life_cycle.window_size)].add( - int(pool_group_id) - ) + for pool_group_id, window_size, _ in self._stats_life_cycle_metadata().values(): + if window_size is not None: + pool_groups_by_window[window_size].add(pool_group_id) return pool_groups_by_window + def _get_and_reset_iteration_peak_block_stats(self, cache_level: CacheLevel): + get_peak_stats = getattr(self.impl, "get_and_reset_iteration_peak_block_stats", None) + if get_peak_stats is not None: + return get_peak_stats(cache_level) + return [ + PoolGroupPeakBlockStats( + available=stats.available, + unavailable=stats.total - stats.available, + evictable=stats.evictable, + ) + for stats in self._get_storage_statistics(cache_level) + ] + @staticmethod def _windows_by_pool_group( pool_groups_by_window: dict[int, set[int]], @@ -2345,7 +2378,7 @@ def _build_iteration_stats( return stats def _collect_iteration_stats_deltas( - self, raw_iteration_stats, storage + self, raw_iteration_stats, life_cycle_metadata ) -> tuple[dict, dict, dict, dict]: reuse_deltas_by_window: dict[int, KVCacheIterationStatsDelta] = {} reuse_deltas_by_life_cycle: dict[int, KVCacheIterationStatsDelta] = {} @@ -2353,9 +2386,7 @@ def _collect_iteration_stats_deltas( pool_group_deltas: dict[int, KVCacheIterationStatsDelta] = {} for life_cycle_id, delta in raw_iteration_stats.items(): - life_cycle = self.impl._life_cycles.get_life_cycle(life_cycle_id) - pool_group_id = int(storage.get_pool_group_index(life_cycle_id)) - window_size = self._stats_life_cycle_window_size(life_cycle) + pool_group_id, window_size, _ = life_cycle_metadata[int(life_cycle_id)] pool_group_delta = self._filter_iteration_stats_delta( delta, KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS @@ -2421,9 +2452,15 @@ def _build_pool_group_iteration_stats( secondary_peak_stats_by_level, pool_group_delta, ) -> KVCacheV2PoolGroupIterationStats: + primary_pool_group_stats = primary_stats[pool_group_id] + slot_size = ( + primary_pool_group_stats.slot_sizes + if hasattr(primary_pool_group_stats, "slot_sizes") + else primary_pool_group_stats.slot_size + ) return KVCacheV2PoolGroupIterationStats( pool_group_id=pool_group_id, - slot_size=tuple(primary_stats[pool_group_id].slot_size), + slot_size=tuple(slot_size), window_sizes=windows_by_pool_group.get(pool_group_id, ()), stats=self._build_iteration_stats( (pool_group_id,), @@ -2439,21 +2476,19 @@ def _build_pool_group_iteration_stats( def _build_life_cycle_iteration_stats( self, life_cycle_id: int, - storage, + life_cycle_metadata, primary_stats, secondary_stats_by_level, primary_peak_stats, secondary_peak_stats_by_level, reuse_delta, ) -> KVCacheV2LifeCycleIterationStats: - typed_life_cycle_id = LifeCycleId(life_cycle_id) - life_cycle = self.impl._life_cycles.get_life_cycle(typed_life_cycle_id) - pool_group_id = int(storage.get_pool_group_index(typed_life_cycle_id)) + pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] return KVCacheV2LifeCycleIterationStats( life_cycle_id=life_cycle_id, pool_group_id=pool_group_id, - window_size=self._stats_life_cycle_window_size(life_cycle), - kind="attention" if isinstance(life_cycle, AttnLifeCycle) else "ssm", + window_size=window_size, + kind=kind, stats=self._build_iteration_stats( (), primary_stats, @@ -2467,7 +2502,7 @@ def _build_life_cycle_iteration_stats( def get_kv_cache_stats(self): kv_cache_stats = KvCacheStats() - pool_group_stats = self.impl._storage.get_statistics(GPU_LEVEL) + pool_group_stats = self._get_storage_statistics(GPU_LEVEL) max_num_blocks = sum(stat.total for stat in pool_group_stats) free_num_blocks = sum(stat.available for stat in pool_group_stats) committed_stats = self.impl.get_committed_stats() @@ -2509,29 +2544,29 @@ def get_iteration_stats(self): if not self.enable_stats: return None - storage = self.impl._storage + life_cycle_metadata = self._stats_life_cycle_metadata() pool_groups_by_window = self._storage_pool_groups_by_window() windows_by_pool_group = self._windows_by_pool_group(pool_groups_by_window) raw_iteration_stats = self.impl.get_and_reset_iteration_stats() - primary_peak_stats = self.impl.get_and_reset_iteration_peak_block_stats(GPU_LEVEL) + primary_peak_stats = self._get_and_reset_iteration_peak_block_stats(GPU_LEVEL) + num_cache_levels = len(self.impl.cache_tier_list) secondary_peak_stats_by_level = [ - self.impl.get_and_reset_iteration_peak_block_stats(CacheLevel(level)) - for level in range(1, int(storage.num_cache_levels)) + self._get_and_reset_iteration_peak_block_stats(CacheLevel(level)) + for level in range(1, num_cache_levels) ] ( reuse_deltas_by_window, reuse_deltas_by_life_cycle, pool_group_deltas_by_window, pool_group_deltas, - ) = self._collect_iteration_stats_deltas(raw_iteration_stats, storage) + ) = self._collect_iteration_stats_deltas(raw_iteration_stats, life_cycle_metadata) windows = set(pool_groups_by_window) windows.update(reuse_deltas_by_window) windows.update(pool_group_deltas_by_window) - primary_stats = storage.get_statistics(GPU_LEVEL) + primary_stats = self._get_storage_statistics(GPU_LEVEL) secondary_stats_by_level = [ - storage.get_statistics(CacheLevel(level)) - for level in range(1, int(storage.num_cache_levels)) + self._get_storage_statistics(CacheLevel(level)) for level in range(1, num_cache_levels) ] stats_by_window = { @@ -2566,7 +2601,7 @@ def get_iteration_stats(self): stats_by_life_cycle = { life_cycle_id: self._build_life_cycle_iteration_stats( life_cycle_id, - storage, + life_cycle_metadata, primary_stats, secondary_stats_by_level, primary_peak_stats, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md b/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md index 5bec58daf1e5..c9a096a03f5a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md @@ -49,7 +49,7 @@ make all ### Debug Mode -Set `TLLM_KV_CACHE_MANAGER_V2_DEBUG=1` to enable debug assertions (`NDEBUG=False`). Default is release mode (`NDEBUG=True`). +Set `TLLM_DEBUG_MODE=1` to enable debug assertions (`NDEBUG=False`). Default is release mode (`NDEBUG=True`). ## Architecture diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 28276acb310b..34399289b9ae 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -13,9 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -from . import rawref -from ._block_radix_tree import ReuseScope, gen_multimodal_cache_key_tokens -from ._common import ( +from . import rawref # noqa: F401 +from ._block_radix_tree import ReuseScope, gen_multimodal_cache_key_tokens # noqa: F401 +from ._common import ( # noqa: F401 + BAD_PAGE_INDEX, + CACHE_LEVEL1, + GPU_LEVEL, NDEBUG, CacheLevel, CacheTier, @@ -23,11 +26,13 @@ LayerId, MemAddress, PageIndexMode, + PageStatus, Priority, + SlidingWindowSize, TokenId, TokenIdExt, ) -from ._config import ( +from ._config import ( # noqa: F401 AttentionLayerConfig, BatchDesc, BufferConfig, @@ -41,17 +46,21 @@ SsmLayerConfig, SwaScratchReuseConfig, ) -from ._core import ( +from ._core import ( # noqa: F401 DEFAULT_BEAM_INDEX, AggregatedPageDesc, BeamIndex, + ExpandedBuffer, KVCacheManager, PageIndexConverter, + PoolDesc, + PoolGroupDesc, PoolGroupPeakBlockStats, ScratchDesc, _KVCache, ) -from ._event_manager import ( +from ._core._kv_cache import _Status as KvCacheStatus # noqa: F401 +from ._event_manager import ( # noqa: F401 KVCacheCreatedData, KVCacheEvent, KVCacheEventDiff, @@ -62,56 +71,85 @@ KVCacheUpdatedData, UniqueToken, ) -from ._life_cycle_registry import LayerGroupId, LifeCycleId -from ._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta -from ._storage import BufferId +from ._exceptions import CuError, OutOfMemoryError, OutOfPagesError # noqa: F401 +from ._life_cycle_registry import AttnLifeCycle, LayerGroupId, LifeCycleId # noqa: F401 +from ._stats import ( # noqa: F401 + _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, + KVCacheIterationStatsDelta, + KVCacheStatsDelta, +) +from ._storage import BufferId # noqa: F401 +from ._storage._config import CoalescedBuffer, SlotDesc, SlotDescVariant # noqa: F401 +from ._storage._core import PoolGroupIndex, PoolIndex # noqa: F401 +from ._utils import HalfOpenRange, exact_div, typed_range # noqa: F401 __all__ = [ - "LifeCycleId", - "LayerGroupId", - "TokenId", - "TokenIdExt", - "KVCacheManager", - "_KVCache", + "AggregatedPageDesc", + "AttentionLayerConfig", + "BAD_PAGE_INDEX", + "CACHE_LEVEL1", + "BatchDesc", + "BeamIndex", + "BufferConfig", + "BufferId", + "CoalescedBuffer", + "CacheLevel", + "CacheTier", + "CacheTierConfig", + "CudaStream", + "DEFAULT_BEAM_INDEX", + "DataRole", + "DiskCacheTierConfig", + "ExpandedBuffer", + "GPU_LEVEL", + "GpuCacheTierConfig", + "HalfOpenRange", + "HostCacheTierConfig", + "KVCacheDesc", "KVCacheCreatedData", "KVCacheEvent", "KVCacheEventDiff", "KVCacheEventManager", + "KVCacheManager", + "KVCacheManagerConfig", "KVCacheRemovedData", "KVCacheStoredBlockData", "KVCacheStoredData", "KVCacheUpdatedData", - "UniqueToken", - "BeamIndex", - "DEFAULT_BEAM_INDEX", + "KvCacheStatus", + "LayerGroupId", "LayerId", - "Priority", - "ReuseScope", - "CacheLevel", - "CacheTier", - "CudaStream", + "LifeCycleId", "MemAddress", "NDEBUG", - "KVCacheManagerConfig", - "SwaScratchReuseConfig", - "AttentionLayerConfig", - "SsmLayerConfig", - "BufferConfig", - "DataRole", - "DiskCacheTierConfig", - "GpuCacheTierConfig", - "HostCacheTierConfig", - "BatchDesc", - "CacheTierConfig", - "KVCacheDesc", - "gen_multimodal_cache_key_tokens", - "rawref", - "AggregatedPageDesc", - "BufferId", + "OutOfPagesError", "PageIndexConverter", "PoolGroupPeakBlockStats", "PageIndexMode", + "PageStatus", + "PoolDesc", + "PoolGroupDesc", + "PoolGroupIndex", + "PoolIndex", + "Priority", + "ReuseScope", "ScratchDesc", "KVCacheIterationStatsDelta", "KVCacheStatsDelta", + "SlidingWindowSize", + "SlotDesc", + "SlotDescVariant", + "SsmLayerConfig", + "SwaScratchReuseConfig", + "TokenId", + "TokenIdExt", + "UniqueToken", + "AttnLifeCycle", + "CuError", + "OutOfMemoryError", + "_KVCache", + "exact_div", + "gen_multimodal_cache_key_tokens", + "rawref", + "typed_range", ] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 5d7c50e676cb..61ddc4bdbb4f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -62,6 +62,7 @@ BeamIndex = NewType("BeamIndex", int) MemAddress = NewType("MemAddress", int) Priority = NewType("Priority", int) PoolGroupIndex = NewType("PoolGroupIndex", int) +PoolIndex = NewType("PoolIndex", int) # From _stats.py @dataclass(slots=True) @@ -165,7 +166,6 @@ class SwaScratchReuseConfig: @dataclass(slots=True) class KVCacheManagerConfig: tokens_per_block: int - vocab_size: int cache_tiers: list[CacheTierConfig] layers: list[LayerConfig] max_util_for_resume: float = ... @@ -345,6 +345,10 @@ class _KVCache: ) -> None: ... @property def num_committed_tokens(self) -> int: ... + @property + def committed_tokens(self) -> list[TokenIdExt]: ... + @property + def reuse_scope(self) -> ReuseScope: ... def stop_committing(self) -> None: ... def suspend(self) -> None: ... def resume(self, cuda_stream: CudaStream | None = None) -> bool: ... @@ -365,14 +369,10 @@ class _KVCache: def tokens_per_block(self) -> int: ... @dataclass(slots=True, frozen=True) -class MemoryPoolDesc: - base: MemAddress - page_size: int - -@dataclass(slots=True, frozen=True) -class MemoryPoolGroupDesc: - num_pages: int - pools: Sequence[MemoryPoolDesc] +class PoolDesc: + pool_index: PoolIndex + base_address: MemAddress + slot_bytes: int class BufferId(NamedTuple): layer_id: LayerId @@ -396,6 +396,36 @@ class AggregatedPageDesc: layer_group_id: LayerGroupId buffers: Sequence[ExpandedBuffer] +@dataclass(slots=True, frozen=True) +class CoalescedBuffer: + single_buffer_size: int + buffer_ids: Sequence[BufferId] + @property + def size(self) -> int: ... + @property + def num_buffers(self) -> int: ... + +@dataclass(slots=True, frozen=True) +class SlotDescVariant: + coalesced_buffers: Sequence[CoalescedBuffer] + @property + def layer_group_id(self) -> LayerGroupId: ... + @property + def slot_size_list(self) -> Sequence[int]: ... + +@dataclass(slots=True, frozen=True) +class SlotDesc: + variants: Sequence[SlotDescVariant] + @property + def slot_size_list(self) -> Sequence[int]: ... + +@dataclass(slots=True, frozen=True) +class PoolGroupDesc: + pool_group_index: PoolGroupIndex + num_slots: int + slot_desc: SlotDesc + pools: Sequence[PoolDesc] + # From _core/_kv_cache_manager.py @dataclass(slots=True, frozen=True) class ScratchDesc: @@ -467,6 +497,8 @@ class KVCacheManager: @property def event_manager(self) -> Any | None: ... @property + def init_config(self) -> KVCacheManagerConfig: ... + @property def allow_seq_rebasing(self) -> bool: ... @property def enable_partial_match(self) -> bool: ... @@ -481,6 +513,8 @@ class KVCacheManager: @property def all_buffer_ids(self) -> Iterator[BufferId]: ... def get_aggregated_pages(self, buffers: Iterable[BufferId]) -> Iterator[AggregatedPageDesc]: ... + @property + def pool_group_descs(self) -> Sequence[PoolGroupDesc]: ... def clamp_max_seq_len_for_mem(self, batch_size: int, token_num_upper_bound: int) -> int: ... def adjust(self) -> None: ... @property diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index a52e4f456df7..38798e17b03a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -28,6 +28,7 @@ filled_list, find_index, map_optional, + typed_enumerate, unwrap_rawref, ) @@ -156,40 +157,50 @@ def get_tree(block: "RootBlock | Block") -> "BlockRadixTree": return tree -def remove_subtree(root: "RootBlock | Block") -> list[rawref.ref["CommittedPage"]]: +def detach_next(parent: "Block | RootBlock", key: BlockKey) -> "Block | None": + child = parent.next.pop(key, None) + if child is None: + return None + + child._prev = rawref.NULL + if isinstance(parent, RootBlock) and not parent.next: + tree = parent._prev() + if tree is not None and parent.key in tree.next: + detached_root = tree.next.pop(parent.key) + parent._prev = rawref.NULL + assert detached_root is parent + return child + + +def remove_subtree(root: "Block") -> None: # taking O(1) space # remove leaf blocks one by one, in post-order - ret: list[rawref.ref["CommittedPage"]] = [] + # Each block's pages are reclaimed eagerly via _release_pages() while the + # StorageManager is still alive, rather than deferring to ~Block()/__del__(). + # An external reference (e.g. a caller holding a matched Block) can keep a Block + # alive past StorageManager teardown, after which page.manager would be dangling. removed_block_hashes: list[BlockKey] = [] tree = try_get_tree(root) event_manager = tree.event_manager if tree is not None else None - block: "RootBlock | Block" = root + block: Block = root while True: if block.next: block = next(iter(block.next.values())) else: - if isinstance(block, Block): - removed_block_hashes.append(block.key) - ret.extend(p for p in block.storage if p is not None) - block.storage = filled_list(None, block.num_life_cycles) - assert isinstance(block, RootBlock) or all(page is None for page in block.storage), ( - "Storage is not cleared, yet" - ) + block._release_pages() + removed_block_hashes.append(block.key) if block._prev() is None: assert block is root break - prev_block: Block | RootBlock | BlockRadixTree = block.prev - # Because Block.__del__() may remove RootBlock from BlockRadixTree, we need to check here. - # It may not be in prev_block.next when block is RootBlock. - if block.key in prev_block.next: - prev_block.next.pop(block.key) + prev_block: Block | RootBlock = block.prev + detached = detach_next(prev_block, block.key) + assert detached is block if block is root: break - assert not isinstance(prev_block, BlockRadixTree) + assert isinstance(prev_block, Block) block = prev_block if event_manager is not None: event_manager.add_removed_event(removed_block_hashes) - return ret def traverse_post_order(root: "Block") -> Iterator["Block"]: @@ -351,22 +362,38 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N to_remove.append(k) event_manager = get_tree(prev).event_manager if to_remove else None for k in to_remove: - b = prev.next.pop(k) + b = detach_next(prev, k) + assert isinstance(b, Block) if event_manager is not None: event_manager.add_removed_event(b.key) assert b.is_orphan # _KVCache may still hold it. # prev.next keeps a strong ref to this _Block, so no need to remove self from prev.next in __del__(). prev.next[self.key] = self - def __del__(self) -> None: - for ref in self.storage: + def _release_pages(self) -> None: + """Reclaim every page held by this block. + + Nulls each page's back-pointer and, for pages still scheduled for eviction, + removes them from the eviction controller (releasing their storage slots). + Idempotent: afterwards ``storage`` holds no pages, so it is safe to call again + from ``__del__``. + + This must run during radix-tree teardown (``remove_subtree``/``clear``) rather + than being deferred to ``__del__``, so that page reclamation does not depend on + this ``Block`` object's destruction timing. An external reference can keep the + ``Block`` alive past ``StorageManager`` teardown, after which ``page.manager`` + would be a dangling reference. + """ + for lc_idx, ref in typed_enumerate(self.storage): if ref is not None and ref() is not None: page = unwrap_rawref(ref) + self.unlink_page(lc_idx) if page.status == PageStatus.DROPPABLE: if page.scheduled_for_eviction: page.manager.exclude_from_eviction(page) - if self._prev() is not None and isinstance(self.prev, RootBlock) and not self.prev.next: - self.prev.prev.next.pop(self.prev.key) + + def __del__(self) -> None: + self._release_pages() self.__rawref__.invalidate() def _partial_match_this_node(self, tokens: TokenBlock) -> int: @@ -386,42 +413,50 @@ def num_life_cycles(self) -> LifeCycleId: def prev(self) -> "Block | RootBlock": return unwrap_rawref(self._prev) - def unset_page( - self, lc_idx: LifeCycleId, lc: LifeCycle, expected_page: "CommittedPage | None" = None - ) -> None: - ref = self.storage[lc_idx] - if ref is None: - return - if expected_page is not None and ref() is not expected_page: - return - ordinal = self.ordinal + def unlink_page( + self, lc_idx: LifeCycleId, expected_page: "CommittedPage | None" = None + ) -> bool: + page_ref = self.storage[lc_idx] + if page_ref is None: + return False + # Only unlink when the slot still holds the expected page. During rebase + # another block with the same key may have replaced the stored page, and + # unlinking then would clobber the newer page's back-pointer. + if expected_page is not None and page_ref() is not expected_page: + return False + page = page_ref() + if page is not None: + page.block = rawref.NULL self.storage[lc_idx] = None - tree = try_get_tree(self) + return True + + @staticmethod + def clear_stale_blocks_after_page_unlink( + start: "Block", lc_idx: LifeCycleId, lc: LifeCycle + ) -> None: + assert start.storage[lc_idx] is None + ordinal = start.ordinal + tree = try_get_tree(start) event_manager = tree.event_manager if tree is not None else None if type(lc) is AttnLifeCycle and (lc.window_size is None or ordinal < lc.num_sink_blocks): - pages = remove_subtree(self) - for r in pages: - if r() is not None: - page = unwrap_rawref(r) - assert page.status == PageStatus.DROPPABLE - if page.scheduled_for_eviction: - page.manager.exclude_from_eviction(page) + remove_subtree(start) elif event_manager is not None: - event_manager.add_removed_life_cycle_event(self.key, int(lc_idx)) + event_manager.add_removed_life_cycle_event(start.key, int(lc_idx)) # It's possible to implement more sophisticated logic to remove useless blocks for SWA, e.g. # check if consecutive available blocks is sufficient for window_size. (TRTLLM-8802) # But for simplicity, we leave it for now. - curr = self + curr = start while ( (isinstance(curr, Block) and curr.storage[lc_idx] is None) and not curr.next and curr._prev() is not None ): - if curr.key in curr.prev.next: - curr.prev.next.pop(curr.key) - if event_manager is not None: - event_manager.add_removed_event(curr.key) - curr = curr.prev + prev = curr.prev + detached = detach_next(prev, curr.key) + assert detached is curr + if event_manager is not None: + event_manager.add_removed_event(curr.key) + curr = prev @property def tokens_per_block(self) -> int: @@ -435,7 +470,9 @@ def is_full(self) -> bool: @property def is_orphan(self) -> bool: - return self.key not in self.prev.next or self.prev.next[self.key] is not self + prev = self._prev() + assert prev is None or (self.key in prev.next and prev.next[self.key] is self) + return prev is None class BlockRadixTree: @@ -489,15 +526,16 @@ def event_manager(self) -> "KVCacheEventManager | None": def num_life_cycles(self) -> LifeCycleId: return self.life_cycles.size - def clear(self) -> list[rawref.ref["CommittedPage"]]: + def clear(self) -> None: # taking O(1) space # remove leaf blocks one by one, in post-order - ret: list[rawref.ref["CommittedPage"]] = [] + # ~Block() / __del__() handles page cleanup. + # detach_next() auto-prunes empty RootBlocks from the tree. while self.next: - block = next(iter(self.next.values())) - ret.extend(remove_subtree(block)) + root = next(iter(self.next.values())) + while root.next: + remove_subtree(next(iter(root.next.values()))) assert not self.next - return ret def _num_matched_tokens(self, matched: list[tuple[Block, int]]) -> int: if not matched: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_common.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_common.py index 6e8029acfa33..5782e59df03b 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_common.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_common.py @@ -18,7 +18,7 @@ from dataclasses import dataclass from typing import Final, NewType -NDEBUG: Final[int] = int(os.environ.get("TLLM_KV_CACHE_MANAGER_V2_DEBUG", "0")) == 0 +NDEBUG: Final[bool] = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] != "1" class PageStatus(enum.IntEnum): diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py index 4a0e9a5f7064..8f374e7a37f7 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py @@ -185,8 +185,6 @@ class KVCacheManagerConfig: """ tokens_per_block: int - # if you have p-tuning tokens, include them. Only needed for multi-modal. - vocab_size: int # cache tiers are sorted from warm to cold. The first one must be GPU memory. cache_tiers: list[CacheTierConfig] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py index f46ca567537d..1b9a1a59ed0a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.py @@ -245,7 +245,13 @@ def __enter__(self) -> "StagingBuffer": if self.min_size > manager.size: raise ValueError(f"Requested min_size {self.min_size} is too large for the manager") with manager.mutex: - self._size = min(self.max_size, manager._suggest_next_max_size_unsafe()) + # If the tail cannot satisfy min_size, wrap to the front before allocating. + available = manager._suggest_next_max_size_unsafe() + if self.min_size > available: + manager.next = 0 + available = manager._suggest_next_max_size_unsafe() + assert self.min_size <= available + self._size = max(min(self.max_size, available), self.min_size) self.start_grain = manager.next manager.next += self.num_grains assert manager.next <= manager.num_grains diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py index 09157ba67d36..ea7360d69c83 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py @@ -17,8 +17,11 @@ from ._kv_cache import _KVCache from ._kv_cache_manager import ( AggregatedPageDesc, + ExpandedBuffer, KVCacheManager, PageIndexConverter, + PoolDesc, + PoolGroupDesc, PoolGroupPeakBlockStats, ScratchDesc, ) @@ -29,7 +32,10 @@ "BeamIndex", "DEFAULT_BEAM_INDEX", "AggregatedPageDesc", + "ExpandedBuffer", "PageIndexConverter", + "PoolDesc", + "PoolGroupDesc", "PoolGroupPeakBlockStats", "ScratchDesc", ] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index 191f2398e92a..63e2244d9b98 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -985,6 +985,14 @@ def commit( def num_committed_tokens(self) -> int: return len(self._committed_tokens) + @property + def committed_tokens(self) -> list[TokenIdExt]: + return list(self._committed_tokens) + + @property + def reuse_scope(self) -> ReuseScope: + return self._reuse_scope + # Users promise to not commit any more tokens. For cases where we shouldn't reuse generated tokens # (eg. CoT), this helps us drop (instead of evict) out-of-window blocks for SWA layers. # If there is a uncommitted block containing committed tokens, we will commit the block immediately. @@ -1086,7 +1094,7 @@ def resume(self, cuda_stream: CudaStream | None = None) -> bool: if self._never_resumed and ( type(life_cycles[lc_idx]) is SsmLifeCycle or has_partial ): - deferred_slots[lc_idx] = slot_lst.pop(0) + deferred_slots[lc_idx] = slot_lst.pop() scratch_slots_to_add[lc_idx] = slot_lst stream_wait_events( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py index 546ce0b96cdb..af54f32fe4f7 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py @@ -15,10 +15,10 @@ import time from collections import defaultdict -from collections.abc import Callable, Sequence +from collections.abc import Callable, Iterable, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Iterable, Iterator, cast +from typing import TYPE_CHECKING, Iterator, cast from .. import rawref from .._block_radix_tree import BlockRadixTree, ReuseMatch, ReuseScope @@ -40,7 +40,7 @@ from .._life_cycle_registry import LayerGroupId, LifeCycle, LifeCycleId, LifeCycleRegistry from .._page import Page, _PageHolder from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta -from .._storage._config import BufferId, create_storage_config +from .._storage._config import BufferId, SlotDesc, create_storage_config from .._storage._core import PoolGroupIndex, PoolIndex, SlotId from .._storage_manager import StorageManager from .._utils import ( @@ -64,15 +64,18 @@ @dataclass(slots=True, frozen=True) -class MemoryPoolDesc: - base: MemAddress - page_size: int +class PoolDesc: + pool_index: PoolIndex + base_address: MemAddress + slot_bytes: int @dataclass(slots=True, frozen=True) -class MemoryPoolGroupDesc: - num_pages: int - pools: TypedIndexList[PoolIndex, MemoryPoolDesc] +class PoolGroupDesc: + pool_group_index: PoolGroupIndex + num_slots: int + slot_desc: SlotDesc + pools: TypedIndexList[PoolIndex, PoolDesc] @dataclass(slots=True, frozen=True) @@ -292,12 +295,7 @@ def shutdown(self) -> None: self._storage.destroy() def clear_reusable_blocks(self) -> None: - for ref in self._radix_tree.clear(): - assert unwrap_rawref(ref).status == PageStatus.DROPPABLE - self._storage.exclude_from_eviction(unwrap_rawref(ref)) - for level in self._storage._levels: - for pg_idx in typed_range(level.storage.num_pool_groups): - assert level.controller.num_evictable_pages(pg_idx) == 0 + self._radix_tree.clear() def get_mem_pool_base_address( self, layer_id: LayerId, data_role: DataRole, index_mode: PageIndexMode | None = None @@ -667,54 +665,73 @@ def get_aggregated_pages(self, buffers: Iterable[BufferId]) -> Iterator[Aggregat Returns: A iterator of aggregated buffers. """ - # Group by (life_cycle, pool_index) groups = defaultdict[tuple[LifeCycleId, PoolIndex], list[tuple[Range, ExpandedBuffer]]]( list[tuple[Range, ExpandedBuffer]] ) buffer_attr_map = self._storage._buffer_attr - for b in buffers: - attr = buffer_attr_map[b] - size = attr.size + for buffer in buffers: + attr = buffer_attr_map[buffer] start = attr.offset key = (attr.life_cycle_id, attr.pool_index) - groups[key].append((Range(start, start + size), ExpandedBuffer(b, attr.expansion))) + groups[key].append( + (Range(start, start + attr.size), ExpandedBuffer(buffer, attr.expansion)) + ) storage = self._storage._levels[GPU_LEVEL].storage lc2pg = self._storage._life_cycle_grouping for (lc, pool_idx), group in groups.items(): pg_idx = lc2pg[lc] - # Sort by start offset - group.sort(key=lambda x: x[0].start) - # Merge contiguous + group.sort(key=lambda item: item[0].start) current_start, current_end, current_buffers = ( group[0][0].start, group[0][0].end, [group[0][1]], ) - # cache stride and pool_base for this group stride = storage.slot_size(pg_idx)[pool_idx] pool_base = int(cast(int, storage.slot_address(pg_idx, pool_idx, SlotId(0)))) - for i in range(1, len(group)): - next_range, next_buf = group[i] + for next_range, next_buffer in group[1:]: if next_range.start == current_end: current_end = next_range.end - current_buffers.append(next_buf) - else: - base = MemAddress(pool_base + current_start) - yield AggregatedPageDesc( - base, current_end - current_start, stride, lc, tuple(current_buffers) - ) - current_start, current_end, current_buffers = ( - next_range.start, - next_range.end, - [next_buf], - ) - # Flush last + current_buffers.append(next_buffer) + continue + + base = MemAddress(pool_base + current_start) + yield AggregatedPageDesc( + base, current_end - current_start, stride, lc, tuple(current_buffers) + ) + current_start, current_end, current_buffers = ( + next_range.start, + next_range.end, + [next_buffer], + ) base = MemAddress(pool_base + current_start) yield AggregatedPageDesc( base, current_end - current_start, stride, lc, tuple(current_buffers) ) + @property + def pool_group_descs(self) -> TypedIndexList[PoolGroupIndex, PoolGroupDesc]: + storage = self._storage + + def get_pool_group_desc(pg_idx: PoolGroupIndex) -> PoolGroupDesc: + slot_size_list = storage.slot_size(pg_idx) + pools = make_typed( + lambda pool_idx: PoolDesc( + pool_index=pool_idx, + base_address=storage.get_mem_pool_base_address(pg_idx, pool_idx), + slot_bytes=slot_size_list[pool_idx], + ), + storage.num_pools(pg_idx), + ) + return PoolGroupDesc( + pool_group_index=pg_idx, + num_slots=storage.num_slots(pg_idx), + slot_desc=storage._slot_desc_list[pg_idx], + pools=pools, + ) + + return make_typed(get_pool_group_desc, storage.num_pool_groups) + @property def _current_gpu_ratio(self) -> TypedIndexList[PoolGroupIndex, float]: return self._storage.get_ratio_list(GPU_LEVEL) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py index c308790f9c29..ad1cb3c12012 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py @@ -524,6 +524,12 @@ def _drop_hash_cache(self, block_hash: bytes) -> None: self._v1_hash_compatible_keys.discard(block_hash) self._v1_root_attrs_by_block_key.pop(block_hash, None) + @staticmethod + def _resolve_page_ref(page_ref: Any) -> Any: + if page_ref is None: + return None + return page_ref() if callable(page_ref) else page_ref + @staticmethod def _normalize_token(token: TokenIdExt) -> UniqueToken: if isinstance(token, bytes): @@ -541,7 +547,7 @@ def _stored_block_from_radix_block( continue if page_ref is None: continue - page = page_ref() + page = self._resolve_page_ref(page_ref) if page is None: continue cache_level = page.cache_level @@ -565,7 +571,7 @@ def _life_cycle_ids_from_radix_block(block: Any) -> set[int]: return { life_cycle_id for life_cycle_id, page_ref in enumerate(block.storage) - if page_ref is not None and page_ref() is not None + if page_ref is not None and KVCacheEventManager._resolve_page_ref(page_ref) is not None } def _parent_hash_from_radix_block(self, block: Any) -> EventBlockHash | None: @@ -577,7 +583,7 @@ def _parent_hash_from_radix_block(self, block: Any) -> EventBlockHash | None: def _hash_from_radix_block(self, block: Any) -> EventBlockHash: if self._hash_algo == KV_CACHE_HASH_ALGO_V1: return self._v1_hash_from_radix_block(block) - return self._normalize_block_hash(block.key) + return self._normalize_block_hash(getattr(block, "event_key", block.key)) def _v1_hash_from_radix_block(self, block: Any) -> int: key = bytes(block.key) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py index 169184e2f589..cf1035c221eb 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.py @@ -198,6 +198,8 @@ def evict( ret = make_typed(lambda _: list[EvictablePage](), self.num_pool_groups) try: for pg_idx, count in typed_enumerate(min_num_pages): + if count < 0: + raise ValueError("Eviction count must be non-negative") policy = self._policies[pg_idx] if (len(policy) + len(ret[pg_idx])) < count: raise OutOfPagesError(f"Not enough pages to evict in group {pg_idx}") diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py new file mode 100644 index 000000000000..8b063677111d --- /dev/null +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import sys +from typing import Any + + +def _cpp_introspection_module() -> Any | None: + package = sys.modules.get(__package__) + if package is None: + return None + return getattr(package, "_cpp_introspection", None) + + +def active_page_stats(kv_cache: Any) -> tuple[list[int], list[int]]: + """Return active pages and unscheduled evictable active pages by cache level.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + counts, unscheduled_evictable = cpp_introspection.active_page_stats(kv_cache) + return list(counts), list(unscheduled_evictable) + + storage = kv_cache.manager._storage + counts = [0] * storage.num_cache_levels + unscheduled_evictable = [0] * storage.num_cache_levels + for ordinal, beam_idx, lc_idx in kv_cache._active_pages(): + block_page = kv_cache._page(ordinal, beam_idx, lc_idx) + if block_page is None: + continue + + page = block_page.page + level = page.cache_level + counts[level] += 1 + if storage.is_evictable(page) and not page.scheduled_for_eviction: + unscheduled_evictable[level] += 1 + return counts, unscheduled_evictable + + +def all_tree_pages_droppable(manager: Any) -> bool: + """Return whether every page reachable from the radix tree is droppable.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return bool(cpp_introspection.all_tree_pages_droppable(manager)) + + from ._block_radix_tree import traverse_post_order + from ._common import PageStatus + from ._utils import unwrap_rawref + + for root_block in manager._radix_tree.next.values(): + for block0 in root_block.next.values(): + for block in traverse_post_order(block0): + for page in block.storage: + if page is not None and unwrap_rawref(page).status != PageStatus.DROPPABLE: + return False + return True + + +def is_commit_allowed(kv_cache: Any) -> bool: + """Return whether the KV cache still allows token commits.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return bool(cpp_introspection.is_commit_allowed(kv_cache)) + return kv_cache._commit_state == kv_cache.CommitState.ALLOWED + + +def current_gpu_ratio(manager: Any) -> list[float]: + """Return the current GPU pool-group ratio list.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return list(cpp_introspection.current_gpu_ratio(manager)) + return list(manager._current_gpu_ratio) + + +def set_num_sampled_kv_caches(manager: Any, value: int) -> None: + """Set the sampled-KV-cache counter that gates auto-tuner rebalancing.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + cpp_introspection.set_num_sampled_kv_caches(manager, value) + return + manager._num_sampled_kv_caches = value + + +def set_last_adjustment_time(manager: Any, value: float) -> None: + """Set the last pool-rebalance timestamp that gates the auto-tuner cooldown.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + cpp_introspection.set_last_adjustment_time(manager, value) + return + manager._last_adjustment_time = value + + +def set_target_ratio_list_gpu(manager: Any, ratios: list[float]) -> None: + """Override the target GPU pool-group ratio list (drives the next rebalance).""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + cpp_introspection.set_target_ratio_list_gpu(manager, list(ratios)) + return + manager._target_ratio_list_gpu = list(ratios) + + +def force_rebalance_precondition(manager: Any, skew: float = 2.0) -> None: + """Force the V2 auto-tuner to do real pool-resize work on the next rebalance. + + Bypasses the sample-count / cooldown gates and perturbs the target GPU + ratio so it differs from the current ratio by more than the auto-tuner's + adjustment threshold. Requires a model with >=2 pool groups (e.g. a VSWA + model) and raises ``ValueError`` otherwise, so a future model change can't + silently turn a dependent test into a no-op. Backend-agnostic white-box + hook intended for accuracy tests, not production code. + """ + current = current_gpu_ratio(manager) + if len(current) < 2: + raise ValueError( + f"force_rebalance_precondition requires >=2 pool groups; got {len(current)}. " + "Check that VSWA is actually configured for this model." + ) + set_num_sampled_kv_caches(manager, 2001) + set_last_adjustment_time(manager, 0.0) + skewed = [current[0] * skew] + current[1:] + total = sum(skewed) + set_target_ratio_list_gpu(manager, [x / total for x in skewed]) + + +def storage_statistics(manager: Any, cache_level: int = 0) -> list[Any]: + """Return storage statistics by pool group for a cache level.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return list(cpp_introspection.storage_statistics(manager, cache_level)) + return list(manager._storage.get_statistics(cache_level)) + + +def storage_utilization(manager: Any, cache_level: int = 0) -> list[float]: + """Return storage utilization by pool group for a cache level.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return list(cpp_introspection.storage_utilization(manager, cache_level)) + return list(manager._storage.get_utilization(cache_level)) + + +def grains_for_slots(num_slots: int, slot_size_list: list[int], granularity: int) -> int: + """Return the grain count required for a pool group slot count.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return int(cpp_introspection.grains_for_slots(num_slots, slot_size_list, granularity)) + + from ._storage._core import CacheLevelStorage + + return int(CacheLevelStorage._grains_for_slots(num_slots, slot_size_list, granularity)) + + +def grains_to_slots(pg_grains: int, slot_size_list: list[int], granularity: int) -> tuple[int, int]: + """Return (slot count, consumed grains) for a pool group grain budget.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + slots, used = cpp_introspection.grains_to_slots(pg_grains, slot_size_list, granularity) + return int(slots), int(used) + + from ._storage._core import CacheLevelStorage + + slots, used = CacheLevelStorage._grains_to_slots(pg_grains, slot_size_list, granularity) + return int(slots), int(used) + + +def ratio_to_slot_count_list( + quota: int, + slot_size_lists: list[list[int]], + ratio_list: list[float], + granularity: int, + min_slots: list[int], +) -> list[int]: + """Return slot counts by pool group for a quota and ratio list.""" + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return list( + cpp_introspection.ratio_to_slot_count_list( + quota, slot_size_lists, ratio_list, granularity, min_slots + ) + ) + + from ._storage._core import CacheLevelStorage + + return list( + CacheLevelStorage.ratio_to_slot_count_list( + quota, slot_size_lists, ratio_list, granularity, min_slots + ) + ) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py index cc4867558fc5..43bac9243622 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py @@ -228,6 +228,18 @@ def check_page(p: "BlockPage") -> bool: @dataclass(slots=True) class CommittedPage(Page): + """A committed page is immutable — all access after commit is read-only. + + We intentionally do not add a separate read_event to track read completion. + The inherited Slot.ready_event serves double duty: after commit or migration it + represents write completion; after _UniqPageLock is destroyed it is set to the + merged finish events of all prior readers. This means a new reader may + unnecessarily wait for a prior reader (read-after-read on immutable data), but + this is functionally correct, only occurs when the lock is fully released between + reuses, and saves one event field per committed page — a worthwhile tradeoff given + the potentially huge number of committed pages in the system. + """ + block: rawref.ref["Block"] __rawref__: rawref.ref["CommittedPage"] @@ -262,11 +274,11 @@ def __del__(self) -> None: block = self.block() # block may be None when rebase happens, i.e. another block with the same key is committed, # replacing it, but the page is still used by a _KVCache. - if block is not None: - block.unset_page( + if block is not None and block.unlink_page(self.life_cycle, self): + Block.clear_stale_blocks_after_page_unlink( + block, self.life_cycle, self.manager._life_cycles.get_life_cycle(self.life_cycle), - self, ) Page.__del__(self) self.__rawref__.invalidate() @@ -378,6 +390,9 @@ def __del__(self) -> None: page = self.page if not NDEBUG: assert_critical(page.cache_level == CacheLevel(0) and not page.scheduled_for_eviction) + # Set ready_event to the merged finish events of all readers. For committed (read-only) + # pages, this means the next reader will wait for prior reads to complete, which is + # unnecessary but correct. See the CommittedPage docstring for rationale. page.ready_event = merge_events(self.finish_events) assert self.holder is not None self.holder._lock = None diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py index a02b55dbbf60..292876ff4985 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py @@ -76,3 +76,8 @@ def iter_cache_hit_rate(self) -> float: if self.iter_reused_blocks == 0 or total == 0: return 0.0 return self.iter_reused_blocks / total + + +_KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple( + field.name for field in fields(KVCacheIterationStatsDelta) +) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py index ab33a0eb5dac..e0d7bcc05092 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py @@ -96,7 +96,7 @@ def resize(self, new_num_slots: int) -> None: pass @abc.abstractmethod - def slot_address(self, slot: int) -> Address: + def slot_address(self, slot: SlotId) -> Address: pass def __del__(self) -> None: @@ -136,8 +136,8 @@ def extend_by_one_phys_mem(self) -> int: return self.num_slots @override - def slot_address(self, slot: int) -> MemAddress: - return MemAddress(int(self._vm.address) + self.slot_size * slot) + def slot_address(self, slot: SlotId) -> MemAddress: + return MemAddress(int(self._vm.address) + self.slot_size * int(slot)) @property @override @@ -172,8 +172,8 @@ def resize(self, new_num_slots: int) -> None: self._host_mem.resize(self.aligned_size(new_num_slots)) @override - def slot_address(self, slot: int) -> MemAddress: - return MemAddress(self._host_mem._address + self.slot_size * slot) + def slot_address(self, slot: SlotId) -> MemAddress: + return MemAddress(self._host_mem._address + self.slot_size * int(slot)) @property @override @@ -232,7 +232,7 @@ def resize(self, new_num_slots: int) -> None: resize_file(self.fd, file_size) @override - def slot_address(self, slot: int) -> DiskAddress: + def slot_address(self, slot: SlotId) -> DiskAddress: assert slot < self.num_slots return DiskAddress(self.fd, slot * self.slot_size) @@ -370,6 +370,8 @@ def allocate(self) -> Slot: # and when we don't have enough free slots, we will free these newly allocated slots by appending # them to the back of the recycled slot queue, which may impact perf. def allocate_multiple(self, num_slots: int) -> list[Slot]: + if num_slots < 0: + raise LogicError("SlotAllocator.allocate_multiple: slot count must be non-negative") if self.num_free_slots < num_slots: raise OutOfPagesError("Not enough free slots") return [self.allocate() for _ in range(num_slots)] @@ -395,8 +397,7 @@ def num_slots(self) -> int: def expand(self, new_num_slots: int) -> None: assert NDEBUG or self._check() assert self._target_capacity == self._capacity - old_num_slots = self._capacity - assert new_num_slots > old_num_slots + assert new_num_slots > self._capacity self._occupied_mask.resize(new_num_slots) self._capacity = new_num_slots self._target_capacity = self._capacity @@ -511,7 +512,7 @@ def destroy(self) -> None: if self._destroyed: return allocator = self._slot_allocator - if allocator._capacity != 0: + if allocator.num_slots != 0: allocator._synchronize() allocator.prepare_for_shrink(0) allocator.finish_shrink() @@ -525,7 +526,7 @@ def num_pools(self) -> PoolIndex: @property def num_slots(self) -> int: - num_slots = self._slot_allocator._capacity + num_slots = self._slot_allocator.num_slots assert num_slots <= self._get_num_slots_from_pools() return num_slots @@ -604,7 +605,10 @@ def __init__( slot_size_list, lambda slot_size: GpuSlotPool( slot_size, - round_down(int(total_gpu_memory * slot_size / max_slot_size), phys_mem_size), + max( + round_down(int(total_gpu_memory * slot_size / max_slot_size), phys_mem_size), + round_up(num_slots * slot_size, phys_mem_size), + ), shared_phys_mem_pool, num_slots, ), @@ -763,7 +767,7 @@ def _grains_to_slots( min_pool_grains = typed_map(slot_size_list, lambda s: div_up(s, granularity)) if pg_grains < sum(min_pool_grains): return (0, 0) - num_slots: int = 1 << 63 + num_slots = 1 << 63 remaining_pg_grains = pg_grains pool_idx_lst = sorted(typed_range(num_pools), key=lambda i: slot_size_list[i]) for j, pool in enumerate(pool_idx_lst): 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 540c8bd76df3..765d3bdca068 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 OutOfPagesError +from ._exceptions import LogicError, OutOfPagesError from ._life_cycle_registry import ( AttnLifeCycle, LifeCycleId, @@ -327,6 +327,8 @@ def new_slots( lc2pg = self._life_cycle_grouping pg_num_slots = filled_list(0, self.num_pool_groups) for lc in typed_range(self.num_life_cycles): + if num_slots[lc] < 0: + raise LogicError("StorageManager.new_slots: slot count must be non-negative") pg_num_slots[lc2pg[lc]] += num_slots[lc] storage = self._levels[level].storage if any( @@ -360,6 +362,10 @@ def new_slots_for_pool_group( migration_recorder: MigrationRecorder | None = None, drop_recorder: DropRecorder | None = None, ) -> list[Slot]: + if num_slots < 0: + raise LogicError( + "StorageManager.new_slots_for_pool_group: slot count must be non-negative" + ) storage = self._levels[level].storage if num_slots > storage.get_num_free_slots(pg_idx): num_slots_list = filled_list(0, self.num_pool_groups) diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py index 95eabb212ee4..da750eeca0e4 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py @@ -123,6 +123,14 @@ def beg(self) -> Idx: def end(self) -> Idx: return self[1] + def __eq__(self, other: object) -> bool: + if not isinstance(other, HalfOpenRange): + return NotImplemented + return (not self and not other) or tuple.__eq__(self, other) + + def __hash__(self) -> int: + return hash((0, 0)) if not self else tuple.__hash__(self) + def __bool__(self) -> bool: return self[0] < self[1] diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py b/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py index f2f5d9e31c12..4e40456a4dbe 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py @@ -78,6 +78,7 @@ "kv_cache_manager_v2/_cuda_virt_mem.py", "kv_cache_manager_v2/_event_manager.py", "kv_cache_manager_v2/_exceptions.py", + "kv_cache_manager_v2/_introspection.py", "kv_cache_manager_v2/_life_cycle_registry.py", "kv_cache_manager_v2/_page.py", "kv_cache_manager_v2/_storage_manager.py", diff --git a/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py b/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py index 121a2bc65bac..6542c7cdfae0 100644 --- a/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py +++ b/tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py @@ -38,31 +38,17 @@ def _inject_pool_ratio_mismatch(llm: LLM, *, skew: float = 2.0) -> None: """Force the V2 auto-tuner to do real pool-resize work on the next rebalance call. - Bypasses the 2000-sample / 120s cooldown gates by stomping counters, - then perturbs ``_target_ratio_list_gpu`` so it differs from - ``_current_gpu_ratio`` by more than the 1.25x threshold inside - ``_need_adjustment``. - - Requires a model with >=2 pool groups (e.g. Gemma-3-1B with VSWA). - Asserts the precondition so a future model change can't silently - turn this test into a no-op. + Delegates to the backend-agnostic KVCacheManagerV2 introspection hook, which + bypasses the sample-count / cooldown gates and perturbs the target GPU ratio + past the auto-tuner's adjustment threshold. The hook requires a model with + >=2 pool groups (e.g. Gemma-3-1B with VSWA) and raises otherwise, so a future + model change can't silently turn this test into a no-op. """ + from tensorrt_llm.runtime.kv_cache_manager_v2 import _introspection + executor = llm._executor.engine kv_cache_manager = executor.kv_cache_manager - impl = kv_cache_manager.impl - - impl._num_sampled_kv_caches = 2001 - impl._last_adjustment_time = 0.0 - - current = list(impl._current_gpu_ratio) - assert len(current) >= 2, ( - f"Ratio injection requires >=2 pool groups; got {len(current)}. " - "Check that VSWA is actually configured for this model." - ) - - skewed = [current[0] * skew] + list(current[1:]) - total = sum(skewed) - impl._target_ratio_list_gpu = [x / total for x in skewed] + _introspection.force_rebalance_precondition(kv_cache_manager.impl, skew=skew) # --------------------------------------------------------------------------- # diff --git a/tests/integration/defs/perf/create_perf_comparison_report.py b/tests/integration/defs/perf/create_perf_comparison_report.py index d36ec0a23dca..d9c66ff58fbc 100644 --- a/tests/integration/defs/perf/create_perf_comparison_report.py +++ b/tests/integration/defs/perf/create_perf_comparison_report.py @@ -130,7 +130,16 @@ def parse_perf_data( def generate_perf_compare_report(perf_files: list[str], output_path: str) -> None: - name_mapping, merged, suffixes = parse_perf_data(perf_files) + # When all perf tests in a stage are skipped/waived (e.g. an nvbugs waive) + # or reused from a previous pipeline, no perf CSV is produced. Skip report + # generation gracefully instead of raising FileNotFoundError, mirroring the + # missing-CSV handling in sanity_perf_check.py. + existing_files = [f for f in (perf_files or []) if Path(f).exists()] + if not existing_files: + print("No perf CSV files found (perf tests skipped/waived); " + "skipping perf comparison report generation.") + return + name_mapping, merged, suffixes = parse_perf_data(existing_files) generate_plots(Path(output_path), name_mapping, merged, suffixes) 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 59112a9c7790..4d8dc494a898 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 @@ -221,7 +221,7 @@ def test_deepseek_v4_pool_ratio_overrides_typical_step_and_constraints(): KvCacheConfig(pool_ratio=[0.2, 0.3, 0.5], avg_seq_len=256) ) - assert config.initial_pool_ratio == [0.2, 0.3, 0.5] + assert config.initial_pool_ratio == pytest.approx([0.2, 0.3, 0.5]) assert config.typical_step is None assert config.constraints == [] diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 2f4d697f6364..76a0d8db25b3 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -510,7 +510,15 @@ def get_block_data( """Unified block data retrieval for both V1 and V2 KVCacheManager.""" if use_v2: layer_grouping = kv_cache_manager.impl.layer_grouping - local_layer_indices = layer_grouping[layer_group_id] + # Read layers in ascending global-layer order so this verification does + # not depend on the KV-cache manager's internal layer_grouping order + # (an implementation detail, not an API contract). The merge below packs + # ranks at ascending layer offsets, so the per-rank stack must also be + # ascending by global layer. + local_layer_indices = sorted( + layer_grouping[layer_group_id], + key=lambda lid: kv_cache_manager.pp_layers[lid], + ) all_layer_data = [] for local_layer_idx in local_layer_indices: diff --git a/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py b/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py index 2f3c8fb80460..70a4e48d6c85 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py +++ b/tests/unittest/kv_cache_manager_v2_tests/fake_engine.py @@ -161,7 +161,7 @@ def _check_pages( ) pool = manager.get_mem_pool_base_address(layer_id, role, index_mode) stride = manager.get_page_stride(layer_id, role) - lc_id = manager._storage._layer_to_life_cycle_ids[layer_id] + lc_id = manager.get_layer_group_id(layer_id) if is_ssm: # SSM: only one page (the SSM slot), only check the last history token if not history: @@ -237,7 +237,7 @@ def _write_new_tokens( ) pool = manager.get_mem_pool_base_address(layer_id, role, index_mode) stride = manager.get_page_stride(layer_id, role) - lc_id = manager._storage._layer_to_life_cycle_ids[layer_id] + lc_id = manager.get_layer_group_id(layer_id) if is_ssm: # SSM: write only the last input token at position 0 of the SSM page ssm_idx = kv_cache.get_ssm_block_base_index(lc_id, beam) diff --git a/tests/unittest/kv_cache_manager_v2_tests/kernels.py b/tests/unittest/kv_cache_manager_v2_tests/kernels.py index 9f69a185f714..d5955026a52f 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/kernels.py +++ b/tests/unittest/kv_cache_manager_v2_tests/kernels.py @@ -73,6 +73,7 @@ def get_program(debug: bool, max_tokens: int) -> ObjectCode: using uint32_t = unsigned int; using uint16_t = unsigned short; +using uint64_t = unsigned long long; struct alignas(16) Value { uint32_t token; @@ -143,13 +144,23 @@ def get_program(debug: bool, max_tokens: int) -> ObjectCode: __nanosleep(sleepTime); } } + +// Spin until the host sets *flag to non-zero (or a bail-out iteration budget is +// exhausted so a buggy/leaked gate cannot hang the GPU forever). Used to keep a +// stream deterministically busy from the host side in ordering tests. +extern "C" __global__ void spinUntilFlag(volatile uint32_t const* flag, uint64_t maxIters) +{ + for (uint64_t i = 0; *flag == 0u && i < maxIters; ++i) { + __nanosleep(1'000'000u); + } +} """ macros = [("MAX_TOKENS", str(max_tokens))] program_options = ProgramOptions(std="c++17", lineinfo=True, debug=debug, define_macro=macros) # type: ignore[arg-type] if not debug: program_options.use_fast_math = True prog = Program(code, code_type="c++", options=program_options) - mod = prog.compile("cubin", name_expressions=("fillValues", "checkValues")) + mod = prog.compile("cubin", name_expressions=("fillValues", "checkValues", "spinUntilFlag")) return mod @@ -158,7 +169,7 @@ def get_kernel(name: str, num_tokens: int) -> tuple[Kernel, int]: @lru_cache(maxsize=None) def impl(name: str, max_tokens: int) -> Kernel: - assert name in ("fillValues", "checkValues") + assert name in ("fillValues", "checkValues", "spinUntilFlag") assert max_tokens != 0 and (max_tokens & (max_tokens - 1)) == 0, ( "max_tokens must be a power of 2" ) @@ -355,3 +366,39 @@ def debug_dump_tokens( for j in range(1, values_per_token): assert token[j] == token[0] yield value + + +class HostGate: + """Deterministically blocks CUDA streams until the host calls release(). + + Enqueues a device-side spin kernel that polls a host-mapped flag. Work + enqueued on a gated stream after block_stream() cannot start until + release() is called (or the kernel's bail-out iteration budget expires, + which makes dependent ordering assertions fail loudly instead of hanging). + """ + + def __init__(self) -> None: + result, ptr = drv.cuMemHostAlloc(4, drv.CU_MEMHOSTALLOC_DEVICEMAP) + _unwrap(result) + self._host_ptr = int(ptr) + ctypes.c_uint32.from_address(self._host_ptr).value = 0 + result, device_ptr = drv.cuMemHostGetDevicePointer(self._host_ptr, 0) + _unwrap(result) + self._device_ptr = int(device_ptr) + + def block_stream(self, stream: CudaStream, max_iters: int = 100_000) -> None: + kernel, _ = get_kernel("spinUntilFlag", 1) + args = (self._device_ptr, max_iters) + arg_types = (ctypes.c_void_p, ctypes.c_uint64) + _unwrap( + drv.cuLaunchKernel(kernel._handle, 1, 1, 1, 1, 1, 1, 0, stream, (args, arg_types), 0) + ) + + def release(self) -> None: + ctypes.c_uint32.from_address(self._host_ptr).value = 1 + + def close(self) -> None: + if self._host_ptr: + self.release() + _unwrap(drv.cuMemFreeHost(self._host_ptr)) + self._host_ptr = 0 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py b/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py index d475ea791060..489896279247 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.py @@ -84,8 +84,8 @@ def tok(self) -> TokenIdExt: return TokenId(next(self._tok)) def _page_addr(self, kv, block_ordinal): - lc = self.manager._storage._layer_to_life_cycle_ids[LID] - base_idx = list(kv.get_base_page_indices(lc, DEFAULT_BEAM_INDEX))[block_ordinal] + lc = self.manager.get_layer_group_id(LID) + base_idx = int(list(kv.get_base_page_indices(lc, DEFAULT_BEAM_INDEX))[block_ordinal]) pool = self.manager.get_mem_pool_base_address(LID, KEY) stride = self.manager.get_page_stride(LID, KEY) scale = self.manager.get_page_index_scale(LID, KEY) 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 8c8daa74b115..fedcf41450a5 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 @@ -26,7 +26,7 @@ from importlib.util import find_spec from random import randbytes from statistics import median -from typing import TYPE_CHECKING, Iterator, NamedTuple, cast +from typing import TYPE_CHECKING, Iterator, NamedTuple, cast, get_type_hints if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: from kv_cache_manager_v2 import ( @@ -51,23 +51,21 @@ SwaScratchReuseConfig, TokenId, TokenIdExt, + _introspection, _KVCache, ) - from kv_cache_manager_v2._block_radix_tree import Hasher, traverse_post_order + from kv_cache_manager_v2._block_radix_tree import Hasher from kv_cache_manager_v2._common import ( BAD_PAGE_INDEX, GPU_LEVEL, CacheTier, MemAddress, PageIndexMode, - PageStatus, SlidingWindowSize, ) from kv_cache_manager_v2._copy_engine import CopyTask, batched_copy from kv_cache_manager_v2._exceptions import OutOfPagesError - from kv_cache_manager_v2._life_cycle_registry import LifeCycleRegistry, SsmLifeCycle - from kv_cache_manager_v2._storage._config import create_storage_config - from kv_cache_manager_v2._storage._core import CacheLevelStorage, SlotAllocator + from kv_cache_manager_v2._storage._core import CacheLevelStorage, PoolGroupBase, SlotAllocator from kv_cache_manager_v2._storage_manager import StorageManager from kv_cache_manager_v2._utils import ( CachedCudaStream, @@ -106,30 +104,23 @@ SwaScratchReuseConfig, TokenId, TokenIdExt, + _introspection, _KVCache, ) - from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import ( - Hasher, - traverse_post_order, - ) + from tensorrt_llm.runtime.kv_cache_manager_v2._block_radix_tree import Hasher from tensorrt_llm.runtime.kv_cache_manager_v2._common import ( BAD_PAGE_INDEX, GPU_LEVEL, CacheTier, MemAddress, PageIndexMode, - PageStatus, SlidingWindowSize, ) from tensorrt_llm.runtime.kv_cache_manager_v2._copy_engine import CopyTask, batched_copy from tensorrt_llm.runtime.kv_cache_manager_v2._exceptions import OutOfPagesError - from tensorrt_llm.runtime.kv_cache_manager_v2._life_cycle_registry import ( - LifeCycleRegistry, - SsmLifeCycle, - ) - from tensorrt_llm.runtime.kv_cache_manager_v2._storage._config import create_storage_config from tensorrt_llm.runtime.kv_cache_manager_v2._storage._core import ( CacheLevelStorage, + PoolGroupBase, SlotAllocator, ) from tensorrt_llm.runtime.kv_cache_manager_v2._storage_manager import StorageManager @@ -154,7 +145,18 @@ with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): from fake_engine import FakeEngine, Role, Step - from kernels import enable_kernel_delay + from kernels import HostGate, enable_kernel_delay + + +def get_cached_cuda_event_type(): + if find_spec("kv_cache_manager_v2") is not None: + from kv_cache_manager_v2._utils import CachedCudaEvent + + return CachedCudaEvent + from tensorrt_llm.runtime.kv_cache_manager_v2._utils import CachedCudaEvent + + return CachedCudaEvent + seed = int.from_bytes(os.urandom(8), "little") print(f"seed: {seed}") @@ -205,14 +207,36 @@ def wrapper(self, *args, **kwargs): return wrapper +class TestTypedSlotIds(unittest.TestCase): + def test_num_slots_accessors_return_int(self) -> None: + self.assertIs(get_type_hints(SlotAllocator.num_slots.fget)["return"], int) + self.assertIs(get_type_hints(SlotAllocator.num_free_slots.fget)["return"], int) + self.assertIs(get_type_hints(SlotAllocator.num_occupied_slots.fget)["return"], int) + self.assertIs(get_type_hints(PoolGroupBase.num_slots.fget)["return"], int) + self.assertIs(get_type_hints(PoolGroupBase.num_free_slots.fget)["return"], int) + self.assertIs(get_type_hints(CacheLevelStorage.num_slots)["return"], int) + self.assertIs(get_type_hints(CacheLevelStorage.get_num_free_slots)["return"], int) + self.assertIs(get_type_hints(StorageManager.num_slots)["return"], int) + + self.assertIs(get_type_hints(SlotAllocator.allocate_multiple)["num_slots"], int) + self.assertIs(get_type_hints(PoolGroupBase.allocate_multiple)["num_slots"], int) + self.assertIs(get_type_hints(CacheLevelStorage.allocate_multiple)["num_slots"], int) + self.assertIs(get_type_hints(StorageManager.new_slots_for_pool_group)["num_slots"], int) + + allocator = SlotAllocator(3) + self.assertEqual(allocator.num_slots, 3) + self.assertEqual(allocator.num_free_slots, 3) + self.assertEqual(allocator.num_occupied_slots, 0) + + class TestCacheLevelStorage(unittest.TestCase): def test_grains_to_slots_refines_proportional_lower_bound(self) -> None: granularity = 16 << 20 slot_size_list = [16_252_928, 4_063_232] min_slots = 157 - grains = CacheLevelStorage._grains_for_slots(min_slots, slot_size_list, granularity) - slots, used = CacheLevelStorage._grains_to_slots(grains, slot_size_list, granularity) + grains = _introspection.grains_for_slots(min_slots, slot_size_list, granularity) + slots, used = _introspection.grains_to_slots(grains, slot_size_list, granularity) self.assertGreaterEqual(slots, min_slots) self.assertLessEqual(used, grains) @@ -226,11 +250,11 @@ def test_ratio_to_slot_count_list_preserves_min_slots(self) -> None: ] min_slots = [157, 3907, 157] total_min_grains = sum( - CacheLevelStorage._grains_for_slots(slots, sizes, granularity) + _introspection.grains_for_slots(slots, sizes, granularity) for slots, sizes in zip(min_slots, slot_size_lists) ) - slot_counts = CacheLevelStorage.ratio_to_slot_count_list( + slot_counts = _introspection.ratio_to_slot_count_list( total_min_grains * granularity, slot_size_lists, [0.2, 0.5, 0.3], @@ -275,7 +299,6 @@ def create_config( cache_tiers = [t for t in cache_tiers if t.quota > 0] return KVCacheManagerConfig( tokens_per_block=tokens_per_block, - vocab_size=4096, cache_tiers=[t for t in cache_tiers if t.quota > 0], layers=[ AttentionLayerConfig( @@ -504,6 +527,76 @@ def test_sol_mem_utilization(self) -> None: # This also tests eviction to disk. self.assertRaises(OutOfPagesError, lambda: self.run_naive(seq_len + 1, 1, False)) + def test_resume_rejects_if_any_pool_group_exceeds_threshold(self) -> None: + cfg = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=4 << 20)], + max_util_for_resume=0.9, + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=(1 << 20) + 1)], + sliding_window_size=32, + num_sink_tokens=0, + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[BufferConfig(role=Role.KEY, size=1024)], + sliding_window_size=None, + ), + ], + typical_step=BatchDesc(kv_caches=[KVCacheDesc(capacity=32, history_length=0)]), + constraints=[BatchDesc(kv_caches=[KVCacheDesc(capacity=32, history_length=0)])], + ) + self.manager = KVCacheManager(cfg) + + def stat_slot_sizes(stat) -> list[int]: + if hasattr(stat, "slot_sizes"): + return stat.slot_sizes + return stat.slot_size + + def overall_utilization() -> float: + numerator = 0 + denominator = 0 + for stat in _introspection.storage_statistics(self.manager): + slot_size = sum(stat_slot_sizes(stat)) + numerator += slot_size * stat.unavailable + denominator += slot_size * stat.total + return numerator / denominator + + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prior_caches: list[_KVCache] = [] + try: + # The worst-case SWA slot reservation means a single sequence cannot + # push a pool group past max_util_for_resume, so resume sequences until + # the big-slot SWA pool group crosses the threshold. + for _ in range(64): + if max(_introspection.storage_utilization(self.manager, GPU_LEVEL)) > ( + cfg.max_util_for_resume + ): + break + kv_cache = self.manager.create_kv_cache() + if not kv_cache.resume(stream): + kv_cache.close() + break + self.assertTrue(kv_cache.resize(cfg.tokens_per_block)) + prior_caches.append(kv_cache) + + utilizations = _introspection.storage_utilization(self.manager, GPU_LEVEL) + self.assertGreater(max(utilizations), cfg.max_util_for_resume) + self.assertLess(overall_utilization(), cfg.max_util_for_resume) + + # One pool group is now over the limit, so a further resume is rejected. + rejected_cache = self.manager.create_kv_cache() + prior_caches.append(rejected_cache) + self.assertFalse(rejected_cache.resume(stream)) + self.assertEqual(rejected_cache.status, _KVCache.Status.SUSPENDED) + finally: + for kv_cache in prior_caches: + if kv_cache.status != _KVCache.Status.CLOSED: + kv_cache.close() + @parameterized.expand([(1,), (2,), (4,)]) # @assert_no_ref_cycle def test_cache_reuse(self, num_reusable_requests: int) -> None: @@ -526,15 +619,10 @@ def test_cache_reuse(self, num_reusable_requests: int) -> None: req.kv_cache.close() s.take_finish_event() - for root_block in self.manager._radix_tree.next.values(): - for block0 in root_block.next.values(): - for block in traverse_post_order(block0): - for page in block.storage: - if page is not None: - assert unwrap_rawref(page).status == PageStatus.DROPPABLE + self.assertTrue(_introspection.all_tree_pages_droppable(self.manager)) req0 = reusable_requests[0] - prompt1 = req0.kv_cache._committed_tokens[: (seq_len // 2 - 7)] + prompt1 = req0.kv_cache.committed_tokens[: (seq_len // 2 - 7)] # request id must be same as req0 because we wrote it into the kv cache. req1 = self.Request( next(req_id_gen), @@ -559,7 +647,6 @@ def test_commit_min_snapshot_reuses_swa_post_commit_prefix(self) -> None: prompt = [TokenId(i) for i in range(tokens_per_block * 4)] cfg = KVCacheManagerConfig( tokens_per_block=tokens_per_block, - vocab_size=4096, cache_tiers=[GpuCacheTierConfig(quota=16 << 20)], layers=[ AttentionLayerConfig( @@ -617,7 +704,7 @@ def test_reuse_scope_isolates_reuse(self) -> None: def commit_for(reuse_scope: ReuseScope | None) -> None: kv_cache = self.manager.create_kv_cache(reuse_scope, tokens[:-1]) - self.assertEqual(kv_cache._reuse_scope, reuse_scope or default_scope) + self.assertEqual(kv_cache.reuse_scope, reuse_scope or default_scope) with TemporaryCudaStream([]) as stream_holder: stream = cast(CudaStream, stream_holder.handle) self.assertTrue(kv_cache.resume(stream)) @@ -626,12 +713,13 @@ def commit_for(reuse_scope: ReuseScope | None) -> None: if uncommitted: kv_cache.commit(uncommitted) kv_cache.stop_committing() + stream_holder.take_finish_event() kv_cache.close() def num_reused(reuse_scope: ReuseScope | None) -> int: probed = self.manager.probe_reuse(reuse_scope, tokens[:-1]) kv_cache = self.manager.create_kv_cache(reuse_scope, tokens[:-1]) - self.assertEqual(kv_cache._reuse_scope, reuse_scope or default_scope) + self.assertEqual(kv_cache.reuse_scope, reuse_scope or default_scope) ret = kv_cache.num_committed_tokens kv_cache.close() self.assertEqual(probed, ret) @@ -649,6 +737,140 @@ def num_reused(reuse_scope: ReuseScope | None) -> int: self.assertGreater(num_reused(None), 0) self.assertGreater(num_reused(default_scope), 0) + def test_create_kv_cache_accepts_sequence_input_tokens(self) -> None: + self.prepare(8 << 20, 0, 0, 2, None, 0, tokens_per_block=4, kv_buf_size=1024) + prompt = [self.next_token() for _ in range(8)] + + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + kv_cache = self.manager.create_kv_cache() + try: + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(len(prompt), len(prompt))) + kv_cache.commit(prompt) + kv_cache.stop_committing() + finally: + if kv_cache.status != _KVCache.Status.CLOSED: + kv_cache.close() + s.take_finish_event() + + kv_cache = self.manager.create_kv_cache(input_tokens=tuple(prompt)) + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + try: + self.assertTrue(kv_cache.resume(stream)) + self.assertEqual(kv_cache.num_committed_tokens, len(prompt)) + finally: + if kv_cache.status != _KVCache.Status.CLOSED: + kv_cache.close() + s.take_finish_event() + + def test_create_kv_cache_custom_priority_callback_gets_lifecycle(self) -> None: + self.prepare(8 << 20, 0, 0, 2, 128, 0, tokens_per_block=4, kv_buf_size=1024) + seen_life_cycles = [] + + def custom_priority_callback(_ordinal, life_cycle): + seen_life_cycles.append(life_cycle) + self.assertNotIsInstance(life_cycle, int) + self.assertTrue(hasattr(life_cycle, "get_stale_range")) + return 42 + + kv_cache = self.manager.create_kv_cache(custom_priority_callback=custom_priority_callback) + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + try: + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(4, 4)) + finally: + if kv_cache.status != _KVCache.Status.CLOSED: + kv_cache.close() + s.take_finish_event() + + self.assertTrue(seen_life_cycles) + + def test_cached_cuda_event_constructor_and_null(self) -> None: + cached_cuda_event = get_cached_cuda_event_type() + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + event = cached_cuda_event(stream) + self.assertFalse(event.is_closed()) + event.wait_in_stream(stream) + event.synchronize() + self.assertTrue(event.is_closed()) + + null_event = cached_cuda_event.NULL + self.assertTrue(null_event.is_closed()) + self.assertTrue(null_event.query_complete()) + null_event.synchronize() + null_event.wait_in_stream(stream) + s.take_finish_event() + + def test_base_page_index_external_buffer_validation(self) -> None: + self.prepare(8 << 20, 0, 0, 2, None, 0, tokens_per_block=4, kv_buf_size=1024) + kv_cache = self.manager.create_kv_cache() + try: + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(8)) + num_blocks = kv_cache.num_blocks + + undersized = array.array("i", [BAD_PAGE_INDEX]) * (num_blocks - 1) + with self.assertRaises((AssertionError, ValueError)): + kv_cache.set_base_page_index_buf( + DEFAULT_BEAM_INDEX, LayerGroupId(0), memoryview(undersized) + ) + + oversized = array.array("i", [123]) * (num_blocks + 2) + kv_cache.set_base_page_index_buf( + DEFAULT_BEAM_INDEX, LayerGroupId(0), memoryview(oversized) + ) + self.assertEqual(list(oversized[num_blocks:]), [BAD_PAGE_INDEX, BAD_PAGE_INDEX]) + kv_cache.close() + s.take_finish_event() + finally: + if kv_cache.status != _KVCache.Status.CLOSED: + kv_cache.close() + + def test_buffer_id_tuple_hash_protocol(self) -> None: + buffer_id = BufferId(LayerId(1), Role.KEY) + same_buffer_id = BufferId(LayerId(1), Role.KEY) + as_tuple = (LayerId(1), Role.KEY) + + self.assertEqual(tuple(buffer_id), as_tuple) + self.assertEqual(buffer_id[0], as_tuple[0]) + self.assertEqual(buffer_id[-1], as_tuple[1]) + self.assertEqual(len(buffer_id), 2) + self.assertEqual(buffer_id, as_tuple) + self.assertEqual(as_tuple, buffer_id) + self.assertEqual(buffer_id, same_buffer_id) + self.assertEqual(hash(buffer_id), hash(as_tuple)) + self.assertEqual({buffer_id: 7}[same_buffer_id], 7) + self.assertEqual({buffer_id: 7}[as_tuple], 7) + with self.assertRaises(AttributeError): + buffer_id.layer_id = LayerId(2) + + def test_shrink_capacity_truncates_base_page_indices(self) -> None: + self.prepare(8 << 20, 0, 0, 2, None, 0, tokens_per_block=4, kv_buf_size=1024) + kv_cache = self.manager.create_kv_cache() + layer_group = self.manager.get_layer_group_id(LayerId(0)) + + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + try: + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(8, 0)) + self.assertEqual(kv_cache.num_blocks, 2) + self.assertEqual(len(kv_cache.get_base_page_indices(layer_group)), 2) + + self.assertTrue(kv_cache.resize(4, 0)) + self.assertEqual(kv_cache.num_blocks, 1) + self.assertEqual(len(kv_cache.get_base_page_indices(layer_group)), 1) + finally: + if kv_cache.status != _KVCache.Status.CLOSED: + kv_cache.close() + s.take_finish_event() + @parameterized.expand(list(itertools.product([False, True], repeat=3))) # @assert_no_ref_cycle def test_naive( @@ -789,7 +1011,7 @@ def update_batch(self, stream: CudaStream) -> None: for kv_cache, _, _ in removed: seq_len = self.seq_len_dict[kv_cache] if seq_len < self.avg_length * 3: - self.past_sequences.append(kv_cache._committed_tokens[:seq_len]) + self.past_sequences.append(kv_cache.committed_tokens[:seq_len]) kv_cache.close() self.seq_len_dict.pop(kv_cache) self.num_finished += 1 @@ -821,11 +1043,7 @@ def update_batch(self, stream: CudaStream) -> None: step = suspended[-1] kv_cache = step.kv_cache ok = kv_cache.resume(stream) - if ( - ok - and not self.enable_reuse - and kv_cache._commit_state == _KVCache.CommitState.ALLOWED - ): + if ok and not self.enable_reuse and _introspection.is_commit_allowed(kv_cache): kv_cache.stop_committing() ok = ok and kv_cache.resize(len(step.history) + len(step.input), None) if ok: @@ -1256,7 +1474,6 @@ def test_complex_model_0(self) -> None: config = KVCacheManagerConfig( tokens_per_block=128, - vocab_size=1024, cache_tiers=[ GpuCacheTierConfig(quota=1024 * 1024 * 1024), HostCacheTierConfig(quota=8000 << 20), @@ -1343,7 +1560,6 @@ def test_complex_model_1(self) -> None: config = KVCacheManagerConfig( tokens_per_block=128, - vocab_size=129280, cache_tiers=[GpuCacheTierConfig(quota=212549334)], layers=layers, typical_step=typical_step, @@ -1364,22 +1580,6 @@ def test_resize_quota(self) -> None: stream_holder = CachedCudaStream() stream = cast(CudaStream, stream_holder.handle) - def count_active_pages_by_level(kv_cache: _KVCache) -> list[int]: - counts = [0] * self.manager._storage.num_cache_levels - for ordinal, beam_idx, lc_idx in kv_cache._active_pages(): - block_page = kv_cache._page(ordinal, beam_idx, lc_idx) - assert block_page is not None - counts[block_page.page.cache_level] += 1 - return counts - - def assert_prefetched_pages_are_evictable(kv_cache: _KVCache) -> None: - for ordinal, beam_idx, lc_idx in kv_cache._active_pages(): - block_page = kv_cache._page(ordinal, beam_idx, lc_idx) - assert block_page is not None - page = block_page.page - if page.cache_level == HOST_LEVEL and self.manager._storage.is_evictable(page): - self.assertTrue(page.scheduled_for_eviction) - # First commit some blocks to fill all levels of cache. This helps test the case where shrinking # the quota will drop some pages from the last-level cache. for _ in range(11): @@ -1437,18 +1637,21 @@ def assert_prefetched_pages_are_evictable(kv_cache: _KVCache) -> None: success = self.manager.resize(HOST_LEVEL, 128 << 20) assert success prefetch_target = kv_cache_lst[1] - prefetch_counts_before = count_active_pages_by_level(prefetch_target) + # _introspection.active_page_stats returns (active counts, unscheduled evictable counts) by cache level. + prefetch_counts_before, _ = _introspection.active_page_stats(prefetch_target) self.assertGreater(prefetch_counts_before[DISK_LEVEL], 0) success = prefetch_target.prefetch(HOST_LEVEL) self.assertEqual(success, True) - prefetch_counts_after = count_active_pages_by_level(prefetch_target) + prefetch_counts_after, unscheduled_evictable_after = _introspection.active_page_stats( + prefetch_target + ) self.assertEqual(prefetch_counts_after[GPU_LEVEL], prefetch_counts_before[GPU_LEVEL]) self.assertEqual(prefetch_counts_after[DISK_LEVEL], 0) self.assertEqual( prefetch_counts_after[HOST_LEVEL], prefetch_counts_before[HOST_LEVEL] + prefetch_counts_before[DISK_LEVEL], ) - assert_prefetched_pages_are_evictable(prefetch_target) + self.assertEqual(unscheduled_evictable_after[HOST_LEVEL], 0) # Now both requests can resume for kv_cache in kv_cache_lst: success = kv_cache.resume(stream) @@ -1479,7 +1682,6 @@ def test_hetero_tokens_per_block(self) -> None: ] self.cfg = KVCacheManagerConfig( tokens_per_block=128, - vocab_size=1024, cache_tiers=[ GpuCacheTierConfig(quota=256 << 20), HostCacheTierConfig(quota=1 << 30), @@ -1521,6 +1723,122 @@ def test_hetero_tokens_per_block(self) -> None: kv_cache.close() +class TestKVCacheReusePerformance(TestKVCacheManagerV2): + """Test class for measuring KV cache reuse performance.""" + + def test_cache_reuse_performance(self, profile: bool = False) -> None: + """Performance test for KV cache reuse (prefill only). + + - First pass: 20 requests with 1000 tokens per prompt (cold cache). + - Second pass: Re-run the same 20 requests to achieve 100% cache hit rate. + """ + self.prepare( + gpu_quota=512 << 20, + host_quota=512 << 20, + disk_quota=1 << 30, + num_layers=36, + window_size=None, + sink_tokens=0, + tokens_per_block=32, + kv_buf_size=8192, + ) + + num_requests = 20 + prompt_len = 1000 + + prompts = [] + for _ in range(num_requests): + prompt = [self.next_token() for _ in range(prompt_len)] + prompts.append(prompt) + + def run_requests(prompts: list[list[TokenIdExt]]) -> dict: + """Run all requests (prefill only) and return performance metrics.""" + results = { + "total_time": 0.0, + "num_reused_tokens": 0, + "num_computed_tokens": 0, + } + + tic_total = time.perf_counter() + + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + + requests = [] + + for req_id, prompt in enumerate(prompts): + kv_cache = self.manager.create_kv_cache(None, prompt) + num_reused = kv_cache.num_committed_tokens + + success = kv_cache.resume(stream) + assert success, f"Failed to resume cache for request {req_id}" + + results["num_reused_tokens"] += num_reused + results["num_computed_tokens"] += prompt_len - num_reused + + if not kv_cache.resize(prompt_len + 1): + raise OutOfPagesError(f"Not enough pages for request {req_id}") + + input_tokens = prompt[num_reused:] + + requests.append(Step(kv_cache, input_tokens, prompt[:num_reused])) + + for r in requests: + r.kv_cache.commit(r.input) + r.kv_cache.close() + + s.take_finish_event().synchronize() + + toc_total = time.perf_counter() + results["total_time"] = toc_total - tic_total + + return results + + profiler1 = None + profiler2 = None + if profile: + import cProfile + + profiler1 = cProfile.Profile() + profiler2 = cProfile.Profile() + + # First pass: No cache reuse expected + if profiler1 is not None: + profiler1.enable() + run_requests(prompts) + if profiler1 is not None: + profiler1.disable() + + # Second pass: 100% cache reuse expected + if profiler2 is not None: + profiler2.enable() + results_pass2 = run_requests(prompts) + if profiler2 is not None: + profiler2.disable() + + if PRINT_TIME: + print(f"total_time = {results_pass2['total_time']}") + + # Verify 100% hit rate on second pass + total_tokens_pass2 = ( + results_pass2["num_reused_tokens"] + results_pass2["num_computed_tokens"] + ) + actual_hit_rate = ( + (results_pass2["num_reused_tokens"] / total_tokens_pass2 * 100) + if total_tokens_pass2 > 0 + else 0 + ) + assert abs(actual_hit_rate - 100.0) < 0.01, ( + f"Expected 100% hit rate on second pass, got {actual_hit_rate:.2f}%" + ) + + if profile: + profiler1.print_stats(sort="cumtime") + profiler2.print_stats(sort="cumtime") + profiler1.dump_stats("kv_cache_reuse_pass1.prof") + profiler2.dump_stats("kv_cache_reuse_pass2.prof") + + class TestSSMSupport(unittest.TestCase): """Tests for basic SSM (State Space Model / Mamba) support in KVCacheManager v2.""" @@ -1577,7 +1895,6 @@ def _make_ssm_config( lid += 1 return KVCacheManagerConfig( tokens_per_block=tokens_per_block, - vocab_size=1024, cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)], layers=layers, enable_partial_reuse=enable_partial_reuse, @@ -1592,10 +1909,11 @@ def test_suspend_and_resume_with_ssm(self) -> None: stream_holder = CachedCudaStream() stream = cast(CudaStream, stream_holder.handle) kv_cache.resume(stream) + # Find the SSM layer group ID from the config. ssm_lg = None - for lc_id, lc in self.manager._life_cycles.items(): - if isinstance(lc, SsmLifeCycle): - ssm_lg = LayerGroupId(lc_id) + for layer in cfg.layers: + if isinstance(layer, SsmLayerConfig): + ssm_lg = self.manager.get_layer_group_id(layer.layer_id) break assert ssm_lg is not None # Grow some capacity @@ -1913,6 +2231,88 @@ def test_ssm_reuse_config_validation(self) -> None: self._make_ssm_config(commit_min_snapshot=False) +class TestClampMaxSeqLenForMem(unittest.TestCase): + TOKENS_PER_BLOCK = 32 + SLOT_SIZE = 2 << 20 + + def setUp(self) -> None: + init_cuda_once() + gc.collect() + gc.disable() + self.managers: list[KVCacheManager] = [] + + def tearDown(self) -> None: + for manager in self.managers: + manager.shutdown() + gc.enable() + + def _make_manager(self, sliding_window_sizes: list[int | None]) -> KVCacheManager: + layers = [ + AttentionLayerConfig( + layer_id=LayerId(layer_id), + buffers=[BufferConfig(role=Role.KEY, size=self.SLOT_SIZE)], + sliding_window_size=window_size, + num_sink_tokens=0 if window_size is not None else None, + ) + for layer_id, window_size in enumerate(sliding_window_sizes) + ] + manager = KVCacheManager( + KVCacheManagerConfig( + tokens_per_block=self.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=len(sliding_window_sizes) * self.SLOT_SIZE)], + layers=layers, + ) + ) + self.managers.append(manager) + return manager + + def test_clamp_max_seq_len_for_mem_zero_upper_bound(self): + manager = self._make_manager([None]) + + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=1, token_num_upper_bound=0), 0 + ) + + def test_clamp_max_seq_len_for_mem_single_feasible_block(self): + manager = self._make_manager([None]) + + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=1, token_num_upper_bound=32), 32 + ) + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=1, token_num_upper_bound=64), 32 + ) + + def test_clamp_max_seq_len_for_mem_batch_consumes_remaining_slots(self): + manager = self._make_manager([None]) + + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=2, token_num_upper_bound=64), 0 + ) + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=3, token_num_upper_bound=64), 0 + ) + + def test_clamp_max_seq_len_for_mem_sliding_window_reuses_slot(self): + manager = self._make_manager([self.TOKENS_PER_BLOCK]) + + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=1, token_num_upper_bound=96), 96 + ) + + def test_clamp_max_seq_len_for_mem_multiple_pool_groups(self): + manager = self._make_manager([self.TOKENS_PER_BLOCK, None]) + + # Worst-case SWA slot reservation sizes the pool to 4 slots (SWA floor 2 + + # full-attention floor 2), so a single sequence fits the full 96 tokens. + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=1, token_num_upper_bound=96), 96 + ) + self.assertEqual( + manager.clamp_max_seq_len_for_mem(batch_size=2, token_num_upper_bound=96), 32 + ) + + class TestInitRatioConfig(unittest.TestCase): """Tests for init_ratio computation from typical_step and constraints.""" @@ -1977,7 +2377,6 @@ def _make_config( lid += 1 return KVCacheManagerConfig( tokens_per_block=self.TOKENS_PER_BLOCK, - vocab_size=4096, cache_tiers=cache_tiers, layers=layers, typical_step=typical_step, @@ -1990,7 +2389,7 @@ def test_default_init_ratio(self): """Without typical_step or constraints, uses hardcoded fallback.""" cfg = self._make_config() manager = KVCacheManager(cfg) - ratio = manager._current_gpu_ratio + ratio = _introspection.current_gpu_ratio(manager) self.assertEqual(len(ratio), 2) self.assertAlmostEqual(sum(ratio), 1.0, places=6) # Windowed layers need fewer blocks than non-windowed at history=2048. @@ -2002,7 +2401,7 @@ def test_typical_step_short_sequences(self): step = BatchDesc(kv_caches=[KVCacheDesc(capacity=64, history_length=32)] * 64) cfg = self._make_config(typical_step=step) manager = KVCacheManager(cfg) - ratio = manager._current_gpu_ratio + ratio = _introspection.current_gpu_ratio(manager) self.assertEqual(len(ratio), 2) self.assertAlmostEqual(sum(ratio), 1.0, places=6) # Short sequences (32 tokens < window 128): no stale blocks. @@ -2015,7 +2414,7 @@ def test_typical_step_long_sequences(self): step = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4000)] * 32) cfg = self._make_config(typical_step=step) manager = KVCacheManager(cfg) - ratio = manager._current_gpu_ratio + ratio = _introspection.current_gpu_ratio(manager) self.assertEqual(len(ratio), 2) self.assertAlmostEqual(sum(ratio), 1.0, places=6) # Windowed layers (window=128) have many stale blocks, non-windowed keep all. @@ -2029,11 +2428,11 @@ def test_constraints_floor_typical_step(self): constraint = BatchDesc(kv_caches=[KVCacheDesc(capacity=256, history_length=128)] * 256) cfg_unconstrained = self._make_config(typical_step=typical) mgr_unconstrained = KVCacheManager(cfg_unconstrained) - ratio_unconstrained = mgr_unconstrained._current_gpu_ratio + ratio_unconstrained = _introspection.current_gpu_ratio(mgr_unconstrained) cfg_constrained = self._make_config(typical_step=typical, constraints=[constraint]) mgr_constrained = KVCacheManager(cfg_constrained) - ratio_constrained = mgr_constrained._current_gpu_ratio + ratio_constrained = _introspection.current_gpu_ratio(mgr_constrained) self.assertGreater(ratio_constrained[0], ratio_unconstrained[0]) self.assertAlmostEqual(sum(ratio_constrained), 1.0, places=6) @@ -2050,27 +2449,62 @@ def test_initial_pool_ratio_overrides_typical_step_and_constraints(self): initial_pool_ratio=[0.8, 0.2], ) manager = KVCacheManager(cfg) - ratio = manager._current_gpu_ratio + ratio = _introspection.current_gpu_ratio(manager) self.assertGreater(ratio[0], ratio[1]) self.assertAlmostEqual(sum(ratio), 1.0, places=6) manager.shutdown() - def test_initial_pool_ratio_length_must_match_pool_groups(self): - cfg = self._make_config(initial_pool_ratio=[1.0]) - life_cycles = LifeCycleRegistry(cfg) - storage_config = create_storage_config(cfg) - - with self.assertRaisesRegex(ValueError, "initial_pool_ratio length"): - StorageManager( - life_cycles, - storage_config, - cfg.tokens_per_block, - cfg.swa_scratch_reuse, - typical_batch=cfg.typical_step, - constraints=cfg.constraints, - initial_pool_ratio=cfg.initial_pool_ratio, - ) + @parameterized.expand( + [ + ("empty", [], "initial_pool_ratio length"), + ("wrong_length", [1.0], "initial_pool_ratio length"), + ("zero", [0.0, 1.0], "initial_pool_ratio values must be positive"), + ("negative", [-0.1, 1.1], "initial_pool_ratio values must be positive"), + ("wrong_sum", [0.4, 0.5], "initial_pool_ratio values must sum to 1.0"), + ] + ) + def test_invalid_initial_pool_ratio(self, _name: str, ratio: list[float], error: str): + cfg = self._make_config(initial_pool_ratio=ratio) + + with self.assertRaisesRegex(ValueError, error): + KVCacheManager(cfg) + + def test_ratio_slot_count_rounding_matches_python(self): + grain = 2 << 20 + cfg = KVCacheManagerConfig( + tokens_per_block=self.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=5 * grain)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=grain - 1)], + sliding_window_size=self.TOKENS_PER_BLOCK, + num_sink_tokens=0, + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[BufferConfig(role=Role.KEY, size=grain)], + ), + ], + constraints=[ + BatchDesc(kv_caches=[KVCacheDesc(capacity=self.TOKENS_PER_BLOCK, history_length=0)]) + ], + ) + manager = KVCacheManager(cfg) + + def stat_slot_sizes(stat) -> list[int]: + if hasattr(stat, "slot_sizes"): + return stat.slot_sizes + return stat.slot_size + + slots_by_size = { + tuple(stat_slot_sizes(stat)): stat.total + for stat in _introspection.storage_statistics(manager) + } + self.assertEqual(slots_by_size[(grain - 1,)], 2) + self.assertEqual(slots_by_size[(grain,)], 3) + manager.shutdown() @parameterized.expand([(0,), (64,), (50,), (256,)]) def test_constraint_guarantees_batch_can_run(self, system_prompt_length: int): @@ -2129,7 +2563,7 @@ def test_constraint_guarantees_batch_can_run(self, system_prompt_length: int): manager = KVCacheManager(cfg) # Verify constraint clamping: each pool group has enough slots. - stats = manager._storage.get_statistics() + stats = _introspection.storage_statistics(manager) self.assertGreaterEqual( stats[0].total, slots_pg0, @@ -2256,7 +2690,7 @@ def test_typical_covers_constraint_ratio_unchanged(self): typical_step=typical, ) mgr_no_constraint = KVCacheManager(cfg_no_constraint) - ratio_no_constraint = mgr_no_constraint._current_gpu_ratio + ratio_no_constraint = _introspection.current_gpu_ratio(mgr_no_constraint) # Ratio with constraint that typical already covers. cfg_with_constraint = self._make_config( @@ -2265,7 +2699,7 @@ def test_typical_covers_constraint_ratio_unchanged(self): constraints=[constraint], ) mgr_with_constraint = KVCacheManager(cfg_with_constraint) - ratio_with_constraint = mgr_with_constraint._current_gpu_ratio + ratio_with_constraint = _introspection.current_gpu_ratio(mgr_with_constraint) # Ratios should be identical since typical covers the constraint. for i in range(len(ratio_no_constraint)): @@ -2306,8 +2740,8 @@ def test_typical_step_scratch_reduces_windowed_ratio(self): ) mgr_no = KVCacheManager(cfg_no) mgr_yes = KVCacheManager(cfg_yes) - ratio_no = mgr_no._current_gpu_ratio - ratio_yes = mgr_yes._current_gpu_ratio + ratio_no = _introspection.current_gpu_ratio(mgr_no) + ratio_yes = _introspection.current_gpu_ratio(mgr_yes) # With scratch: PG0 (windowed) needs far fewer slots. self.assertLess(ratio_yes[0], ratio_no[0]) @@ -2374,7 +2808,7 @@ def test_constraint_with_scratch_accounts_for_scratch(self): manager = KVCacheManager(cfg) # Verify constraint clamping: each pool group has enough slots. - stats = manager._storage.get_statistics() + stats = _introspection.storage_statistics(manager) self.assertGreaterEqual( stats[0].total, slots_pg0, @@ -2418,7 +2852,6 @@ def _prepare_scratch( kv_buf_size = 8192 self.cfg = KVCacheManagerConfig( tokens_per_block=tokens_per_block, - vocab_size=4096, cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)], layers=[ AttentionLayerConfig( @@ -2437,6 +2870,70 @@ def _prepare_scratch( self.engine = FakeEngine(self.cfg) self.manager = KVCacheManager(self.cfg) + def test_excess_scratch_slot_waits_for_ready_event_on_new_stream(self): + num_layers = 512 + self._prepare_scratch( + num_layers=num_layers, + window_size=32, + tokens_per_block=32, + gpu_quota=16 << 20, + ) + producer_prompt = [self.next_token() for _ in range(64)] + consumer_prompt = [self.next_token() for _ in range(256)] + producer = self.manager.create_kv_cache(None, producer_prompt) + consumer = self.manager.create_kv_cache(None, consumer_prompt) + producer_stream_holder = CachedCudaStream() + consumer_stream_holder = CachedCudaStream() + producer_stream = cast(CudaStream, producer_stream_holder.handle) + consumer_stream = cast(CudaStream, consumer_stream_holder.handle) + cached_cuda_event = get_cached_cuda_event_type() + producer_marker = None + # Deterministically hold the producer stream open until released from + # the host, so the ordering assertions below cannot pass vacuously + # just because the producer happened to finish early. + gate = HostGate() + + try: + self.assertTrue(producer.resume(producer_stream)) + self.assertTrue(producer.resize(64)) + with enable_kernel_delay(): + for _ in range(8): + self.engine.execute([Step(producer, producer_prompt, [])], producer_stream) + gate.block_stream(producer_stream) + producer_marker = cached_cuda_event(producer_stream) + producer.close() + + self.assertTrue(consumer.resume(producer_stream)) + self.assertTrue(consumer.resize(256)) + self.assertTrue(consumer.has_scratch_slots) + + consumer.cuda_stream = consumer_stream + self.assertTrue(consumer.resize(288, 256)) + self.assertFalse(consumer.has_scratch_slots) + + consumer_marker = cached_cuda_event(consumer_stream) + # While the producer gate is held, the consumer must not be able + # to complete: its scratch->committed migration is ordered after + # the producer's ready event, which is gated. + self.assertFalse(producer_marker.query_complete()) + self.assertFalse(consumer_marker.query_complete()) + gate.release() + consumer_marker.synchronize() + self.assertTrue(producer_marker.query_complete()) + finally: + gate.release() + producer_stream_holder.synchronize() + consumer_stream_holder.synchronize() + if producer_marker is not None and not producer_marker.is_closed(): + producer_marker.synchronize() + if producer.status != _KVCache.Status.CLOSED: + producer.close() + if consumer.status != _KVCache.Status.CLOSED: + consumer.close() + producer_stream_holder.synchronize() + consumer_stream_holder.synchronize() + gate.close() + def test_request_scratch_toggle_for_two_round_inference(self): self._prepare_scratch(num_layers=8, window_size=32, tokens_per_block=32, gpu_quota=16 << 20) prompt = [self.next_token() for _ in range(256)] @@ -2708,7 +3205,6 @@ def test_scratch_chunk_size_variation( self.cfg = KVCacheManagerConfig( tokens_per_block=tokens_per_block, - vocab_size=4096, cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)], layers=[ AttentionLayerConfig( diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py new file mode 100644 index 000000000000..fa43bf2d0992 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import pytest +import torch + +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + GPU_LEVEL, + AttentionLayerConfig, + BufferConfig, + GpuCacheTierConfig, + KVCacheIterationStatsDelta, + KVCacheManager, + KVCacheManagerConfig, + KVCacheStatsDelta, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") + + +def _make_config(*, enable_stats: bool = True) -> KVCacheManagerConfig: + return KVCacheManagerConfig( + tokens_per_block=4, + cache_tiers=[GpuCacheTierConfig(quota=4 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=0, + buffers=[BufferConfig(role="key", size=4096)], + ) + ], + enable_stats=enable_stats, + ) + + +def test_stats_delta_arithmetic() -> None: + stats = KVCacheStatsDelta(4, 3, 2, 1) + delta = KVCacheStatsDelta(1, 2, 3, 4) + stats.add(delta) + assert stats == KVCacheStatsDelta(5, 5, 5, 5) + stats.subtract(delta) + assert stats == KVCacheStatsDelta(4, 3, 2, 1) + copied = stats.copy() + stats.clear() + assert stats.empty + assert copied == KVCacheStatsDelta(4, 3, 2, 1) + + iteration = KVCacheIterationStatsDelta(iter_reused_blocks=3, iter_missed_blocks=1) + assert iteration.iter_cache_hit_rate == 0.75 + iteration.clear() + assert iteration.empty + assert iteration.iter_cache_hit_rate == 0.0 + + +@pytest.mark.parametrize("enable_stats", [False, True]) +def test_manager_stats_config_and_api(enable_stats: bool) -> None: + manager = KVCacheManager(_make_config(enable_stats=enable_stats)) + cache = None + try: + assert manager.init_config.enable_stats is enable_stats + assert manager.get_committed_stats() == KVCacheStatsDelta() + assert manager.get_and_reset_iteration_stats() == {} + peak_stats = manager.get_and_reset_iteration_peak_block_stats(GPU_LEVEL) + assert len(peak_stats) == 1 + assert peak_stats[0].available >= 0 + assert peak_stats[0].unavailable >= 0 + assert peak_stats[0].evictable >= 0 + + manager.mark_stats_dirty(11) + manager.mark_stats_dirty(None) + assert manager.get_dirty_stats_kv_cache_ids() == {11} + manager.mark_stats_excluded(11) + assert manager.is_stats_excluded(11) + assert manager.get_dirty_stats_kv_cache_ids() == set() + manager.clear_stats_excluded(11) + assert not manager.is_stats_excluded(11) + + cache = manager.create_kv_cache(id=17, expected_prompt_length=8) + manager.mark_stats_dirty(17) + assert cache.commit_pending_stats() == KVCacheStatsDelta() + assert manager.get_dirty_stats_kv_cache_ids() == set() + cache.discard_pending_stats() + + # Exercise the collection path: allocate blocks and commit the pending + # stats. With stats enabled the allocation must be visible in both the + # per-request and manager-level committed stats; with stats disabled + # everything must stay empty. + stream = torch.cuda.Stream() + assert cache.resume(stream.cuda_stream) + assert cache.resize(8) + request_stats = cache.commit_pending_stats() + committed = manager.get_committed_stats() + iteration = manager.get_and_reset_iteration_stats() + if enable_stats: + assert request_stats.alloc_total_blocks > 0 + assert committed.alloc_total_blocks > 0 + assert any(not delta.empty for delta in iteration.values()) + else: + assert request_stats == KVCacheStatsDelta() + assert committed == KVCacheStatsDelta() + assert iteration == {} + stream.synchronize() + finally: + if cache is not None: + cache.close() + manager.shutdown() diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 7b6dafb97bce..6038fe3d4906 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -501,6 +501,10 @@ def test_KvCacheConfig_declaration(): kv_cache_event_hash_algo="auto").kv_cache_event_hash_algo == "auto" assert KvCacheConfig(kv_cache_event_hash_algo="v1_block_key" ).kv_cache_event_hash_algo == "v1_block_key" + assert KvCacheConfig(kv_cache_event_hash_algo="v2_sha256" + ).kv_cache_event_hash_algo == "v2_sha256" + assert KvCacheConfig(kv_cache_event_hash_algo="v2_sha256_64" + ).kv_cache_event_hash_algo == "v2_sha256_64" assert pybind_config.enable_partial_reuse == True assert pybind_config.copy_on_partial_reuse == True assert pybind_config.attention_dp_events_gather_period_ms == 10