diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index adb134953938..19600e0fd52a 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -144,6 +144,10 @@ struct BufferConfig std::optional tokensPerBlockOverride; bool isSparse = false; //!< Whether history uses the sparse attention lifecycle. + + //! Buffers may coalesce only when their group and expanded size match. + //! This does not change their lifecycle or shared slot allocation. + int coalescingGroup = 0; }; // --------------------------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp index c9e3718b14b9..a5d17a7817e1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp @@ -154,13 +154,13 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) LifeCycleRegistry registry{config}; int tokensPerBlock = config.tokensPerBlock; - // Map: lifeCycleId → (bufferSize → list of BufferId). - // Outer map key is LifeCycleId; inner map key is expanded buffer size. + // Map: lifeCycleId → ((coalescingGroup, expanded buffer size) → list of BufferId). // NOTE: Python uses insertion-ordered dict/defaultdict here. The std::map // grouping below intentionally remains sorted, which can make observable // pool-group indices/layout order differ from Python. Later lookups use // StorageConfig mappings, so this is not expected to affect correctness. - std::map>> bufferGroups; + using CoalescingKey = std::pair; + std::map>> bufferGroups; std::unordered_map expansionMap; for (auto const& layer : config.layers) @@ -179,7 +179,7 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) BufferId bid{cfg.layerId, buf.role}; expansionMap[bid] = exp; - bufferGroups[lcId][expandedSize].push_back(bid); + bufferGroups[lcId][{buf.coalescingGroup, expandedSize}].push_back(bid); } }, layer); @@ -189,14 +189,14 @@ StorageConfig createStorageConfig(KVCacheManagerConfig const& config) // Each variant: coalesced buffers sorted by size descending. std::vector slotGroups; slotGroups.reserve(bufferGroups.size()); - for (auto const& [lcId, sizeMap] : bufferGroups) + for (auto const& [lcId, coalescingMap] : bufferGroups) { SlotDescVariant var; var.lifeCycleId = lcId; - for (auto const& [sz, bufIds] : sizeMap) + for (auto const& [key, bufIds] : coalescingMap) { CoalescedBuffer cb; - cb.singleBufferSize = sz; + cb.singleBufferSize = key.second; cb.bufferIds = bufIds; var.coalescedBuffers.push_back(std::move(cb)); } diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index d56f2bb5fc83..c503e0250648 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1616,12 +1616,14 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) .def("validate", &kv::DiskCacheTierConfig::validate) DEF_COPY(kv::DiskCacheTierConfig); nb::class_(m, "BufferConfig") - .def(nb::init, bool>(), nb::arg("role"), nb::arg("size"), - nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false) + .def(nb::init, bool, int>(), nb::arg("role"), nb::arg("size"), + nb::arg("tokens_per_block_override") = std::nullopt, nb::arg("is_sparse") = false, + nb::arg("coalescing_group") = 0) .def_rw("role", &kv::BufferConfig::role) .def_rw("size", &kv::BufferConfig::size) .def_rw("tokens_per_block_override", &kv::BufferConfig::tokensPerBlockOverride) - .def_rw("is_sparse", &kv::BufferConfig::isSparse) DEF_COPY(kv::BufferConfig); + .def_rw("is_sparse", &kv::BufferConfig::isSparse) + .def_rw("coalescing_group", &kv::BufferConfig::coalescingGroup) DEF_COPY(kv::BufferConfig); nb::class_(m, "AttentionLayerConfig") .def(nb::init, std::optional, std::optional, diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index c2c50aff9d1b..824c1f559edb 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1717,14 +1717,17 @@ KvCachePoolPointers buildKvCachePoolPointers(at::Tensor const& hostKvCachePoolPo } else { - TORCH_CHECK(hostKvCachePoolPointers.dim() == 2); + TORCH_CHECK(hostKvCachePoolPointers.dim() == 2 || hostKvCachePoolPointers.dim() == 3); + // Non-FP4 layers use only the data pointers in a mixed-format pool table. + auto const dataPointers + = hostKvCachePoolPointers.dim() == 3 ? hostKvCachePoolPointers.select(2, 0) : hostKvCachePoolPointers; pointers.primaryPoolPtr = reinterpret_cast( reinterpret_cast( - readHostTensor2D(hostKvCachePoolPointers, poolIndex, 0, "host_kv_cache_pool_pointers")) + readHostTensor2D(dataPointers, poolIndex, 0, "host_kv_cache_pool_pointers")) + intraPoolOffset); pointers.secondaryPoolPtr = reinterpret_cast( reinterpret_cast( - readHostTensor2D(hostKvCachePoolPointers, poolIndex, 1, "host_kv_cache_pool_pointers")) + readHostTensor2D(dataPointers, poolIndex, 1, "host_kv_cache_pool_pointers")) + intraPoolOffset); } return pointers; diff --git a/tensorrt_llm/_torch/attention/backends/flashinfer.py b/tensorrt_llm/_torch/attention/backends/flashinfer.py index 1325c50aeb85..c7df494c58d4 100644 --- a/tensorrt_llm/_torch/attention/backends/flashinfer.py +++ b/tensorrt_llm/_torch/attention/backends/flashinfer.py @@ -1211,20 +1211,23 @@ def _post_init_with_buffers(self, buffers) -> None: self._max_num_blocks_per_seq = self.kv_cache_manager.max_blocks_per_seq # Layers may share one page-index list only when they are in the - # same pool AND have the same page-index scale: VSWA splits pools, + # same pool, page-index scale, and KV factor: VSWA splits pools, # and per-layer geometry (e.g. Gemma4 sliding/global head_dim) # splits scales even when all windows collapse to max_seq_len and - # is_vswa is False. Guarded on V2-specific attributes (V1 managers - # lack the per-pool infrastructure). + # is_vswa is False. Mixed K-only/KV layers can use different view + # indices even with the same physical scale. Guarded on V2-specific + # attributes (V1 managers lack the per-pool infrastructure). mgr = self.kv_cache_manager get_scale = getattr(mgr, 'get_layer_page_index_scale', None) + get_kv_factor = getattr(mgr, 'get_layer_kv_factor', None) layer_space: Dict[int, int] = {} if hasattr(mgr, 'layer_to_pool_mapping_dict') and get_scale: space_ids = {} for layer_idx in attention_layer_indices: layer_offset = mgr.layer_offsets[layer_idx] key = (mgr.layer_to_pool_mapping_dict[layer_offset], - get_scale(layer_idx)) + get_scale(layer_idx), + get_kv_factor(layer_idx) if get_kv_factor else None) if (self._has_mixed_precision_kv() and layer_offset in mgr.high_layers): key = ("mixed_precision_staging", ) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py index c52c6d665126..5b450c3742c9 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py @@ -987,6 +987,7 @@ def _get_batch_cache_indices_by_pool_id( is_kv_aggregate: bool = True, num_blocks_per_seq: Optional[Sequence[int]] = None, index_scale: Optional[int] = None, + kv_factor: Optional[int] = None, ): """Return page indices; padded entries remain ``BAD_PAGE_INDEX`` (-1). @@ -1015,6 +1016,8 @@ def _get_batch_cache_indices_by_pool_id( index_scale: Kept for compatibility with the base virtual method; M3 bypasses the V1 block-id conversion entirely, so any caller-supplied scale is ignored alongside ``index_scales``. + kv_factor: Kept for compatibility with the base virtual method; + M3 returns raw slot IDs without KV aggregation. """ res = [] for req_idx, req_id in enumerate(request_ids): diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 806e104d148a..54b9deae7fc3 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -283,6 +283,101 @@ class BlockReusePolicy(StrEnum): PER_CONVERSATION = "per_conversation" +class KVCacheLayerOwner(StrEnum): + """Model component that reads and writes a cache layer.""" + + TARGET = "target" + DRAFT = "draft" + + +@dataclass(frozen=True) +class KVCacheLayerSpec: + """Rank-local attention-layer storage consumed by costing and construction. + + Tuple position is the C++ local layer id. ``global_layer_id`` is the model + layer id used by attention metadata. + """ + + global_layer_id: int + owner: KVCacheLayerOwner + dtype: DataType + num_kv_heads: int + head_dim: int + buffer_roles: tuple[DataRole, ...] + sliding_window_size: Optional[int] = None + num_sink_tokens: Optional[int] = None + reuse_policy: AttentionReusePolicy = AttentionReusePolicy.REQUIRED + + def __post_init__(self) -> None: + if self.global_layer_id < 0: + raise ValueError(f"global_layer_id must be non-negative, got {self.global_layer_id}") + if self.num_kv_heads <= 0: + raise ValueError(f"num_kv_heads must be positive, got {self.num_kv_heads}") + if self.head_dim <= 0: + raise ValueError(f"head_dim must be positive, got {self.head_dim}") + if not self.buffer_roles: + raise ValueError("buffer_roles must not be empty") + if len(set(self.buffer_roles)) != len(self.buffer_roles): + raise ValueError( + f"layer {self.global_layer_id} contains duplicate buffer roles: {self.buffer_roles}" + ) + supported_roles = { + Role.KEY, + Role.VALUE, + Role.KEY_BLOCK_SCALE, + Role.VALUE_BLOCK_SCALE, + } + unsupported_roles = set(self.buffer_roles) - supported_roles + if unsupported_roles: + raise ValueError( + f"layer {self.global_layer_id} has unsupported standard buffer roles: " + f"{sorted(map(str, unsupported_roles))}" + ) + if Role.KEY not in self.buffer_roles: + raise ValueError(f"layer {self.global_layer_id} must contain a key buffer") + scale_roles = {Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE} + if any(role in scale_roles for role in self.buffer_roles) and self.dtype != DataType.NVFP4: + raise ValueError("block-scale roles require NVFP4 KV-cache dtype") + if self.dtype == DataType.NVFP4: + if self.head_dim % 2 != 0: + raise ValueError(f"NVFP4 head_dim must be even, got {self.head_dim}") + expected_scale_roles = {Role.KEY_BLOCK_SCALE} + if Role.VALUE in self.buffer_roles: + expected_scale_roles.add(Role.VALUE_BLOCK_SCALE) + if scale_roles.intersection(self.buffer_roles) != expected_scale_roles: + raise ValueError( + f"NVFP4 layer {self.global_layer_id} must contain matching block-scale roles" + ) + if self.sliding_window_size is not None and self.sliding_window_size <= 0: + raise ValueError( + f"sliding_window_size must be positive when set, got {self.sliding_window_size}" + ) + if self.num_sink_tokens is not None and self.num_sink_tokens < 0: + raise ValueError( + f"num_sink_tokens must be non-negative when set, got {self.num_sink_tokens}" + ) + + def bytes_per_token(self, role: Optional[DataRole] = None) -> int: + """Return storage bytes per token for one role or the whole layer.""" + if role is None: + return sum(self.bytes_per_token(buffer_role) for buffer_role in self.buffer_roles) + if role not in self.buffer_roles: + raise ValueError(f"role {role} is not present in layer {self.global_layer_id}") + elements = self.num_kv_heads * self.head_dim + if role in (Role.KEY, Role.VALUE): + return get_size_in_bytes(elements, self.dtype) + if role in (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE): + if elements % 16 != 0: + raise ValueError(f"NVFP4 layer {self.global_layer_id} size must be divisible by 16") + return get_size_in_bytes(elements // 16, DataType.FP8) + raise ValueError(f"unsupported buffer role: {role}") + + @property + def kv_factor(self) -> int: + """Number of data planes exposed by the layer's paged tensor view.""" + return 2 if Role.VALUE in self.buffer_roles else 1 + + def _swa_endpoint_priority( prompt_length: int, tokens_per_block: int, rewind_tokens: int ) -> Callable[[int, object], int]: @@ -517,26 +612,34 @@ def _estimate_swa_cache_size( context: bool, scratch: bool, generation_capacity_headroom: int = BASE_GENERATION_TOKEN_COUNT, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, ) -> tuple[int, int]: tokens_per_block = int(tokens_per_block) size_per_token = 0 size_per_request = 0 scratch_keys = set() - for layer_size, window_size in zip(layer_sizes, attention_windows): + for layer_idx, (layer_size, window_size) in enumerate(zip(layer_sizes, attention_windows)): if window_size is not None and window_size > 0: + layer = layer_specs[layer_idx] if layer_specs is not None else None + sink_blocks = math.ceil((layer.num_sink_tokens or 0) / tokens_per_block) if layer else 0 # Match AttnLifeCycle.get_stale_range(): the live interval contains # window_size + generation_capacity_headroom - 1 tokens. Across all # page offsets, that interval touches at most the count below. window_blocks = ( math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1 ) - window_tokens = window_blocks * tokens_per_block + window_tokens = (sink_blocks + window_blocks) * tokens_per_block if not context: size_per_request += window_tokens * layer_size elif not scratch: size_per_token += layer_size else: - scratch_key = (int(window_size), layer_size) + scratch_key = ( + int(window_size), + layer_size, + sink_blocks, + layer.reuse_policy if layer else None, + ) if scratch_key in scratch_keys: size_per_request += window_tokens * layer_size else: @@ -573,6 +676,7 @@ def _estimate_cache_size_components( *, scratch: bool, generation_capacity_headroom: int, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, ) -> tuple[int, int, int]: """Return context/generation bytes per token and generation bytes per request. @@ -582,7 +686,12 @@ def _estimate_cache_size_components( """ full_attn_size = _estimate_full_attn_size_per_token(layer_sizes, attention_windows) context_swa_size, _ = _estimate_swa_cache_size( - layer_sizes, attention_windows, tokens_per_block, context=True, scratch=scratch + layer_sizes, + attention_windows, + tokens_per_block, + context=True, + scratch=scratch, + layer_specs=layer_specs, ) generation_swa_size, generation_swa_per_request = _estimate_swa_cache_size( layer_sizes, @@ -591,6 +700,7 @@ def _estimate_cache_size_components( context=False, scratch=False, generation_capacity_headroom=generation_capacity_headroom, + layer_specs=layer_specs, ) return ( full_attn_size + context_swa_size, @@ -797,6 +907,30 @@ def normalize_window_size(window_size: Optional[int]) -> Optional[int]: return layer_sizes, attention_windows +def _get_layer_spec_cache_size_components( + layer_specs: Sequence[KVCacheLayerSpec], + *, + max_seq_len: Optional[int] = None, +) -> tuple[List[int], List[Optional[int]]]: + """Return the shared costing inputs for an explicit layer-spec sequence.""" + if not layer_specs: + raise ValueError("layer_specs must contain at least one rank-local layer") + global_layer_ids = [layer.global_layer_id for layer in layer_specs] + if len(set(global_layer_ids)) != len(global_layer_ids): + raise ValueError(f"layer_specs contains duplicate global layer IDs: {global_layer_ids}") + return ( + [layer.bytes_per_token() for layer in layer_specs], + [ + None + if max_seq_len is not None + and layer.sliding_window_size is not None + and layer.sliding_window_size >= max_seq_len + else layer.sliding_window_size + for layer in layer_specs + ], + ) + + class _MmRunMetadata(NamedTuple): """CPU tensors for exact multimodal run lookup. @@ -1126,6 +1260,31 @@ def _update_kv_cache_draft_token_location( "update_kv_cache_draft_token_location requires all local layers " f"in one KV pool, got pools {sorted(local_pool_ids)}" ) + layer_specs = getattr(cache_manager, "layer_specs", None) + if layer_specs is not None: + if len({(layer.head_dim, layer.dtype) for layer in layer_specs}) != 1: + raise ValueError( + "update_kv_cache_draft_token_location requires uniform head_dim and dtype " + "for explicit layer specs" + ) + if any(set(layer.buffer_roles) != {Role.KEY, Role.VALUE} for layer in layer_specs): + raise ValueError( + "update_kv_cache_draft_token_location requires explicit layer specs " + "with only KEY and VALUE buffers" + ) + # The kernel advances each layer base by one contiguous K/V pair. + if any( + int(cache_manager.kv_cache_pool_mapping[layer_idx][1]) != layer_idx + for layer_idx in range(cache_manager.num_local_layers) + ): + raise ValueError( + "update_kv_cache_draft_token_location requires explicit layers " + "in contiguous local-layer order within the KV pool" + ) + layer = layer_specs[0] + head_size_in_bytes = get_size_in_bytes(layer.head_dim, layer.dtype) + else: + head_size_in_bytes = int(cache_manager.head_dim * kv_cache_dtype_byte_size) pool_idx = local_pool_ids.pop() torch.ops.tensorrt_llm.update_kv_cache_draft_token_location( accepted_draft_token_offsets, @@ -1136,7 +1295,7 @@ def _update_kv_cache_draft_token_location( # Use TP-sharded num_kv_heads (per-rank) instead of the unsharded # total so the C++ kernel computes correct strides and grid dims. cache_manager.num_kv_heads_per_layer[0], - int(cache_manager.head_dim * kv_cache_dtype_byte_size), + head_size_in_bytes, cache_manager.max_total_draft_tokens, max_attention_window, rewind_draft_token_separate_adjustments, @@ -1360,10 +1519,26 @@ def __init__( cold_page_codec_provider: Optional[object] = None, joint_kv_cache_reuse: bool = False, max_cuda_graph_batch_size: Optional[int] = None, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> None: self.mapping = mapping self.dtype = dtype + self._declared_target_layer_count = num_layers + self.layer_specs = tuple(layer_specs) if layer_specs is not None else None + if ( + self.layer_specs is not None + and type(self)._extra_buffers_per_layer is not KVCacheManagerV2._extra_buffers_per_layer + ): + raise ValueError( + f"{type(self).__name__} does not support layer_specs with " + "an overridden _extra_buffers_per_layer." + ) + if ( + self.layer_specs is not None + and getattr(kv_cache_config, "mixed_precision_kv", None) is not None + ): + raise ValueError("layer_specs does not support mixed-precision KV cache.") self._init_mixed_precision( kv_cache_config, kwargs.get("pretrained_config"), max_seq_len, dtype ) @@ -1378,13 +1553,34 @@ def __init__( self._layer_group_life_cycles: Optional[List[AttnLifeCycle]] = None self.kv_cache_type = kv_cache_type - self.pp_layers, self.num_layers = get_pp_layers( - num_layers, - mapping, - spec_config=spec_config, - layer_mask=layer_mask, - ) + if self.layer_specs is None: + self.pp_layers, self.num_layers = get_pp_layers( + num_layers, + mapping, + spec_config=spec_config, + layer_mask=layer_mask, + ) + else: + self.pp_layers = [layer.global_layer_id for layer in self.layer_specs] + if len(set(self.pp_layers)) != len(self.pp_layers): + raise ValueError(f"layer-spec global layer ids must be unique: {self.pp_layers}") + self.num_layers = max(num_layers, max(self.pp_layers, default=-1) + 1) self.is_draft = is_draft + self._legacy_draft_layer_ids = { + global_layer_id + for global_layer_id in self.pp_layers + if is_draft + or ( + spec_config is not None + and layer_mask is None + and global_layer_id >= self._declared_target_layer_count + ) + } + if self.layer_specs is not None: + if is_draft and any( + layer.owner is not KVCacheLayerOwner.DRAFT for layer in self.layer_specs + ): + raise ValueError("a draft mirror manager cannot own target cache layers") # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies) # can distinguish the throwaway estimation-phase managers from the @@ -1395,7 +1591,12 @@ def __init__( # Set True by a compression manager; generation-step resize then leaves history untouched. self.kv_compression_manages_history: bool = False self.enable_swa_scratch_reuse = ( - kv_cache_config.enable_swa_scratch_reuse and not self.is_draft + kv_cache_config.enable_swa_scratch_reuse + and not self.is_draft + and not ( + self.layer_specs is not None + and any(layer.owner is KVCacheLayerOwner.DRAFT for layer in self.layer_specs) + ) ) block_reuse_config = kv_cache_config.block_reuse_config self.block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) @@ -1460,10 +1661,12 @@ def __init__( "all CP ranks, so publishing it requires a group-wide " "protocol. Set kv_cache_config.enable_block_reuse=False." ) - if is_draft: + if is_draft or any( + layer.owner is KVCacheLayerOwner.DRAFT for layer in self.layer_specs or () + ): raise ValueError( "KVCacheManagerV2 does not support a draft cache " - "manager with helix context parallelism." + "manager or draft-owned layers with helix context parallelism." ) if mapping.enable_attention_dp: raise ValueError( @@ -1530,6 +1733,13 @@ def __init__( self.pp_layers, ) assert len(self.max_attention_window_vec) == self.num_local_layers + if self.layer_specs is not None: + self.max_attention_window_vec = [ + None + if layer.sliding_window_size is None or layer.sliding_window_size >= max_seq_len + else layer.sliding_window_size + for layer in self.layer_specs + ] # A D-token match backoff needs D more physical SWA history at the # committed endpoint so commit_min_snapshot still covers the rewind. @@ -1596,7 +1806,12 @@ def __init__( mm_token_id_offset=vocab_size, ) - if isinstance(num_kv_heads, int): + if self.layer_specs is not None: + self.num_kv_heads_per_layer = [layer.num_kv_heads for layer in self.layer_specs] + self.total_num_kv_heads_per_layer = [0] * self.num_layers + for layer in self.layer_specs: + self.total_num_kv_heads_per_layer[layer.global_layer_id] = layer.num_kv_heads + elif isinstance(num_kv_heads, int): self.num_kv_heads_per_layer = [ (num_kv_heads + tp_size - 1) // tp_size for _ in range(self.num_local_layers) ] @@ -1629,7 +1844,9 @@ def append_to_kv_heads_per_layer( append_to_kv_heads_per_layer(self.total_num_kv_heads_per_layer, kv_head) # Build per-layer head_dim (similar to num_kv_heads_per_layer) - if isinstance(head_dim, int): + if self.layer_specs is not None: + self.head_dim_per_layer = [layer.head_dim for layer in self.layer_specs] + elif isinstance(head_dim, int): self.head_dim_per_layer = [head_dim for _ in range(self.num_local_layers)] else: assert len(head_dim) == self.num_layers, ( @@ -1793,9 +2010,16 @@ def append_to_kv_heads_per_layer( def create_cold_page_codec(cache_config: object) -> Optional[object]: if cold_page_codec_provider is None: return None + if ( + self.layer_specs is not None + and len({layer.dtype for layer in self.layer_specs}) > 1 + ): + raise NotImplementedError( + "cold-page compression does not yet support mixed-dtype KV-cache layer specs" + ) return cold_page_codec_provider.create_cold_page_codec( cache_config, - runtime_dtype=self.dtype, + runtime_dtype=self.layer_specs[0].dtype if self.layer_specs else self.dtype, pp_layers=self.pp_layers, num_kv_heads_per_layer=self.num_kv_heads_per_layer, head_dim_per_layer=self.head_dim_per_layer, @@ -2291,7 +2515,10 @@ def _fill_fresh_kv_pages(self, request_id: int) -> None: # Same conversion the attention backends index the buffer # with: layers in one pool can carry different scales. scale = self.get_layer_page_index_scale(layer_idx) - pages = [page * scale // self.kv_factor for page in fresh] + layer_offset = self.layer_offsets[layer_idx] + kv_factor = self._get_layer_kv_factor_by_local_id(layer_offset) + view_stride = self._get_cache_view_page_stride(scale, kv_factor) + pages = [page * scale // view_stride for page in fresh] pages = [page for page in pages if 0 <= page < buffer.shape[0]] if pages: if _fill_kv_pages(buffer, pages, self._fresh_page_fill): @@ -2327,12 +2554,38 @@ def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: When present, role B must be addressable from role A using a constant page-index offset. """ + if self.layer_specs is not None: + layer_id = self._pool_layer_ids_by_role[(pool_id, Role.KEY)] + roles = self._get_layer_roles_by_local_id(int(layer_id)) + return Role.KEY, Role.VALUE if Role.VALUE in roles else None role_b = None if self.kv_cache_type == CacheTypeCpp.SELFKONLY else Role.VALUE return Role.KEY, role_b + def _get_layer_kv_factor_by_local_id(self, local_layer_id: int) -> int: + if self.layer_specs is not None: + return self.layer_specs[local_layer_id].kv_factor + return self.kv_factor + + def _get_layer_page_table_roles( + self, local_layer_id: int, pool_id: int + ) -> Tuple[DataRole, Optional[DataRole]]: + if self.layer_specs is None: + return self._get_pool_roles(pool_id) + return Role.KEY, (Role.VALUE if self.layer_specs[local_layer_id].kv_factor == 2 else None) + + def _get_layer_dtype_by_local_id(self, local_layer_id: int) -> DataType: + if self.layer_specs is not None: + return self.layer_specs[local_layer_id].dtype + return self.dtype + + def _get_layer_roles_by_local_id(self, local_layer_id: int) -> tuple[DataRole, ...]: + if self.layer_specs is not None: + return self.layer_specs[local_layer_id].buffer_roles + return tuple(self._get_buffer_roles_for_layer(local_layer_id)) + def _get_block_scale_role(self, role_a: DataRole, layer_id: int) -> Optional[DataRole]: """Select block scales for a role and local layer in the pool mapping.""" - if self.dtype != DataType.NVFP4 or role_a != Role.KEY: + if self._get_layer_dtype_by_local_id(layer_id) != DataType.NVFP4 or role_a != Role.KEY: return None return Role.KEY_BLOCK_SCALE @@ -2345,10 +2598,14 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] + uses_block_scale_pointers = any( + self._get_layer_dtype_by_local_id(layer_id) == DataType.NVFP4 + for layer_id in range(self.num_local_layers) + ) if self._use_per_layer_page_tables: for layer_id in typed_range(LayerId(self.num_local_layers)): pool_id = self.impl.get_layer_group_id(layer_id) - role_a, _ = self._get_pool_roles(pool_id) + role_a, _ = self._get_layer_page_table_roles(layer_id, pool_id) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( @@ -2357,7 +2614,7 @@ def _build_pool_mapping_tensors(self): 0, ] ) - if self.dtype == DataType.NVFP4: + if uses_block_scale_pointers: block_scale_role = self._get_block_scale_role(role_a, layer_id) block_scale_pool_pointers_list.append( [ @@ -2378,7 +2635,7 @@ def _build_pool_mapping_tensors(self): layer_id, role_a, PageIndexMode.SHARED ) kv_cache_pool_pointers_list.append([key_base_addr, 0]) - if self.dtype == DataType.NVFP4: + if uses_block_scale_pointers: # The KEY/scale pointers are a (rep-layer, 0-offset) origin # against which each layer's kv_cache_pool_mapping offset is # resolved. The block-scale origin must reproduce the SAME @@ -2395,7 +2652,7 @@ def _build_pool_mapping_tensors(self): rep_offset = self._kv_pool_mapping_offset(layer_id, pool_id, key_base_addr) scale_stride = ( self.get_layer_bytes_per_token(layer_id, block_scale_role) - * self.kv_factor + * self._get_layer_kv_factor_by_local_id(layer_id) * self.tokens_per_block ) scale_base_addr = ( @@ -2421,13 +2678,16 @@ def _build_pool_mapping_tensors(self): ) offset_divisor = self.impl.get_page_stride(layer_id, role_a) if role_b is not None: - offset_divisor *= self.kv_factor + offset_divisor *= self._get_layer_kv_factor_by_local_id(layer_id) offset = exact_div( addr_offset, offset_divisor, ) - if self.dtype != DataType.NVFP4 or role_a != Role.KEY: + if ( + self._get_layer_dtype_by_local_id(layer_id) != DataType.NVFP4 + or role_a != Role.KEY + ): block_scale_offset = None else: block_scale_role = self._get_block_scale_role(role_a, layer_id) @@ -2444,7 +2704,7 @@ def _build_pool_mapping_tensors(self): block_scale_offset = exact_div( block_scale_addr_offset, self.get_layer_bytes_per_token(layer_id, block_scale_role) - * self.kv_factor + * self._get_layer_kv_factor_by_local_id(layer_id) * self.tokens_per_block, ) @@ -2455,7 +2715,7 @@ def _build_pool_mapping_tensors(self): kv_cache_pool_mapping_list.append([layer_group_id, offset]) - if self.dtype == DataType.NVFP4: + if uses_block_scale_pointers: for pool_id, block_scale_pool_pointers in enumerate(block_scale_pool_pointers_list): pool_pointers = kv_cache_pool_pointers_list[pool_id] kv_cache_pool_pointers_list[pool_id] = [ @@ -2495,6 +2755,9 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: > 1 for pool_id, layer_ids in enumerate(self.impl.layer_grouping) ) + if self.layer_specs is not None: + storage_layouts = {(layer.dtype, layer.kv_factor) for layer in self.layer_specs} + self._use_per_layer_page_tables |= len(storage_layouts) > 1 if self._use_per_layer_page_tables: self.num_attention_op_pools = self.num_local_layers else: @@ -2555,11 +2818,16 @@ def _kv_pool_mapping_offset( return exact_div( addr_offset, self.get_layer_bytes_per_token(layer_id, Role.KEY) - * self.kv_factor + * self._get_layer_kv_factor_by_local_id(layer_id) * self.tokens_per_block, ) def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]: + if self.layer_specs is not None: + return ( + [layer.bytes_per_token() for layer in self.layer_specs], + list(self.max_attention_window_vec), + ) layer_sizes = [] attention_windows = [] for local_layer_idx in range(self.num_local_layers): @@ -2596,6 +2864,7 @@ def _get_max_tokens_from_quota_impl(self, quota: int) -> float: self.tokens_per_block, scratch=self.enable_swa_scratch_reuse, generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, ) size_per_batch = self._get_generation_request_capacity() * generation_swa_size_per_request if quota < size_per_batch: @@ -2632,6 +2901,7 @@ def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: self.tokens_per_block, scratch=self.enable_swa_scratch_reuse, generation_capacity_headroom=self._generation_kv_capacity_headroom, + layer_specs=self.layer_specs, ) context_tokens = min(max_tokens, self.max_num_tokens) generation_tokens = max_tokens - context_tokens @@ -2791,7 +3061,7 @@ def _prepare_swa_scratch_copy_tensors(self, index_mapper_capacity: int) -> None: for local_layer_idx in range(self.num_local_layers): layer_id = LayerId(local_layer_idx) pool_id = self.layer_to_pool_mapping_dict[layer_id] - role_a, role_b = self._get_pool_roles(pool_id) + role_a, role_b = self._get_layer_page_table_roles(layer_id, pool_id) roles = [ role_a, role_a if role_b is None else role_b, @@ -2980,6 +3250,8 @@ def _copy_batch_block_offsets_per_layer( def _get_buffer_roles_for_layer(self, local_layer_idx: int) -> List[DataRole]: """Return the primary cache roles allocated for one local layer.""" + if self.layer_specs is not None: + return list(self.layer_specs[local_layer_idx].buffer_roles) roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: roles.append(Role.VALUE) @@ -3073,43 +3345,70 @@ def _build_base_config( ) ) - # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional - # per-layer BufferConfig entries — for example a sparse index-K - # buffer — without overriding the K/V/NVFP4 scale wiring above. - # The dict maps local layer id -> list of extra BufferConfig. Each - # extra buffer's role must be unique within the layer (asserted by - # AttentionLayerConfig.__post_init__) and its size must be in bytes - # per block (= bytes_per_token * tokens_per_block). - extra_buffers_per_layer = ( - self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} - ) - - layer_configs: List[AttentionLayerConfig] = [] - for layer_id in typed_range(LayerId(self.num_local_layers)): - buffer_type = self._get_buffer_roles_for_layer(layer_id) - buffers = [ - BufferConfig( - role=role, - size=self.get_layer_bytes_per_token(local_layer_idx=layer_id, data_role=role) - * tokens_per_block, - ) - for role in buffer_type - ] - for extra in extra_buffers_per_layer.get(int(layer_id), ()): - assert extra.role not in buffer_type, ( - f"extra buffer role {extra.role!r} for layer " - f"{int(layer_id)} duplicates a standard K/V/scale role" - ) - buffers.append(extra) - layer_configs.append( + if self.layer_specs is not None: + scale_roles = (Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE) + layer_configs = [ AttentionLayerConfig( - layer_id=layer_id, - buffers=buffers, - sliding_window_size=self.max_attention_window_vec[layer_id], - num_sink_tokens=None, + layer_id=LayerId(local_layer_id), + buffers=[ + BufferConfig( + role=role, + size=layer.bytes_per_token(role) * tokens_per_block, + # Keep NVFP4 data (1), its scales (2), and other data (0) + # separate while preserving matching K/V page order. + coalescing_group=( + 1 + int(role in scale_roles) if layer.dtype == DataType.NVFP4 else 0 + ), + ) + for role in (Role.KEY, Role.VALUE, *scale_roles) + if role in layer.buffer_roles + ], + sliding_window_size=self.max_attention_window_vec[local_layer_id], + num_sink_tokens=layer.num_sink_tokens, + reuse_policy=layer.reuse_policy, ) + for local_layer_id, layer in enumerate(self.layer_specs) + ] + else: + # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional + # per-layer BufferConfig entries — for example a sparse index-K + # buffer — without overriding the K/V/NVFP4 scale wiring above. + # The dict maps local layer id -> list of extra BufferConfig. Each + # extra buffer's role must be unique within the layer (asserted by + # AttentionLayerConfig.__post_init__) and its size must be in bytes + # per block (= bytes_per_token * tokens_per_block). + extra_buffers_per_layer = ( + self._extra_buffers_per_layer(tokens_per_block=tokens_per_block) or {} ) + layer_configs: List[AttentionLayerConfig] = [] + for layer_id in typed_range(LayerId(self.num_local_layers)): + buffer_type = self._get_buffer_roles_for_layer(layer_id) + buffers = [ + BufferConfig( + role=role, + size=self.get_layer_bytes_per_token( + local_layer_idx=layer_id, data_role=role + ) + * tokens_per_block, + ) + for role in buffer_type + ] + for extra in extra_buffers_per_layer.get(int(layer_id), ()): + assert extra.role not in buffer_type, ( + f"extra buffer role {extra.role!r} for layer " + f"{int(layer_id)} duplicates a standard K/V/scale role" + ) + buffers.append(extra) + layer_configs.append( + AttentionLayerConfig( + layer_id=layer_id, + buffers=buffers, + sliding_window_size=self.max_attention_window_vec[layer_id], + num_sink_tokens=None, + ) + ) + rebalance_cfg = kv_cache_config.kv_pool_rebalance_config return KVCacheManagerConfigPy( # Used by the backend only for token<->block arithmetic and @@ -3187,6 +3486,7 @@ def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheMa buffers=buffers, sliding_window_size=layer.sliding_window_size, num_sink_tokens=layer.num_sink_tokens, + reuse_policy=layer.reuse_policy, ) ) else: @@ -3483,6 +3783,64 @@ def get_disagg_ignored_roles(self) -> frozenset[DataRole]: """ return frozenset() + def get_layer_cache_spec(self, global_layer_id: int) -> KVCacheLayerSpec: + """Return an explicitly supplied rank-local layer specification.""" + if self.layer_specs is None: + raise RuntimeError("this manager was not constructed with explicit layer specs") + try: + local_layer_id = self.layer_offsets[global_layer_id] + except KeyError as error: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") from error + return self.layer_specs[local_layer_id] + + def get_layer_owner(self, global_layer_id: int) -> KVCacheLayerOwner: + if self.layer_specs is not None: + return self.get_layer_cache_spec(global_layer_id).owner + if global_layer_id in self._legacy_draft_layer_ids: + return KVCacheLayerOwner.DRAFT + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return KVCacheLayerOwner.TARGET + + def get_layer_dtype(self, global_layer_id: int) -> DataType: + if self.layer_specs is None: + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return self.dtype + return self.get_layer_cache_spec(global_layer_id).dtype + + def get_layer_buffer_roles(self, global_layer_id: int) -> tuple[DataRole, ...]: + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return self._get_layer_roles_by_local_id(self.layer_offsets[global_layer_id]) + + def get_layer_kv_factor(self, global_layer_id: int) -> int: + if global_layer_id not in self.layer_offsets: + raise KeyError(f"global layer {global_layer_id} is not local to this rank") + return self._get_layer_kv_factor_by_local_id(self.layer_offsets[global_layer_id]) + + @property + def target_layer_ids(self) -> tuple[int, ...]: + return tuple( + layer_id + for layer_id in self.pp_layers + if self.get_layer_owner(layer_id) is KVCacheLayerOwner.TARGET + ) + + @property + def draft_layer_ids(self) -> tuple[int, ...]: + return tuple( + layer_id + for layer_id in self.pp_layers + if self.get_layer_owner(layer_id) is KVCacheLayerOwner.DRAFT + ) + + def get_draft_kv_cache_view(self) -> Optional["KVCacheManagerV2"]: + """Return this manager when explicit draft layers share its request timeline.""" + if self.layer_specs is None or not self.draft_layer_ids: + return None + return self + @property def draft_prompt_lookahead(self) -> int: """Prompt tokens the configured draft reads ahead, before reuse-backoff opt-outs.""" @@ -3578,11 +3936,13 @@ def get_buffers( *, mixed_precision_pool: Literal["low", "high"] = "low", ) -> Optional[torch.Tensor]: + """Return a zero-copy KV view; its slot dimension may be strided.""" layer_offset = self.layer_offsets[layer_idx] if layer_offset in self.high_layers: return self._get_mixed_precision_buffers(layer_offset, kv_layout, mixed_precision_pool) addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, PageIndexMode.SHARED) - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + layer_kv_factor = self._get_layer_kv_factor_by_local_id(layer_offset) + if layer_kv_factor == 2: addr_value = self.impl.get_mem_pool_base_address( layer_offset, Role.VALUE, PageIndexMode.SHARED ) @@ -3594,7 +3954,7 @@ def get_buffers( assert kv_layout in ["NHD", "HND"], f"Unsupported kv_layout: {kv_layout}" element_per_container = 1 - dtype = self.dtype + dtype = self._get_layer_dtype_by_local_id(layer_offset) if dtype == DataType.NVFP4: element_per_container = 2 dtype = torch.int8 @@ -3602,28 +3962,42 @@ def get_buffers( layer_head_dim = self.head_dim_per_layer[layer_offset] if kv_layout == "NHD": shape = [ - self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // self.kv_factor, - self.kv_factor, + self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // layer_kv_factor, + layer_kv_factor, self.tokens_per_block, self.num_kv_heads_per_layer[layer_offset], layer_head_dim // element_per_container, ] else: shape = [ - self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // self.kv_factor, - self.kv_factor, + self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) // layer_kv_factor, + layer_kv_factor, self.num_kv_heads_per_layer[layer_offset], self.tokens_per_block, layer_head_dim // element_per_container, ] - return convert_to_torch_tensor( - TensorWrapper( - addr_key, - dtype, - shape, + wrapper = TensorWrapper(addr_key, dtype, shape) + scale = self.get_layer_page_index_scale(layer_idx) + view_stride = self._get_cache_view_page_stride(scale, layer_kv_factor) + if view_stride != layer_kv_factor: + # A slot containing, for example, three buffers cannot be indexed + # as contiguous K/V pairs. Expose only this layer's buffers with + # the physical slot stride, indexed by raw slot IDs. + converter = self.impl.get_page_index_converter(layer_offset, Role.KEY) + page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) + shape[0] = exact_div(page_upper + int(converter.layer_offset), scale) + element_size = torch.empty(0, dtype=wrapper.dtype).element_size() + page_elements = exact_div( + self.impl.get_page_stride(layer_offset, Role.KEY), element_size ) - ) + wrapper.shape = shape + wrapper.strides = [ + view_stride * page_elements, + page_elements, + *[math.prod(shape[i + 1 :]) for i in range(2, len(shape))], + ] + return convert_to_torch_tensor(wrapper) def get_index_k_buffer( self, @@ -3774,18 +4148,33 @@ def get_num_available_tokens( return clamped def get_num_free_blocks(self) -> int: - # NOTE This method is used to get the number of blocks in the primary pool not the FREE blocks. - # However, since we only use this function when the kv cache manager is empty, so it is safe to do so. - # - # "Empty" means no request holds a cache. The guard page, when it is - # enabled, is held by a permanent reservation made at construction, not - # by a request, so it is excluded from the check and subtracted from - # the count: a caller sizing itself off this number must not plan to - # use the page the guard owns. + """Estimate warmup block capacity when only permanent reservations remain. + + Explicit specs count independent one-block sequences across all required + GPU pools. Longer sequences and resume limits are checked during allocation. + """ reserved = set(self.kv_cache_map) & _RESERVED_REQUEST_IDS assert not set(self.kv_cache_map) - reserved, ( "get_num_free_blocks is only used when the kv cache manager is empty" ) + if self.layer_specs is not None: + # Each lifecycle needs one slot per block. Coalesced layers share a + # lifecycle, while separate lifecycles can compete for the same pool. + slots_per_block: dict[int, int] = defaultdict(int) + for pool_id in self.impl.get_life_cycle_pool_group_indices(GPU_LEVEL): + slots_per_block[int(pool_id)] += 1 + + # Available slots already exclude the guard's permanent reservation. + stats = self._get_storage_statistics(GPU_LEVEL) + return min( + ( + int(stats[pool_id].available) // slot_count + for pool_id, slot_count in slots_per_block.items() + ), + default=0, + ) + + # Legacy managers estimate primary-pool capacity from page indices. max_num_pages = max( [ self.impl.get_page_index_upper_bound(layer_id, Role.KEY) @@ -5876,10 +6265,17 @@ def _connector_reports_cached_tokens(self) -> bool: ) def get_block_ids_per_seq(self, request_ids: List[int]) -> torch.Tensor: - block_ids_per_seq = self.get_batch_cache_indices(request_ids) + if self.layer_specs is not None: + block_ids_per_seq = self.get_batch_cache_indices(request_ids, raw_indices=True) + div_factor = 1 + else: + block_ids_per_seq = self.get_batch_cache_indices(request_ids) + div_factor = self.num_local_layers + if self.kv_factor > 1 and self._index_scale_ints[0] % self.kv_factor != 0: + div_factor = 1 block_ids_per_seq_tensors = [ torch.tensor( - [i // self.num_local_layers if i != BAD_PAGE_INDEX else 0 for i in sublist], + [i // div_factor if i != BAD_PAGE_INDEX else 0 for i in sublist], dtype=torch.int, ) for sublist in block_ids_per_seq @@ -6256,6 +6652,15 @@ def get_layer_page_index_scale(self, layer_idx: int) -> int: must not use the pool-level scale.""" return int(self.impl.get_page_index_scale(self.layer_offsets[layer_idx], Role.KEY)) + @staticmethod + def _get_cache_view_page_stride(index_scale: int, kv_factor: int) -> int: + """Physical buffer pages between rows of the KV tensor view. + + Contiguous K/V rows are usable only when every slot starts on a row + boundary. Otherwise, rows span whole slots and use unscaled slot IDs. + """ + return kv_factor if index_scale % kv_factor == 0 else index_scale + def get_batch_cache_indices( self, request_ids: List[int], @@ -6264,7 +6669,10 @@ def get_batch_cache_indices( *, raw_indices: bool = False, ) -> List[List[int]]: - """Return cache indices for a layer, or pool 0 when no layer is given. + """Return indices into a layer's ``get_buffers()`` view. + + With explicit specs, omitting the layer selects pool 0's KEY + representative. Callers accessing other layers must pass their IDs. Set raw_indices for slot-major views that need base slot IDs without page-index scaling or KV aggregation. @@ -6272,15 +6680,22 @@ def get_batch_cache_indices( if layer_idx is None: pool_id = 0 index_scale = 1 if raw_indices else None + kv_factor = ( + self._get_layer_kv_factor_by_local_id(self._pool_layer_ids_by_role[(0, Role.KEY)]) + if self.layer_specs is not None + else self.kv_factor + ) else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] index_scale = 1 if raw_indices else self.get_layer_page_index_scale(layer_idx) + kv_factor = self.get_layer_kv_factor(layer_idx) return self._get_batch_cache_indices_by_pool_id( request_ids, pool_id=pool_id, is_kv_aggregate=not raw_indices, num_blocks_per_seq=num_blocks_per_seq, index_scale=index_scale, + kv_factor=kv_factor, ) def _get_batch_cache_indices_by_pool_id( @@ -6291,16 +6706,17 @@ def _get_batch_cache_indices_by_pool_id( is_kv_aggregate: bool = True, num_blocks_per_seq: Optional[Sequence[int]] = None, index_scale: Optional[int] = None, + kv_factor: Optional[int] = None, ) -> List[List[int]]: - if is_kv_aggregate: - # Div by kv_factor to index kv cache with size - # [num_blocks, kv_factor, tokens_per_block, num_kv_heads, head_dim] - div_factor = self.kv_factor - else: - div_factor = 1 - if index_scale is None: index_scale = int(self.index_scales[pool_id]) + div_factor = ( + self._get_cache_view_page_stride( + index_scale, self.kv_factor if kv_factor is None else kv_factor + ) + if is_kv_aggregate + else 1 + ) res = [] for req_idx, req_id in enumerate(request_ids): @@ -6344,6 +6760,11 @@ def get_batch_cache_indices_flat( if layer_idx is None: pool_id = 0 scale = self._index_scale_ints[pool_id] + kv_factor = ( + self._get_layer_kv_factor_by_local_id(self._pool_layer_ids_by_role[(0, Role.KEY)]) + if self.layer_specs is not None + else self.kv_factor + ) else: pool_id = self.layer_to_pool_mapping_dict[self.layer_offsets[layer_idx]] # Layers sharing a pool can still require different page-index @@ -6351,7 +6772,8 @@ def get_batch_cache_indices_flat( # scale so this flat block table matches get_batch_cache_indices() # and never feeds out-of-range page ids to FlashInfer. scale = self.get_layer_page_index_scale(layer_idx) - div_factor = self.kv_factor + kv_factor = self.get_layer_kv_factor(layer_idx) + div_factor = self._get_cache_view_page_stride(scale, kv_factor) out_tensor = torch.empty(sum(num_blocks), dtype=torch.int32, pin_memory=prefer_pinned()) out = out_tensor.numpy() @@ -6372,6 +6794,8 @@ def get_batch_cache_indices_flat( return out_tensor def get_cache_bytes_per_token(self) -> int: + if self.layer_specs is not None: + return sum(layer.bytes_per_token() for layer in self.layer_specs) data_roles = [Role.KEY] if self.kv_cache_type != CacheTypeCpp.SELFKONLY: data_roles.append(Role.VALUE) @@ -6387,6 +6811,12 @@ def get_cache_bytes_per_token(self) -> int: ) def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: Role): + if self.layer_specs is not None: + layer = self.layer_specs[local_layer_idx] + if data_role == Role.ALL: + return layer.bytes_per_token() + return layer.bytes_per_token(data_role) + if self._is_mixed_precision_layer(local_layer_idx): return self._mixed_precision_layer_bytes_per_token(local_layer_idx, data_role) if self.dtype not in ( @@ -6451,12 +6881,26 @@ def calculate_scaling_factor_size_bytes( def _iter_cache_buffers_for_invalid_check(self) -> Iterable[Tuple[int, torch.Tensor]]: pool_handled = set() for layer_id, layer_offset in self.layer_offsets.items(): - pool_id = self.layer_to_pool_mapping_dict[layer_offset] - if pool_id in pool_handled: - continue buffer = self.get_buffers(layer_id) if buffer is None: continue + if self.layer_specs is not None: + # Inspect each explicit layer's own rows, without reading + # another layer's bytes through the wrong dtype. + scale = self.get_layer_page_index_scale(layer_id) + view_stride = self._get_cache_view_page_stride( + scale, self.get_layer_kv_factor(layer_id) + ) + yield layer_id, buffer[:: exact_div(scale, view_stride)] + continue + if buffer.stride(0) != math.prod(buffer.shape[1:]): + # Slot-strided views cover only one layer, including legacy + # layouts that coalesce extra buffers alongside K/V. + yield layer_id, buffer + continue + pool_id = self.layer_to_pool_mapping_dict[layer_offset] + if pool_id in pool_handled: + continue yield layer_id, buffer pool_handled.add(pool_id) @@ -6481,10 +6925,16 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool # before inference, but that must not wipe the guard-page # sentinel -- doing so would leave the startup warning claiming # a guard that no longer exists. Rewrite the guard page after - # the zero. Each pool is zeroed on one representative layer, so - # only that layer's guard page needs restoring here. + # the zero. Explicit specs expose one view per layer here; + # contiguous legacy views use a single representative. guard_page = self._guard_page_by_layer.get(layer_id) if guard_page is not None and self._guard_page_value is not None: + if self.layer_specs is not None: + scale = self.get_layer_page_index_scale(layer_id) + view_stride = self._get_cache_view_page_stride( + scale, self.get_layer_kv_factor(layer_id) + ) + guard_page = exact_div(guard_page, exact_div(scale, view_stride)) _fill_kv_pages(buffer, [guard_page], self._guard_page_value) torch.cuda.synchronize() @@ -6556,20 +7006,26 @@ def get_cache_size_per_token( max_num_tokens: int = 0, spec_config=None, is_draft: bool = False, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ): - layer_sizes, attention_windows = _get_static_cache_size_layer_components( - model_config, - mapping, - num_layers=num_layers, - max_seq_len=max_seq_len, - kv_cache_config=kv_cache_config, - is_external_draft=( - is_draft - and spec_config is not None - and spec_config.spec_dec_mode.is_external_drafter() - ), - ) + if layer_specs is not None: + layer_sizes, attention_windows = _get_layer_spec_cache_size_components( + layer_specs, max_seq_len=max_seq_len + ) + else: + layer_sizes, attention_windows = _get_static_cache_size_layer_components( + model_config, + mapping, + num_layers=num_layers, + max_seq_len=max_seq_len, + kv_cache_config=kv_cache_config, + is_external_draft=( + is_draft + and spec_config is not None + and spec_config.spec_dec_mode.is_external_drafter() + ), + ) reuse_backoff_enabled = ( kv_cache_config is not None and kv_cache_config.enable_block_reuse @@ -6600,8 +7056,10 @@ def get_cache_size_per_token( kv_cache_config is not None and kv_cache_config.enable_swa_scratch_reuse and not is_draft + and not any(spec.owner == KVCacheLayerOwner.DRAFT for spec in (layer_specs or ())) ), generation_capacity_headroom=generation_capacity_headroom, + layer_specs=layer_specs, ) # The affine slope covers all tokens; context additionally retains SWA # pages for the current token batch beyond the generation windows. @@ -6612,16 +7070,43 @@ def get_cache_size_per_token( bytes_per_slot = _get_single_swa_pool_slot_bytes( layer_sizes, attention_windows, tokens_per_block ) - if is_draft and fixed_cost > 0 and bytes_per_slot is not None: + if is_draft and fixed_cost > 0: + costs = [(fixed_cost, bytes_per_slot)] + if layer_specs is not None: + # Explicit lifecycles can have different slot counts. Normalize + # each layer separately so pool rounding cannot hide a deficit. + costs = [] + for layer, layer_size, window in zip(layer_specs, layer_sizes, attention_windows): + if window is None or window <= 0: + continue + _, retained_bytes = _estimate_swa_cache_size( + [layer_size], + [window], + tokens_per_block, + context=False, + scratch=False, + generation_capacity_headroom=generation_capacity_headroom, + layer_specs=[layer], + ) + costs.append( + ( + layer_size * max_num_tokens + retained_bytes * max_batch_size, + layer_size * tokens_per_block, + ) + ) # Without a config, all windows are full attention, so no SWA slot exists. assert kv_cache_config is not None # The affine intercept is the configured quota needed to preserve # the fixed usable capacity at V2's resume watermark. Keep the # allocator's page rounding inside the manager estimator rather # than extending the generic CacheCost model with pool geometry. - required_slots = math.ceil(fixed_cost / bytes_per_slot) resume_util = float(np.float32(kv_cache_config.max_util_for_resume)) - fixed_cost = math.ceil(required_slots / resume_util) * bytes_per_slot + fixed_cost = sum( + cost + if slot_bytes is None + else math.ceil(math.ceil(cost / slot_bytes) / resume_util) * slot_bytes + for cost, slot_bytes in costs + ) return ( cache_size_per_token, fixed_cost, @@ -6639,6 +7124,7 @@ def _get_per_request_layer_tokens( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, ) -> Optional[tuple[List[tuple[int, int, int]], int]]: """Per-local-layer ``(layer_size, retained_tokens, intercept_tokens)`` plus the block-rounded full-sequence token count. @@ -6665,7 +7151,7 @@ def _get_per_request_layer_tokens( """ if max_seq_len is None or int(max_seq_len) <= 0: return None - if num_layers is None: + if layer_specs is None and num_layers is None: total_hidden_layers = getattr(model_config.pretrained_config, "num_hidden_layers", None) if ( total_hidden_layers is not None @@ -6674,18 +7160,23 @@ def _get_per_request_layer_tokens( return None max_seq_len = int(max_seq_len) tokens_per_block = int(tokens_per_block) - layer_sizes, covered_windows = _get_static_cache_size_layer_components( - model_config, - mapping, - num_layers=num_layers, - max_seq_len=max_seq_len, - kv_cache_config=kv_cache_config, - is_external_draft=( - is_draft - and spec_config is not None - and spec_config.spec_dec_mode.is_external_drafter() - ), - ) + if layer_specs is not None: + layer_sizes, covered_windows = _get_layer_spec_cache_size_components( + layer_specs, max_seq_len=max_seq_len + ) + else: + layer_sizes, covered_windows = _get_static_cache_size_layer_components( + model_config, + mapping, + num_layers=num_layers, + max_seq_len=max_seq_len, + kv_cache_config=kv_cache_config, + is_external_draft=( + is_draft + and spec_config is not None + and spec_config.spec_dec_mode.is_external_drafter() + ), + ) if not layer_sizes: return None # ``covered_windows`` are the windows ``get_cache_size_per_token`` @@ -6706,7 +7197,7 @@ def _get_per_request_layer_tokens( # against the draft pool), so the drafter's own windows do not # bound per-request retention. retained_windows = [None] * len(layer_sizes) - elif all(window is None for window in covered_windows): + elif layer_specs is None and all(window is None for window in covered_windows): # The static estimator only sees windows the config carries. # Runtime manager creation additionally derives per-layer windows # from ``layer_types`` (see @@ -6744,7 +7235,7 @@ def _get_per_request_layer_tokens( * tokens_per_block ) - def window_tokens(window_size: Optional[int]) -> int: + def window_tokens(window_size: Optional[int], sink_tokens: int) -> int: if window_size is None or window_size <= 0 or window_size >= max_seq_len: return full_tokens # Match _estimate_swa_cache_size: the live interval holds @@ -6752,15 +7243,17 @@ def window_tokens(window_size: Optional[int]) -> int: window_blocks = ( math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1 ) - return min(window_blocks * tokens_per_block, full_tokens) + sink_blocks = math.ceil(sink_tokens / tokens_per_block) + return min((sink_blocks + window_blocks) * tokens_per_block, full_tokens) entries = [] - for layer_size, retained_window, covered_window in zip( - layer_sizes, retained_windows, covered_windows + for layer_idx, (layer_size, retained_window, covered_window) in enumerate( + zip(layer_sizes, retained_windows, covered_windows) ): - retained = window_tokens(retained_window) + sink_tokens = (layer_specs[layer_idx].num_sink_tokens or 0) if layer_specs else 0 + retained = window_tokens(retained_window, sink_tokens) covered = ( - window_tokens(covered_window) + window_tokens(covered_window, sink_tokens) if covered_window is not None and 0 < covered_window < max_seq_len else 0 ) @@ -6779,6 +7272,7 @@ def get_cache_bytes_per_request( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> Optional[int]: """Slope-funded pool bytes ONE max-length request pins in this manager. @@ -6811,6 +7305,7 @@ def get_cache_bytes_per_request( max_seq_len=max_seq_len, spec_config=spec_config, is_draft=is_draft, + layer_specs=layer_specs, ) if layer_tokens is None: return None @@ -6839,6 +7334,7 @@ def get_cache_quota_per_request( max_seq_len: Optional[int] = None, spec_config=None, is_draft: bool = False, + layer_specs: Optional[Sequence[KVCacheLayerSpec]] = None, **kwargs, ) -> Optional[int]: """Allocator quota ONE max-length request makes this manager insist on. @@ -6873,6 +7369,7 @@ def get_cache_quota_per_request( max_seq_len=max_seq_len, spec_config=spec_config, is_draft=is_draft, + layer_specs=layer_specs, ) if layer_tokens is None: return None diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 2f1644a37c31..e418a261853e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -263,6 +263,7 @@ class BufferConfig: size: int tokens_per_block_override: int | None = None is_sparse: bool = False + coalescing_group: int = 0 class AttentionReusePolicy(enum.IntEnum): REQUIRED = 0 diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 1aad3d7d79c4..d16299ede1b9 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -1893,7 +1893,7 @@ def __init__(self): class MockMetadata(DSAtrtllmAttentionMetadata): """Mock DSA attention metadata for unit testing indexer operations.""" - def __init__(self, cache_manager): + def __init__(self, cache_manager: DSACacheManager | DSACacheManagerV2) -> None: """Initialize mock metadata with buffers, block tables, and indexer state. `cache_manager` is a parameter rather than a capture of the enclosing @@ -1932,7 +1932,7 @@ def __init__(self, cache_manager): self._tokens_per_block = tpb self.kv_lens_cuda_runtime = kv_lens.cuda() self.indexer_k_cache_block_offsets = torch.zeros( - [batch_size, cache_manager.max_blocks_per_seq], + [batch_size, self.kv_cache_manager.max_blocks_per_seq], device="cuda", dtype=torch.int32, ) @@ -1941,7 +1941,7 @@ def __init__(self, cache_manager): device="cpu", pin_memory=True, ) - block_ids = cache_manager.get_batch_cache_indices(request_ids) + block_ids = self.kv_cache_manager.get_batch_cache_indices(request_ids) for i in range(len(block_ids)): self.host_indexer_k_cache_block_offsets[i, : len(block_ids[i])] = torch.tensor( block_ids[i], dtype=torch.int32 diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index e989db333636..e522fd422acf 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -146,6 +146,7 @@ def test_indexer_runtime_cache_cost_matches_registered_buffers(fp8, snapshot_int from tensorrt_llm.mapping import Mapping manager = object.__new__(Glm5NextCacheManager) + manager.layer_specs = None manager.mapping = Mapping(world_size=1, tp_size=1) manager.pp_layers = [0, 1] manager.layer_offsets = {0: 0, 1: 1} diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py index 0225dea33702..19f9bfdba7ab 100644 --- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py +++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py @@ -173,6 +173,7 @@ def test_nvfp4_shared_draft_constructs_and_converts_heterogeneous_page_tables( @pytest.mark.parametrize("dtype", [DataType.FP8, DataType.NVFP4]) def test_block_scale_role_uses_global_layer_ids(dtype: DataType) -> None: manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2) + manager.layer_specs = None manager.dtype = dtype manager.pp_layers = [60, 7, 2] manager.sparse_layer_ids = {7} @@ -181,6 +182,7 @@ def test_block_scale_role_uses_global_layer_ids(dtype: DataType) -> None: assert manager._get_block_scale_role(Role.KEY, local_layer) == expected assert manager._get_block_scale_role(Role.VALUE, local_layer) is None generic = KVCacheManagerV2.__new__(KVCacheManagerV2) + generic.layer_specs = None generic.dtype = dtype expected = Role.KEY_BLOCK_SCALE if dtype == DataType.NVFP4 else None assert generic._get_block_scale_role(Role.KEY, 0) == expected diff --git a/tests/unittest/_torch/attention/test_fp4_mla.py b/tests/unittest/_torch/attention/test_fp4_mla.py index 0667afe93f91..8aab22347603 100644 --- a/tests/unittest/_torch/attention/test_fp4_mla.py +++ b/tests/unittest/_torch/attention/test_fp4_mla.py @@ -675,6 +675,7 @@ def test_fp4_mla_v2_cache_size_accounts_for_bf16_fallback( def test_fp4_mla_v2_runtime_sizing_accounts_for_pipeline_slots() -> None: manager = object.__new__(Fp4MlaKVCacheManagerV2) + manager.layer_specs = None manager.max_batch_size = 3 manager.mapping = SimpleNamespace(pp_size=2) manager.max_num_tokens = 100 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 7fd42932ef73..49b226eef64e 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -11,7 +11,7 @@ """ from contextlib import nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, replace from types import SimpleNamespace from unittest.mock import MagicMock, Mock, patch @@ -20,11 +20,8 @@ import tensorrt_llm.bindings.internal.batch_manager as batch_manager from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures -from tensorrt_llm._torch.model_config import ( - KVCacheLayerSpec, - ModelConfig, - _get_kv_cache_layer_specs, -) +from tensorrt_llm._torch.model_config import KVCacheLayerSpec as ModelKVCacheLayerSpec +from tensorrt_llm._torch.model_config import ModelConfig, _get_kv_cache_layer_specs from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor._util import ( CacheCost, @@ -36,12 +33,16 @@ ) from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + KVCacheLayerOwner, + KVCacheLayerSpec, KVCacheManagerV2, + Role, _get_static_cache_size_layer_components, ) from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode +from tensorrt_llm.bindings import DataType from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( KvCacheConfig, @@ -50,6 +51,7 @@ TorchLlmArgs, ) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import AttentionReusePolicy pytestmark = pytest.mark.cpu_only @@ -683,7 +685,7 @@ def test_gemma4_hybrid_scales_by_num_pool_groups(): sliding_window=sliding_window, ) hybrid._model_engine.model.model_config.kv_cache_layer_specs = tuple( - KVCacheLayerSpec( + ModelKVCacheLayerSpec( head_dim=256 if layer_type == "sliding_attention" else 128, num_kv_heads=1, num_q_heads=1, @@ -831,7 +833,7 @@ def test_pool_group_estimation_uses_resolved_spec_windows( sliding_window=512, ), kv_cache_layer_specs=tuple( - KVCacheLayerSpec( + ModelKVCacheLayerSpec( head_dim=head_dim, num_q_heads=8, num_kv_heads=2, attention_window=window ) for head_dim, window in [(128, 128), (256, 256), (128, None)] @@ -883,7 +885,7 @@ def test_estimation_budget_uses_temporary_pool_ratio() -> None: model_config.attn_backend = "FLASHINFER" model_config.is_encoder_decoder = False model_config.kv_cache_layer_specs = tuple( - KVCacheLayerSpec( + ModelKVCacheLayerSpec( head_dim=head_dim, num_q_heads=8, num_kv_heads=2, @@ -960,7 +962,7 @@ def test_pool_scaling_prevents_mmmu_pro_underestimation(): sliding_window=sliding_window, ) c._model_engine.model.model_config.kv_cache_layer_specs = tuple( - KVCacheLayerSpec( + ModelKVCacheLayerSpec( head_dim=256 if layer_type == "sliding_attention" else 128, num_kv_heads=1, num_q_heads=1, @@ -1147,7 +1149,7 @@ def test_v2_static_sizing_uses_model_provided_per_layer_geometry( expected_sliding_layer_size: int, ) -> None: layer_specs = tuple( - KVCacheLayerSpec( + ModelKVCacheLayerSpec( head_dim=256 if (layer_idx + 1) % 6 else 512, num_kv_heads=8 if (layer_idx + 1) % 6 else 1, num_q_heads=16, @@ -1247,8 +1249,70 @@ def get_cache_size_per_token(model_config, mapping, **kwargs): assert cost == CacheCost(slope=20, intercept=21) +def test_v2_layer_spec_costing_preserves_sink_and_reuse_policy() -> None: + windowed = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.BF16, + num_kv_heads=1, + head_dim=16, + buffer_roles=(Role.KEY, Role.VALUE), + sliding_window_size=128, + num_sink_tokens=65, + ) + layers = [ + windowed, + replace(windowed, global_layer_id=1, reuse_policy=AttentionReusePolicy.PRIVATE), + replace(windowed, global_layer_id=2, num_sink_tokens=None), + replace(windowed, global_layer_id=3, sliding_window_size=None, num_sink_tokens=None), + ] + config = KvCacheConfig(enable_block_reuse=False, enable_swa_scratch_reuse=True) + kwargs = dict( + tokens_per_block=64, + max_seq_len=4096, + kv_cache_config=config, + layer_specs=layers, + ) + slope, fixed = KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=256, **kwargs + ) + # Each layer costs 64 bytes/token. Different sinks or policies cannot share scratch. + # The SWA layers retain 5, 5, and 3 pages, including two pages for 65 sink tokens. + assert slope == 64 + assert fixed == 256 * 192 + 2 * (320 + 320 + 192) * 64 + + manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = tuple(layers) + manager._has_cp_helix = False + manager.max_attention_window_vec = [128, 128, 128, None] + manager.tokens_per_block = 64 + manager.max_batch_size = 2 + manager.max_num_tokens = 256 + manager._generation_kv_capacity_headroom = 1 + manager.enable_swa_scratch_reuse = True + quota = manager._get_quota_from_max_tokens(512) + assert quota == slope * 512 + fixed + assert manager._get_max_tokens_from_quota(quota) == 512 + assert KVCacheManagerV2._get_per_request_layer_tokens(None, Mapping(), **kwargs) == ( + [(64, 320, 320), (64, 320, 320), (64, 192, 192), (64, 4160, 0)], + 4160, + ) + + # Draft lifecycles need 7 and 21 usable slots: separate resume rounding + # gives 8 + 23 slots; rounding their average together would give only 30. + draft = replace(windowed, owner=KVCacheLayerOwner.DRAFT, num_sink_tokens=0) + kwargs.update( + layer_specs=[draft, replace(draft, global_layer_id=1, num_sink_tokens=448)], + kv_cache_config=KvCacheConfig(enable_block_reuse=False, max_util_for_resume=0.95), + ) + assert KVCacheManagerV2.get_cache_size_per_token( + None, Mapping(), max_batch_size=2, max_num_tokens=64, is_draft=True, **kwargs + ) == (0, (8 + 23) * 64 * 64) + + def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(KVCacheManagerV2) + manager.layer_specs = None manager._has_cp_helix = False manager.num_local_layers = 3 manager.pp_layers = [4, 5, 6] @@ -1768,7 +1832,7 @@ def __init__(self, kv_cache_config: KvCacheConfig, *args: object, **kwargs: obje model_config.quant_config = None if publish_layer_specs: model_config.kv_cache_layer_specs = tuple( - KVCacheLayerSpec( + ModelKVCacheLayerSpec( head_dim=256 if layer_type == _SLIDING else 512, num_kv_heads=8 if layer_type == _SLIDING else 1, num_q_heads=16, diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index 10d1ab4408f0..addb326f0a91 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -28,6 +28,8 @@ from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as kv_cache_v2_module from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( BlockReusePolicy, + KVCacheLayerOwner, + KVCacheLayerSpec, KVCacheManagerV2, Role, _extend_swa_windows_for_reuse, @@ -43,6 +45,7 @@ from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.bindings.executor import KvCacheRetentionConfig +from tensorrt_llm.bindings.internal import thop from tensorrt_llm.bindings.internal.batch_manager import CacheType, LinearCacheType from tensorrt_llm.conversation_params import ConversationParams from tensorrt_llm.llmapi.llm_args import ( @@ -54,6 +57,7 @@ PARDDecodingConfig, ) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.quantization import QuantMode from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BAD_PAGE_INDEX, DEFAULT_BEAM_INDEX, @@ -72,6 +76,7 @@ KVCacheManager, KVCacheManagerConfig, LayerId, + PageIndexMode, ReuseScope, SsmLayerConfig, _introspection, @@ -81,6 +86,19 @@ MAX_SEQ_LEN = 16 +@pytest.mark.cpu_only +def test_nvfp4_layer_spec_rejects_odd_head_dim() -> None: + with pytest.raises(ValueError, match="NVFP4 head_dim must be even"): + KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.NVFP4, + num_kv_heads=16, + head_dim=1, + buffer_roles=(Role.KEY, Role.KEY_BLOCK_SCALE), + ) + + @pytest.mark.cpu_only @pytest.mark.parametrize( "prompt,rewind,ordinal,window,sinks,expected", @@ -315,6 +333,7 @@ def _make_cache_config_for_test( assert len(max_attention_window_vec) == len(pp_layers) cache_manager = object.__new__(KVCacheManagerV2) + cache_manager.layer_specs = None cache_manager.max_beam_width = 1 cache_manager.kv_cache_type = kv_cache_type cache_manager.dtype = dtype @@ -364,6 +383,8 @@ def _make_manager_for_cache_tier_test( joint_reuse: bool = False, mapping: Mapping | None = None, vocab_size: int | None = 16, + dtype: DataType = DataType.HALF, + layer_specs: list[KVCacheLayerSpec] | None = None, ) -> tuple[KVCacheManagerV2, Mock]: impl_constructor = Mock(side_effect=impl_side_effect) if mapping is None: @@ -426,7 +447,8 @@ def build_cache_config( max_seq_len=MAX_SEQ_LEN, max_batch_size=1, mapping=mapping, - dtype=DataType.HALF, + dtype=dtype, + layer_specs=layer_specs, spec_config=spec_config, is_draft=is_draft, is_disagg=is_disagg, @@ -636,6 +658,229 @@ def test_zero_size_layers_are_removed_before_pool_allocation( manager.shutdown() +@pytest.mark.parametrize("kv_first", [False, True], ids=["key-first", "kv-first"]) +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"]) +def test_mixed_role_coalesced_pool_preserves_slot_stride(kv_first: bool, kv_layout: str) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + key_layer_id, kv_layer_id = 0, 1 + key_spec = KVCacheLayerSpec( + global_layer_id=key_layer_id, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.HALF, + num_kv_heads=2, + head_dim=8, + buffer_roles=(Role.KEY,), + ) + kv_spec = replace(key_spec, global_layer_id=kv_layer_id, buffer_roles=(Role.KEY, Role.VALUE)) + layer_specs = [kv_spec, key_spec] if kv_first else [key_spec, kv_spec] + manager = KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=2, + head_dim=8, + tokens_per_block=TOKENS_PER_BLOCK, + max_seq_len=MAX_SEQ_LEN, + max_batch_size=1, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + layer_specs=layer_specs, + ) + try: + key_local_id = manager.layer_offsets[key_layer_id] + kv_local_id = manager.layer_offsets[kv_layer_id] + pool_id = manager.layer_to_pool_mapping_dict[key_local_id] + assert manager.layer_to_pool_mapping_dict[kv_local_id] == pool_id + assert manager.num_pools == 1 + assert manager.get_layer_page_index_scale(key_layer_id) == 3 + assert manager.get_layer_page_index_scale(kv_layer_id) == 3 + + key_buffer = manager.get_buffers(key_layer_id, kv_layout=kv_layout) + kv_buffer = manager.get_buffers(kv_layer_id, kv_layout=kv_layout) + page_stride = manager.impl.get_page_stride(kv_local_id, Role.KEY) + key_base = manager.impl.get_mem_pool_base_address( + kv_local_id, Role.KEY, PageIndexMode.SHARED + ) + value_base = manager.impl.get_mem_pool_base_address( + kv_local_id, Role.VALUE, PageIndexMode.SHARED + ) + page_shape = (TOKENS_PER_BLOCK, 2, 8) if kv_layout == "NHD" else (2, TOKENS_PER_BLOCK, 8) + assert kv_buffer.shape[1:] == (2, *page_shape) + assert kv_buffer.stride(0) * kv_buffer.element_size() == 3 * page_stride + # Slot 1 begins three physical pages later, even though this layer + # exposes only two roles. Reversing the specs also exercises a base + # pointer offset within the coalesced slot. + assert kv_buffer[1, 0].data_ptr() == key_base + 3 * page_stride + assert kv_buffer[1, 1].data_ptr() == value_base + 3 * page_stride + assert key_buffer.stride(0) * key_buffer.element_size() == page_stride + + request_id = 7 + cache = manager._create_kv_cache(request_id, None, None) + assert cache.resume(manager._stream.cuda_stream) + assert cache.resize(3 * TOKENS_PER_BLOCK, 0) + raw_slots = [int(slot) for slot in cache.get_base_page_indices(pool_id)[: cache.num_blocks]] + assert len(raw_slots) == 3 + assert any(slot % 2 == 1 for slot in raw_slots) + for layer_id, buffer, expected in ( + (key_layer_id, key_buffer, [3 * slot for slot in raw_slots]), + (kv_layer_id, kv_buffer, raw_slots), + ): + batch_indices = manager.get_batch_cache_indices([request_id], layer_idx=layer_id) + flat_indices = manager.get_batch_cache_indices_flat( + [request_id], [len(raw_slots)], layer_idx=layer_id + ).tolist() + assert batch_indices == [expected] + assert flat_indices == expected + local_id = manager.layer_offsets[layer_id] + for indices in (batch_indices[0], flat_indices): + for slot, index in zip(raw_slots, indices): + for role_index, role in enumerate(manager.get_layer_buffer_roles(layer_id)): + role_base = manager.impl.get_mem_pool_base_address( + local_id, role, PageIndexMode.SHARED + ) + assert ( + buffer[index, role_index].data_ptr() + == role_base + 3 * slot * page_stride + ) + + # Give each role and slot a distinct value. Writing an odd K/V slot + # must preserve the K-only page and both neighboring K/V slots. + for ordinal, slot in enumerate(raw_slots): + key_buffer[3 * slot, 0].fill_(10 + ordinal) + for ordinal, slot in enumerate(raw_slots): + kv_buffer[slot, 0].fill_(20 + ordinal) + kv_buffer[slot, 1].fill_(30 + ordinal) + for ordinal, slot in enumerate(raw_slots): + assert torch.all(key_buffer[3 * slot, 0] == 10 + ordinal) + assert torch.all(kv_buffer[slot, 0] == 20 + ordinal) + assert torch.all(kv_buffer[slot, 1] == 30 + ordinal) + finally: + manager.shutdown() + + +def test_nvfp4_scale_indices_with_mixed_bf16_storage() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + torch.cuda.init() + nvfp4 = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=DataType.NVFP4, + num_kv_heads=4, + head_dim=128, + buffer_roles=(Role.KEY, Role.VALUE, Role.KEY_BLOCK_SCALE, Role.VALUE_BLOCK_SCALE), + ) + bf16 = replace( + nvfp4, + global_layer_id=1, + owner=KVCacheLayerOwner.DRAFT, + dtype=DataType.BF16, + num_kv_heads=1, + buffer_roles=(Role.KEY, Role.VALUE), + ) + manager = KVCacheManagerV2( + KvCacheConfig( + max_gpu_total_bytes=16 << 20, + host_cache_size=0, + enable_block_reuse=False, + enable_swa_scratch_reuse=False, + ), + CacheType.SELF, + num_layers=2, + num_kv_heads=4, + head_dim=128, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=1, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.NVFP4, + layer_specs=[nvfp4, bf16], + ) + try: + assert manager.num_pools == 1 + (group,) = manager.impl.pool_group_descs + assert len(group.pools) == 3 + assert group.num_slots > 1 + slots = [1, int(group.num_slots) - 1] + for offset, (data_role, scale_role) in enumerate( + ((Role.KEY, Role.KEY_BLOCK_SCALE), (Role.VALUE, Role.VALUE_BLOCK_SCALE)) + ): + data_converter = manager.impl.get_page_index_converter(0, data_role) + scale_converter = manager.impl.get_page_index_converter(0, scale_role) + pages = data_converter(slots, PageIndexMode.PER_LAYER) + assert pages == [2 * slot + offset for slot in slots] + assert scale_converter(slots, PageIndexMode.PER_LAYER) == pages + + scale_base = manager.impl.get_mem_pool_base_address( + 0, scale_role, PageIndexMode.PER_LAYER + ) + scale_pool = next(pool for pool in group.pools if pool.base_address == scale_base) + page_bytes = manager.impl.get_page_stride(0, scale_role) + # Attention reuses data indices for scales, including at the final slot. + for slot, page in zip(slots, pages): + address = scale_base + page * page_bytes + assert address == scale_base + slot * scale_pool.slot_bytes + offset * page_bytes + assert address + page_bytes <= scale_base + group.num_slots * scale_pool.slot_bytes + finally: + manager.shutdown() + + +def test_mixed_nvfp4_bf16_pool_pointer_views() -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + bf16 = torch.empty((2, 1, 32, 128), dtype=torch.bfloat16, device="cuda") + nvfp4 = torch.empty((2, 1, 32, 64), dtype=torch.uint8, device="cuda") + scales = torch.empty((2, 1, 32, 8), dtype=torch.float8_e4m3fn, device="cuda") + pointers = torch.tensor( + [ + [[nvfp4.data_ptr(), scales.data_ptr()], [0, 0]], + # A nonzero scale pointer must be ignored for the BF16 layer. + [[bf16.data_ptr(), scales.data_ptr()], [0, 0]], + ], + dtype=torch.int64, + device="cpu", + ) + mapping = torch.tensor([[0, 0], [1, 0]], dtype=torch.int32, device="cpu") + block_offsets = torch.zeros((2, 1, 2, 1), dtype=torch.int32, device="cuda") + for table, layer_id, quant_mode, expected_data, expected_scales in ( + (pointers, 1, 0, bf16, None), + (pointers, 0, int(QuantMode.NVFP4_KV_CACHE), nvfp4, scales), + (pointers[:, :, 0].contiguous(), 1, 0, bf16, None), + ): + data, blocks, block_scales = thop.build_trtllm_gen_kv_cache_metadata( + host_kv_cache_pool_pointers=table, + host_kv_cache_pool_mapping=mapping, + kv_cache_block_offsets=block_offsets, + layer_idx=layer_id, + num_kv_heads=1, + tokens_per_block=32, + head_dim=128, + kv_factor=2, + total_num_blocks=2, + kv_cache_quant_mode=quant_mode, + batch_start=0, + batch_size=1, + dtype=torch.bfloat16, + ) + assert data.data_ptr() == expected_data.data_ptr() + assert data.shape == expected_data.shape + assert data.dtype == expected_data.dtype + assert blocks.data_ptr() == block_offsets[layer_id].data_ptr() + if expected_scales is None: + assert block_scales is None + else: + assert block_scales is not None + assert block_scales.data_ptr() == expected_scales.data_ptr() + assert block_scales.shape == expected_scales.shape + + def test_zero_size_filter_preserves_ssm_and_nonempty_extra_buffers() -> None: manager = object.__new__(KVCacheManagerV2) manager.pp_layers = [3, 4, 5] @@ -1241,6 +1486,85 @@ def test_cold_codec_provider_receives_draft_role() -> None: assert codec_provider.create_cold_page_codec.call_args.kwargs["is_draft"] is True +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "manager_dtype,layer_dtype", + [ + (DataType.HALF, DataType.BF16), + (DataType.BF16, DataType.HALF), + (DataType.HALF, DataType.FP8), + ], +) +def test_cold_codec_provider_receives_layer_spec_dtype( + manager_dtype: DataType, layer_dtype: DataType +) -> None: + codec_provider = Mock() + layer = KVCacheLayerSpec( + global_layer_id=0, + owner=KVCacheLayerOwner.TARGET, + dtype=layer_dtype, + num_kv_heads=1, + head_dim=1, + buffer_roles=(Role.KEY,), + ) + manager, _ = _make_manager_for_cache_tier_test( + KvCacheConfig(max_gpu_total_bytes=16 << 20), + [Mock()], + cold_page_codec_provider=codec_provider, + dtype=manager_dtype, + layer_specs=[layer], + ) + + assert manager.dtype == manager_dtype + codec_provider.create_cold_page_codec.assert_called_once() + assert codec_provider.create_cold_page_codec.call_args.kwargs["runtime_dtype"] == layer_dtype + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("manager_dtype", [DataType.HALF, DataType.BF16]) +def test_cold_codec_provider_receives_legacy_manager_dtype(manager_dtype: DataType) -> None: + codec_provider = Mock() + _make_manager_for_cache_tier_test( + KvCacheConfig(max_gpu_total_bytes=16 << 20), + [Mock()], + cold_page_codec_provider=codec_provider, + dtype=manager_dtype, + ) + + codec_provider.create_cold_page_codec.assert_called_once() + assert codec_provider.create_cold_page_codec.call_args.kwargs["runtime_dtype"] == manager_dtype + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "layer_dtypes", [(DataType.BF16, DataType.HALF), (DataType.BF16, DataType.FP8)] +) +def test_cold_codec_provider_rejects_mixed_layer_spec_dtypes( + layer_dtypes: tuple[DataType, DataType], +) -> None: + codec_provider = Mock() + layer_specs = [ + KVCacheLayerSpec( + global_layer_id=layer_id, + owner=KVCacheLayerOwner.TARGET, + dtype=layer_dtype, + num_kv_heads=1, + head_dim=1, + buffer_roles=(Role.KEY,), + ) + for layer_id, layer_dtype in enumerate(layer_dtypes) + ] + with pytest.raises(NotImplementedError, match="mixed-dtype KV-cache layer specs"): + _make_manager_for_cache_tier_test( + KvCacheConfig(max_gpu_total_bytes=16 << 20), + [Mock()], + cold_page_codec_provider=codec_provider, + layer_specs=layer_specs, + ) + + codec_provider.create_cold_page_codec.assert_not_called() + + def test_extra_tokens_are_in_context_capacity() -> None: config = _make_cache_config_for_test( KvCacheConfig(avg_seq_len=264), @@ -3027,6 +3351,7 @@ def _make_fresh_fill_manager( manager = object.__new__(KVCacheManagerV2) manager._stream = SimpleNamespace(device=torch.device("cuda:0")) manager._fresh_page_fill = fill_value + manager.layer_specs = None manager._fresh_pages_filled = {} manager._fresh_fill_announced = True manager.kv_cache_map = {1: kv_cache} @@ -3347,6 +3672,7 @@ def test_warmup_zeroing_preserves_guard_page(monkeypatch: pytest.MonkeyPatch) -> manager = object.__new__(KVCacheManagerV2) manager.layer_offsets = {0: 0} + manager.layer_specs = None manager.layer_to_pool_mapping_dict = {0: 0} manager.get_buffers = lambda layer_idx, kv_layout="NHD": buffer manager._guard_page_by_layer = {0: guard_page} diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 30627c33e48e..9c8067141c5e 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2054,6 +2054,7 @@ def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): def test_v2_hybrid_rejects_quota_below_live_state_floor(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.max_batch_size = 2 @@ -2085,6 +2086,7 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.max_batch_size = 2 @@ -2313,6 +2315,7 @@ def test_expect_snapshot_points_binding_round_trip(): def test_v2_hybrid_pool_ratio_controls_allocated_memory(): def allocated_memory(pool_ratio): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.kv_cache_type = CacheTypeCpp.SELF @@ -3903,6 +3906,7 @@ def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): """A V2 hybrid manager stub for _build_cache_config: one KDA layer with the replay caches and one attention layer.""" mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.layer_specs = None mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.kv_cache_type = CacheTypeCpp.SELF diff --git a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py index 322ad40bd6f3..559500719639 100644 --- a/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py +++ b/tests/unittest/disaggregated/test_minimax_m3_kv_transfer.py @@ -70,6 +70,8 @@ def get_base_page_indices(_pool_id): def test_minimax_cache_indices_support_block_count_limit() -> None: manager = object.__new__(MiniMaxM3KVCacheManagerV2) + manager.layer_specs = None + manager.kv_factor = 2 manager.kv_cache_map = {7: _FakeKVCache(), 8: _ShortFakeKVCache()} assert manager._get_batch_cache_indices_by_pool_id([7]) == [[4, 5, 6, -1]] @@ -79,6 +81,12 @@ def test_minimax_cache_indices_support_block_count_limit() -> None: [4, 5, 6, -1], [7, 8, -1, -1], ] + assert manager._get_batch_cache_indices_by_pool_id([7], index_scale=3, kv_factor=2) == [ + [4, 5, 6, -1] + ] + assert manager.get_batch_cache_indices([7, 8]) == [[4, 5, 6, -1], [7, 8, -1, -1]] + assert manager.get_batch_cache_indices([7], num_blocks_per_seq=[2]) == [[4, 5]] + assert manager.get_batch_cache_indices([7], raw_indices=True) == [[4, 5, 6, -1]] assert manager.get_block_ids_per_seq([7]).tolist() == [[4, 5, 6, 0]] assert manager.get_block_ids_per_seq([7, 8]).tolist() == [ [4, 5, 6, 0],