diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 9ca7b7aaa49f..90d9f3b9c7d8 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -28,7 +28,7 @@ unset or when the safety sanitizer rejects the runtime value. ### `TorchLlmArgs` -269 captured fields. +270 captured fields. | Captured key | Annotation | Kind | Converter | Allowed values | |--------------|------------|------|-----------|----------------| @@ -43,6 +43,7 @@ unset or when the safety sanitizer rejects the runtime value. | `attention_dp_config.kv_cache_routing_load_balance_weight` | `` | `value` | | | | `attention_dp_config.kv_cache_routing_match_rate_threshold` | `` | `value` | | | | `attention_dp_config.kv_cache_routing_max_sessions` | `` | `value` | | | +| `attention_dp_config.kv_cache_routing_new_conv_placement` | `Literal['round_robin', 'least_queued']` | `categorical` | | `round_robin`, `least_queued` | | `attention_dp_config.timeout_iters` | `` | `value` | | | | `attn_backend` | `` | `categorical` | allowlist | `VANILLA`, `TRTLLM`, `FLASHINFER`, `FLASHINFER_STAR_ATTENTION` | | `backend` | `Literal['pytorch']` | `categorical` | | `pytorch` | @@ -123,7 +124,7 @@ unset or when the safety sanitizer rejects the runtime value. | `kv_cache_config.mamba_ssm_cache_dtype` | `Literal['auto', 'float16', 'bfloat16', 'float32']` | `categorical` | | `auto`, `float16`, `bfloat16`, `float32` | | `kv_cache_config.mamba_ssm_philox_rounds` | `` | `value` | | | | `kv_cache_config.mamba_ssm_stochastic_rounding` | `` | `value` | | | -| `kv_cache_config.mamba_state_cache_interval` | `` | `value` | | | +| `kv_cache_config.mamba_state_config.periodic_snapshot_interval` | `` | `value` | | | | `kv_cache_config.max_attention_window` | `Optional[List[int]]` | `value` | | | | `kv_cache_config.max_gpu_total_bytes` | `` | `value` | | | | `kv_cache_config.max_tokens` | `Optional[int]` | `value` | | | diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index a44a6a866904..82f7fbe58143 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -74,6 +74,49 @@ scheduler_config: enable_prefix_aware_scheduling: false ``` +### Mamba Snapshot Boundaries + +Hybrid Mamba models must retain the recurrent Mamba state together with the +attention KV prefix. Snapshot policy is grouped under +`kv_cache_config.mamba_state_config`. `periodic_snapshot_interval` controls +periodic boundaries. They are disabled by default; set the interval to a +positive value to enable them. The deprecated +`kv_cache_config.mamba_state_cache_interval` alias remains accepted for +compatibility and is copied to the nested field during validation. New code and +configuration files should use the nested field. The prototype +`additional_snapshot_offsets_from_start` and +`additional_snapshot_offsets_from_end` options add fixed boundaries. Start +offsets count tokens from the beginning of the prompt. End offsets count +backward from the prompt end, and an end offset of `0` selects the final +prompt boundary. The `per_conversation` block reuse policy disables periodic +Mamba snapshots, so configure one or more explicit stable boundaries (usually +an end offset of `0`) when using it with a hybrid Mamba model. For example: + +```yaml +kv_cache_config: + enable_block_reuse: true + use_kv_cache_manager_v2: true + avg_seq_len: 2048 + mamba_state_config: + periodic_snapshot_interval: 0 + additional_snapshot_offsets_from_start: [128] + additional_snapshot_offsets_from_end: [0, 32] +``` + +This retains snapshots after the first 128 tokens, at the end of the prompt, +and before the final 32 prompt tokens. Positions outside a particular prompt +are ignored. Set `avg_seq_len` to the workload's average total sequence length +so V2 can size the attention KV and Mamba state pools in the right proportion. +If neither `avg_seq_len` nor an explicit `pool_ratio` is configured, hybrid +Mamba models warn and fall back to half of `max_seq_len`, which can produce a +suboptimal pool split. Exact explicit boundaries currently require +`MambaHybridCacheManagerV2`, `max_beam_width=1`, and no KV connector. Hybrid +Mamba models select V2 by default when +`use_kv_cache_manager_v2: auto`; set it to `false` to select the V1 C++ +compatibility manager. In disaggregated serving, V2 Mamba requires the Python +NIXL transceiver (`transceiver_runtime: PYTHON`); V1 routes support periodic +snapshots only. + ### KV Cache Salting for Secure Reuse KV cache salting provides a security mechanism to control which requests can reuse cached KV states. When a `cache_salt` parameter is provided with a request, the KV cache system will only allow reuse of cached blocks given the same cache salt value. This prevents potential security issues such as prompt theft attacks, where malicious users might try to infer information from cached states of other users' requests. diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index a3f4e89d8d85..992eec88b7fd 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/attention_backend/trtllm.py @@ -1004,6 +1004,11 @@ def update_spec_dec_param( self.use_spec_decoding = self.is_spec_decoding_enabled self.is_spec_dec_tree = is_spec_dec_tree self.is_spec_dec_dynamic_tree = is_spec_dec_dynamic_tree + # A hybrid model's first executed attention layer can have a nonzero + # cache-local index because recurrent layers precede it. Do not rely + # on the C++ ``layer_idx == 0`` fallback to rebuild the target mask: + # the dynamic draft loop clears that mask before the next target step. + self.force_prepare_spec_dec_tree_mask = is_spec_dec_dynamic_tree # Forward static tree length to FMHA kernel selection. self.max_total_draft_tokens = max_total_draft_tokens diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py index 405810661fba..36aeb34ffdcd 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/attention/peer.py @@ -587,6 +587,11 @@ def head_match(self, peer_ri: RankInfo) -> tuple[bool, bool]: return head_match, is_dup_head def duplicate_head_factors(self, peer_ri: RankInfo) -> tuple[int, int]: + # Head duplication only applies to attention KV. In particular, do + # not divide 0 / 0 for an attention-free hybrid PP stage; MambaPolicy + # computes its TP mapping separately. + if self._ri.attention.kv_heads_per_rank == 0 or peer_ri.attention.kv_heads_per_rank == 0: + return 1, 1 factor_self, factor_peer = self._head_factors(peer_ri) dup_head = max(1, factor_self // factor_peer) peer_dup_head = max(1, factor_peer // factor_self) diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index e8c8816ba4c8..a96fc3e8f496 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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 typing import Dict, List, Optional, Tuple import numpy as np @@ -338,13 +352,17 @@ def _build_layer_ptrs( overlapping_layers: List[int], slot: int, ) -> np.ndarray: - """Build per-layer pointers for a given pool (conv or ssm) and slot.""" - ptrs = [] - for glid in overlapping_layers: - lid = layer_offsets[glid] - ptrs.append( - pool.base_address + lid * pool.num_slots * pool.slot_bytes + slot * pool.slot_bytes - ) + """Build per-layer pointers from a pool's affine layer/slot layout.""" + slot_stride_bytes = pool.slot_stride_bytes + layer_stride_bytes = pool.layer_stride_bytes + assert slot_stride_bytes is not None + assert layer_stride_bytes is not None + ptrs = [ + pool.base_address + + layer_offsets[global_layer_id] * layer_stride_bytes + + slot * slot_stride_bytes + for global_layer_id in overlapping_layers + ] return np.array(ptrs, dtype=np.int64) @staticmethod @@ -431,10 +449,16 @@ def build_mamba_frags( (self_mlg.ssm_states, peer_mlg.ssm_states, False), ]: src_ptrs = MambaPolicy._build_layer_ptrs( - self_pool, self_mlg.mamba_layer_offsets, overlapping_layers, src_slot + self_pool, + self_mlg.mamba_layer_offsets, + overlapping_layers, + src_slot, ) dst_ptrs = MambaPolicy._build_layer_ptrs( - peer_pool, peer_mlg.mamba_layer_offsets, overlapping_layers, dst_slot + peer_pool, + peer_mlg.mamba_layer_offsets, + overlapping_layers, + dst_slot, ) src_region = SpecRegion( diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 4c0dba242199..e2803042c7fc 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -74,6 +74,10 @@ def from_kv_cache_manager( m = kv_cache_manager.mapping kvm = kv_cache_manager enable_attention_dp = m.enable_attention_dp + # Keep AttentionInfo on attention-free PP stages so it can still carry + # the attention-DP topology used by Mamba transfers. A zero head count + # means that this rank has no local attention cache; AttentionPolicy + # must not perform head-ratio arithmetic for such ranks. kv_heads_per_rank = next((h for h in kvm.num_kv_heads_per_layer if h > 0), 0) # Eight is the smallest element count guaranteed to occupy whole bytes # for every supported sub-byte cache dtype (including NVFP4). diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 4aa76c5b9605..9248bf506bce 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -14,9 +14,10 @@ # limitations under the License. from collections import defaultdict -from typing import Dict, List +from typing import Dict, List, Sequence import numpy as np +import torch from tensorrt_llm._torch.disaggregation.base.region import ( DataLayout, @@ -41,7 +42,10 @@ get_physical_pool, ) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManager, + MambaHybridCacheManagerV2, +) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import get_size_in_bytes, nvtx_range from tensorrt_llm.bindings import DataType @@ -106,10 +110,12 @@ def extract( base_ptr = pool.base_address block_size = pool.slot_bytes + block_stride = pool.slot_stride_bytes + assert block_stride is not None # KV cache: filter out invalid block_ids (BAD_PAGE_INDEX = -1) valid = region_ids >= 0 - ptrs = base_ptr + block_size * region_ids[valid] + ptrs = base_ptr + block_stride * region_ids[valid] memory = MemRegionGroup(ptrs=ptrs, bytes_per_region=block_size) return SpecRegion(memory=memory) @@ -134,16 +140,19 @@ def _build_layer_group_for_mamba( base_address=conv_state.data_ptr(), slot_bytes=conv_state.stride(1) * conv_state.element_size(), num_slots=conv_state.shape[1], + layer_stride_bytes=conv_state.stride(0) * conv_state.element_size(), ) ssm_pool = PhysicalPool( base_address=ssm_state.data_ptr(), slot_bytes=ssm_state.stride(1) * ssm_state.element_size(), num_slots=ssm_state.shape[1], + layer_stride_bytes=ssm_state.stride(0) * ssm_state.element_size(), ) # Per-section bytes for conv_state and per-head bytes for ssm_state. - # conv_state layout: [x: d_inner/tp | B: ng*ds/tp | C: ng*ds/tp] x (d_conv-1) + # The section ordering is supplied by the cache manager because Mamba2 + # uses [x | B | C], while GDN uses [Q | K | V]. # ssm_state layout: (nheads/tp, head_dim, d_state) d_conv_m1 = conv_state.shape[3] conv_elem_size = conv_state.element_size() @@ -165,6 +174,91 @@ def _build_layer_group_for_mamba( ) +def _slot_stride_bytes(tensor: torch.Tensor) -> int: + return int(tensor.stride(0) * tensor.element_size()) + + +def _build_v2_mamba_state_pool(states: Sequence[torch.Tensor]) -> PhysicalPool: + """Describe affine layer/slot addressing for one V2 Mamba state role.""" + if not states: + raise ValueError("V2 Mamba state pool requires at least one layer") + + first_state = states[0] + base_address = int(first_state.data_ptr()) + num_slots = int(first_state.shape[0]) + slot_bytes = int(first_state[0].numel() * first_state.element_size()) + slot_stride_bytes = _slot_stride_bytes(first_state) + + num_layers = len(states) + if slot_stride_bytes % num_layers != 0: + raise ValueError("V2 Mamba physical slot must divide evenly across layers") + # Each role appears once per layer in its size-class pool. Equal-size SSM + # and convolution states share that pool and are interleaved, so their + # layer stride includes both role payloads. + layer_stride_bytes = slot_stride_bytes // num_layers + + for layer_offset, state in enumerate(states): + state_slot_bytes = int(state[0].numel() * state.element_size()) + if ( + int(state.shape[0]) != num_slots + or state_slot_bytes != slot_bytes + or _slot_stride_bytes(state) != slot_stride_bytes + ): + raise ValueError("V2 Mamba state tensors must share one slot layout per role") + expected_address = base_address + layer_offset * layer_stride_bytes + if int(state.data_ptr()) != expected_address: + raise ValueError("V2 Mamba state tensors must have a uniform layer stride per role") + + return PhysicalPool( + base_address=base_address, + slot_bytes=slot_bytes, + num_slots=num_slots, + slot_stride_bytes=slot_stride_bytes, + layer_stride_bytes=layer_stride_bytes, + ) + + +def _build_layer_group_for_v2_mamba( + manager: MambaHybridCacheManagerV2, pool_group_idx: int +) -> MambaLayerGroup: + mamba_layer_offsets = { + int(global_layer_id): int(local_layer_id) + for global_layer_id, local_layer_id in manager.mamba_layer_offsets.items() + } + + expected_offsets = list(range(len(mamba_layer_offsets))) + if sorted(mamba_layer_offsets.values()) != expected_offsets: + raise ValueError("V2 Mamba layer offsets must be dense") + if len(manager.all_conv_states) != len(expected_offsets) or len(manager.all_ssm_states) != len( + expected_offsets + ): + raise ValueError("V2 Mamba state tensors must match the layer-offset table") + + first_conv_state = manager.all_conv_states[0] + first_ssm_state = manager.all_ssm_states[0] + conv_pool = _build_v2_mamba_state_pool(manager.all_conv_states) + ssm_pool = _build_v2_mamba_state_pool(manager.all_ssm_states) + if conv_pool.num_slots != ssm_pool.num_slots: + raise ValueError("V2 Mamba convolution and SSM states must have the same number of slots") + + d_conv_m1 = manager.conv_state_shape[1] + conv_elem_size = first_conv_state.element_size() + _, head_dim, d_state = manager.ssm_state_shape + conv_section_bytes = [dim * d_conv_m1 * conv_elem_size for dim in manager.conv_section_dims] + + ssm_elem_size = first_ssm_state.element_size() + ssm_bytes_per_head = head_dim * d_state * ssm_elem_size + + return MambaLayerGroup( + pool_group_idx=pool_group_idx, + mamba_layer_offsets=mamba_layer_offsets, + conv_states=conv_pool, + ssm_states=ssm_pool, + conv_section_bytes=conv_section_bytes, + ssm_bytes_per_head=ssm_bytes_per_head, + ) + + def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManager (V1).""" if kv_cache_manager.dtype == DataType.NVFP4: @@ -419,6 +513,14 @@ def _window_size_for_layer(internal_layer_id: int): 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]) + if isinstance(manager, MambaHybridCacheManagerV2) and any( + manager._is_local_mamba_layer(int(layer_id)) for layer_id in all_internal_layer_ids + ): + layer_groups_by_id[layer_group_id] = _build_layer_group_for_v2_mamba( + manager, storage_pg_to_list_idx[storage_pg_idx] + ) + continue + all_global_layer_ids = _compute_global_layer_ids(manager, layer_group_id) local_layers = [ @@ -519,7 +621,9 @@ def _window_size_for_layer(internal_layer_id: int): raise ValueError(f"Missing V2 layer group descriptor for layer group {layer_group_id}") layer_groups.append(layer_group) - if isinstance(manager, MambaHybridCacheManager): + if isinstance(manager, MambaHybridCacheManager) and not isinstance( + manager, MambaHybridCacheManagerV2 + ): mamba_layer_group_idx = len(pool_groups) mamba_layer_group = _build_layer_group_for_mamba(manager, mamba_layer_group_idx) layer_groups.append(mamba_layer_group) diff --git a/tensorrt_llm/_torch/disaggregation/resource/page.py b/tensorrt_llm/_torch/disaggregation/resource/page.py index 259561cfd354..b15f21764e61 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -76,15 +76,41 @@ class MapperKind(IntEnum): @dataclass class PhysicalPool: + """Affine view of a physical pool over logical layers and slots. + + ``slot_bytes`` is the transferable payload for one ``(layer, slot)`` and + ``num_slots`` is the number of logical slots. The payload address is + ``base_address + layer * layer_stride_bytes + slot * slot_stride_bytes``. + The strides describe the physical layout independently of payload size and + slot count. Their defaults describe dense layer-major storage, where + ``slot_stride_bytes == slot_bytes`` and + ``layer_stride_bytes == num_slots * slot_stride_bytes``. V2 Mamba supplies + both explicitly for its slot-major, role-interleaved pools. + """ + base_address: int # uint64 slot_bytes: int num_slots: int + slot_stride_bytes: Optional[int] = None + layer_stride_bytes: Optional[int] = None + + def __post_init__(self) -> None: + if self.slot_stride_bytes is None: + self.slot_stride_bytes = self.slot_bytes + if self.layer_stride_bytes is None: + self.layer_stride_bytes = self.num_slots * self.slot_stride_bytes + if self.slot_stride_bytes < self.slot_bytes: + raise ValueError("slot_stride_bytes must be greater than or equal to slot_bytes") + if self.layer_stride_bytes < self.slot_bytes: + raise ValueError("layer_stride_bytes must be greater than or equal to slot_bytes") def to_dict(self) -> dict: return { "base_address": int(self.base_address), "slot_bytes": int(self.slot_bytes), "num_slots": int(self.num_slots), + "slot_stride_bytes": int(self.slot_stride_bytes), + "layer_stride_bytes": int(self.layer_stride_bytes), } @staticmethod @@ -93,6 +119,16 @@ def from_dict(data: dict) -> "PhysicalPool": base_address=int(data["base_address"]), slot_bytes=int(data["slot_bytes"]), num_slots=int(data["num_slots"]), + slot_stride_bytes=( + int(data["slot_stride_bytes"]) + if data.get("slot_stride_bytes") is not None + else None + ), + layer_stride_bytes=( + int(data["layer_stride_bytes"]) + if data.get("layer_stride_bytes") is not None + else None + ), ) diff --git a/tensorrt_llm/_torch/disaggregation/resource/utils.py b/tensorrt_llm/_torch/disaggregation/resource/utils.py index 6f58bc2782e4..2e39ed9af17e 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.py @@ -25,7 +25,7 @@ def get_pool_bytes(pool: PhysicalPool) -> int: - """Total bytes across all slots in this pool.""" + """Total transferable payload bytes across all slots in this pool.""" return pool.slot_bytes * pool.num_slots @@ -33,7 +33,8 @@ def get_slot_address(pool: PhysicalPool, slot_id: int) -> int: """Base address of *slot_id*.""" if slot_id >= pool.num_slots: raise ValueError(f"slot_id {slot_id} >= num_slots {pool.num_slots}") - return pool.base_address + slot_id * pool.slot_bytes + assert pool.slot_stride_bytes is not None + return pool.base_address + slot_id * pool.slot_stride_bytes # ------------------------------------------------------------------------- @@ -190,9 +191,22 @@ def get_unique_pool_memory_descs( pool_counter = 0 for lg_idx, lg in enumerate(page_table.layer_groups): if isinstance(lg, MambaLayerGroup): - num_mamba_layers = len(lg.mamba_layer_offsets) - for pool in [lg.conv_states, lg.ssm_states]: - pool_size = num_mamba_layers * pool.num_slots * pool.slot_bytes + # V2 Mamba layer groups reference manager-owned physical pools. + # V1 Mamba state views are standalone and use the first invalid + # pool-group index after the attention groups. + has_physical_pool_group = 0 <= int(lg.pool_group_idx) < len(page_table.pool_groups) + if has_physical_pool_group: + pools_and_sizes = [ + (pool, get_pool_bytes(pool)) + for pool in page_table.pool_groups[int(lg.pool_group_idx)].pools + ] + else: + num_mamba_layers = len(lg.mamba_layer_offsets) + pools_and_sizes = [ + (pool, num_mamba_layers * pool.num_slots * pool.slot_bytes) + for pool in [lg.conv_states, lg.ssm_states] + ] + for pool, pool_size in pools_and_sizes: pool_key = (pool.base_address, pool_size) if pool_key not in unique_pools: unique_pools[pool_key] = pool_counter diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 2c257a39f731..133dd75270dc 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -33,7 +33,10 @@ from tensorrt_llm._torch.distributed.communicator import Distributed from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import KvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManager, + MambaHybridCacheManagerV2, +) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import nvtx_range from tensorrt_llm.bindings import LlmRequestState @@ -140,7 +143,9 @@ def _init_sync_policy(self): def _exchange_rank_info(self): endpoints = cast(list, self._dist.allgather(self._transfer_worker.sender_endpoint)) layer_num = len(self._kv_cache_manager.pp_layers) - if isinstance(self._kv_cache_manager, MambaHybridCacheManager): + if isinstance(self._kv_cache_manager, MambaHybridCacheManager) and not isinstance( + self._kv_cache_manager, MambaHybridCacheManagerV2 + ): layer_num += len(self._kv_cache_manager._impl.mamba_layer_offsets) layer_num_per_pp = cast(list, getattr(self._dist, "pp_allgather")(layer_num)) self._transfer_worker.populate_instance_and_rank_info( @@ -233,7 +238,12 @@ def _create_kv_slice(self, req: LlmRequest) -> KVSlice: groups.append(block_ids) mamba_state_index = None - if isinstance(self._kv_cache_manager, MambaHybridCacheManager): + if isinstance(self._kv_cache_manager, MambaHybridCacheManagerV2): + if self._kv_cache_manager.local_num_mamba_layers > 0: + mamba_state_index = self._kv_cache_manager._request_id_to_state_index[ + req.py_request_id + ] + elif isinstance(self._kv_cache_manager, MambaHybridCacheManager): mamba_state_index = self._kv_cache_manager.mamba_cache_index[req.py_request_id] return KVSlice( diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 3bb491fd1b91..8516a3dbdb1e 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -17,7 +17,7 @@ import re from contextlib import contextmanager from dataclasses import replace -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import torch @@ -985,11 +985,22 @@ def load_weights(self, def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: """Model-specific defaults for NemotronH. - Disables block reuse due to SSM/hybrid architecture constraints. + Uses KV cache manager V2 for the hybrid state layout. Block reuse + remains opt-in because it also requires a Mamba snapshot policy. """ - # TODO: Remove enable_block_reuse=False once KV cache block reuse - # is supported for Mamba/SSM-based models - return {"kv_cache_config": {"enable_block_reuse": False}} + return { + "kv_cache_config": { + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + } + } + + @classmethod + def get_preferred_transceiver_runtime(cls, + pretrained_config: object + | None = None) -> Literal["PYTHON"]: + """Use the Python transceiver for hybrid-state transfers.""" + return "PYTHON" @staticmethod def lora_config(model_dir: str): diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 6fc43129a353..3a5c0612226c 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_5.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_5.py @@ -15,7 +15,7 @@ import re from types import SimpleNamespace -from typing import Dict, List +from typing import Dict, List, Literal import torch from transformers import PretrainedConfig @@ -669,11 +669,17 @@ def get_model_defaults(cls, llm_args): # model class (this VLM wrapper), not on the inner decoder. Both # inner LMs (`Qwen3_5MoeForCausalLM` / `Qwen3_5ForCausalLM`) inherit # `Qwen3NextForCausalLM`'s defaults unchanged, so delegate to it to - # propagate `enable_block_reuse=False` — the hybrid Mamba/SSM path - # doesn't support KV-cache block reuse. Without this the VLM path - # would silently fall back to the global default (block reuse on). + # propagate the V2 manager selection and keep block reuse disabled + # until a recurrent-state snapshot policy is configured. return Qwen3NextForCausalLM.get_model_defaults(llm_args) + @classmethod + def get_preferred_transceiver_runtime( + cls, pretrained_config: object | None = None + ) -> Literal["PYTHON"]: + """Match the hybrid text decoder's Python disaggregated route.""" + return "PYTHON" + def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs): kwargs["vision_model_class"] = Qwen3VisionModel kwargs["disable_fuse_rope"] = kwargs.get("disable_fuse_rope", False) diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_next.py b/tensorrt_llm/_torch/models/modeling_qwen3_next.py index 6a8a2a9495c3..ce33e4c6274d 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -20,7 +20,7 @@ import copy import os from types import SimpleNamespace -from typing import TYPE_CHECKING, Dict, List, Optional +from typing import TYPE_CHECKING, Dict, List, Literal, Optional import torch @@ -1051,9 +1051,24 @@ def __init__( @classmethod def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict: - # TODO: Remove enable_block_reuse=False once KV cache block reuse - # is supported for Mamba/SSM-based models - return {"kv_cache_config": {"enable_block_reuse": False}} + """Use V2 for the hybrid state layout. + + Block reuse remains opt-in because it also requires a recurrent-state + snapshot policy. + """ + return { + "kv_cache_config": { + "enable_block_reuse": False, + "use_kv_cache_manager_v2": True, + } + } + + @classmethod + def get_preferred_transceiver_runtime(cls, + pretrained_config: object + | None = None) -> Literal["PYTHON"]: + """Use the Python transceiver for hybrid-state transfers.""" + return "PYTHON" def load_weights(self, weights: dict, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..4f00fde762c0 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -47,9 +47,9 @@ from ..speculative import (get_num_extra_kv_tokens, get_num_spec_layers, get_spec_decoder, should_use_separate_draft_kv_cache) from ..utils import is_gdn_replay_enabled -from .config_utils import (extract_mamba_kv_cache_params, is_gemma4_hybrid, - is_hybrid_linear, is_mla, is_nemotron_hybrid, - is_qwen3_hybrid) +from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params, + is_gemma4_hybrid, is_hybrid_linear, is_mla, + is_nemotron_hybrid, is_qwen3_hybrid) from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .guided_decoder import GuidedDecoder @@ -58,6 +58,7 @@ from .llm_request import ExecutorResponse, LlmRequestState from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, MixedMambaHybridCacheManager, use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine @@ -87,6 +88,22 @@ def _non_hybrid_kv_cache_manager_cls(config, kv_cache_config: KvCacheConfig): return KVCacheManagerV2 if needs_v2 else KVCacheManager +def _resolve_disagg_transceiver_route( + cache_transceiver_config: Optional[CacheTransceiverConfig], +) -> tuple[Optional[str], Optional[str]]: + """Return the effective backend and runtime used for manager routing.""" + if cache_transceiver_config is None: + return None, None + + backend, _ = cache_transceiver_config._resolve_default_backend() + runtime = cache_transceiver_config.transceiver_runtime + if runtime == "auto": + # Model loading normally resolves ``auto``. Paths that skip model + # defaults use the global C++ fallback, matching transceiver creation. + runtime = None + return backend, runtime + + def get_kv_cache_manager_cls( model_config: ModelConfig, kv_cache_config: KvCacheConfig, @@ -94,14 +111,20 @@ def get_kv_cache_manager_cls( cache_transceiver_config: Optional[CacheTransceiverConfig] = None): """Resolve the concrete KV cache manager class for ``model_config``. - For hybrid mamba models the choice between ``Mixed`` (TRTLLM_USE_PY_MAMBA) - and ``Cpp`` (unified pool with block reuse) is made here. Callers that - don't care about disagg can omit ``is_disagg`` and get the unified-pool - default. + For hybrid mamba models the choice between + ``MambaHybridCacheManagerV2`` and compatibility managers is made here. + Callers that don't care about disagg can omit ``is_disagg`` and get the + unified-pool default. - Env-var overrides (agg mode only — disagg picks its inner impl via - ``cache_transceiver_config.transceiver_runtime``): - * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager with PythonMambaCacheManager. + Model loading resolves ``use_kv_cache_manager_v2="auto"`` to V2 for + supported hybrid Mamba models. An explicit ``False`` selects a + compatibility manager. In disaggregated serving, V2 additionally requires + the Python transceiver with the NIXL backend. Unsupported V2 routes fail + rather than falling back to a different manager. + + Env-var overrides: + * ``TRTLLM_USE_PY_MAMBA=1`` — Mixed manager in aggregated serving. + * ``TLLM_MAMBA_MANAGER_PREFERENCE`` — explicit manager preference. """ config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config @@ -122,44 +145,100 @@ def get_kv_cache_manager_cls( f"Sparse attention algorithm {sparse_attn_algorithm!r} is not " "supported with hybrid Mamba / linear-attention models.") + state_config = kv_cache_config.mamba_state_config + has_additional_snapshots = bool( + state_config.additional_snapshot_offsets_from_start + or state_config.additional_snapshot_offsets_from_end) + use_v2 = kv_cache_config.use_kv_cache_manager_v2 is True + + if has_additional_snapshots and not use_v2: + raise ValueError("Mamba additional snapshot offsets require " + "use_kv_cache_manager_v2=True; V1 supports only " + "periodic_snapshot_interval.") # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. - if use_py_mamba_cache_manager(): + if is_disagg: + backend, runtime = _resolve_disagg_transceiver_route( + cache_transceiver_config) + if use_v2: + if runtime != "PYTHON" or backend != "NIXL": + raise ValueError( + "KV cache manager V2 for hybrid Mamba disaggregated " + "serving requires transceiver_runtime='PYTHON' with " + "backend='NIXL'.") + else: + if (kv_cache_config.enable_block_reuse and runtime == "PYTHON"): + raise ValueError( + "Hybrid Mamba disaggregated serving with block reuse " + "and transceiver_runtime='PYTHON' requires " + "use_kv_cache_manager_v2=True.") + if kv_cache_config.enable_block_reuse: + return CppMambaHybridCacheManager + if runtime == "PYTHON" and backend == "NIXL": + logger.info("Python transceiver detected; using " + "MixedMambaHybridCacheManager for hybrid model") + return MixedMambaHybridCacheManager + return CppMambaHybridCacheManager + + if use_py_mamba_cache_manager() and not is_disagg: + if use_v2: + raise ValueError( + "TRTLLM_USE_PY_MAMBA=1 conflicts with explicit " + "use_kv_cache_manager_v2=True.") if kv_cache_config.enable_block_reuse: raise ValueError( "TRTLLM_USE_PY_MAMBA=1 forces " "MixedMambaHybridCacheManager, which does not support " "block reuse. Disable block reuse or unset " - "TRTLLM_USE_PY_MAMBA to use CppMambaHybridCacheManager.") + "TRTLLM_USE_PY_MAMBA to use the configured cache manager.") logger.info( "Using MixedMambaHybridCacheManager for hybrid mamba model") return MixedMambaHybridCacheManager - if kv_cache_config.enable_block_reuse: - return CppMambaHybridCacheManager - if (cache_transceiver_config is not None - and cache_transceiver_config.transceiver_runtime == "PYTHON"): - logger.info("Python transceiver detected; using " - "MixedMambaHybridCacheManager for hybrid mamba model") - return MixedMambaHybridCacheManager - default_cls = CppMambaHybridCacheManager env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None) if env_override is not None: - if env_override.upper() == 'MIXED': + env_override = env_override.upper() + if env_override == 'MIXED': + if use_v2: + raise ValueError( + "TLLM_MAMBA_MANAGER_PREFERENCE=MIXED conflicts with " + "explicit use_kv_cache_manager_v2=True.") + if kv_cache_config.enable_block_reuse: + raise ValueError( + "TLLM_MAMBA_MANAGER_PREFERENCE=MIXED forces " + "MixedMambaHybridCacheManager, which does not support " + "block reuse. Disable block reuse, use the CPP " + "preference, or explicitly enable KV cache manager " + "V2.") logger.warning( - "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED overrides the default Mamba cache manager to MixedMambaHybridCacheManager. This may lead to increased memory usage due to lack of block reuse, but can be necessary for disaggregated setups or to avoid potential issues with the C++ manager. Set TLLM_MAMBA_MANAGER_PREFERENCE=CPP to use the CppMambaHybridCacheManager instead, which is the default for non-disaggregated setups without block reuse explicitly disabled." - ) + "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED " + "overrides the default Mamba cache manager to " + "MixedMambaHybridCacheManager.") return MixedMambaHybridCacheManager - elif env_override.upper() == 'CPP': + if env_override == 'CPP': + if use_v2: + raise ValueError( + "TLLM_MAMBA_MANAGER_PREFERENCE=CPP conflicts with " + "explicit use_kv_cache_manager_v2=True.") logger.warning( - "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP overrides the default Mamba cache manager to CppMambaHybridCacheManager. This enables block reuse and can reduce memory usage, but may not be compatible with disaggregated setups. Set TLLM_MAMBA_MANAGER_PREFERENCE=MIXED to use the MixedMambaHybridCacheManager instead if you encounter issues with the C++ manager or are running in a disaggregated environment." - ) + "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP " + "overrides the default Mamba cache manager to " + "CppMambaHybridCacheManager.") return CppMambaHybridCacheManager - else: - logger.warning( - f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " - f"Expected 'CPP' or 'MIXED'. Using default {default_cls.__name__}." - ) - return default_cls + logger.warning( + f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " + "Expected 'CPP' or 'MIXED'. Using the configured " + "KV cache manager default.") + + if not use_v2: + return CppMambaHybridCacheManager + + if (kv_cache_config.enable_block_reuse + and kv_cache_config.enable_kv_pool_rebalance): + raise ValueError( + "V2 Mamba block reuse is not compatible with " + "enable_kv_pool_rebalance because the rebalancer does not " + "yet model retained recurrent-state snapshots.") + return MambaHybridCacheManagerV2 elif sparse_attn_config is not None: return get_sparse_attn_kv_cache_manager(sparse_attn_config) else: @@ -358,22 +437,16 @@ def _get_model_kv_cache_manager_cls( cache_transceiver_config=self._cache_transceiver_config) cls = self._fallback_if_unsupported_kv_cache_manager_v2( cls, model_config, kv_cache_config) - # The V1-route hybrid mamba managers (disagg, TRTLLM_USE_CPP_MAMBA, - # TRTLLM_USE_PY_MAMBA, or one-model speculative decoding) keep mamba - # state in a separate cache that doesn't honor block reuse. Warn at - # the routing site so users see the warning where the decision is - # actually made. + # Compatibility managers do not support MTP block reuse. Warn at the + # routing site so users see the concrete manager selected for the + # incompatible combination. if is_hybrid_linear(model_engine.model.model_config.pretrained_config) \ - and kv_cache_config.enable_block_reuse: - uses_v1_mamba_route = self._is_disagg \ - or os.environ.get('TRTLLM_USE_CPP_MAMBA', '0') == '1' \ - or os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' \ - or self._speculative_config is not None - if uses_v1_mamba_route: + and kv_cache_config.enable_block_reuse \ + and self._speculative_config is not None: + if not issubclass(cls, MambaHybridCacheManagerV2): logger.warning( "Block reuse does not work with MTP for hybrid linear models " - "when using the legacy MambaCacheManager (TRTLLM_USE_CPP_MAMBA=1)" - ) + f"when using non-V2 Mamba cache manager {cls.__name__}") return cls def _fallback_if_unsupported_kv_cache_manager_v2( @@ -390,7 +463,7 @@ def _fallback_if_unsupported_kv_cache_manager_v2( if self._kv_connector_manager is not None: incompat.append("kv_connector_manager") if self._max_beam_width is not None and self._max_beam_width > 1: - incompat.append("beam_width > 1") + incompat.append("max_beam_width > 1") if incompat: incompat_str = ", ".join(incompat) # Some models are structurally bound to V2 and cannot fall @@ -415,6 +488,12 @@ def _fallback_if_unsupported_kv_cache_manager_v2( f"Gemma4 hybrid attention requires KVCacheManagerV2, " f"which is not yet supported with {incompat_str}. " f"Disable these features to run Gemma4 hybrid models.") + if is_hybrid_linear(config): + raise NotImplementedError( + "Hybrid Mamba cache managers do not support " + f"{incompat_str}; CppMambaHybridCacheManager does not " + "provide a compatible fallback. Use max_beam_width=1 " + "and disable the KV connector.") # Plain V2 (explicitly enabled or selected by a model default): # V2 was a preference, not a structural requirement, so we can # safely fall back to V1. @@ -446,6 +525,15 @@ def _per_manager_cache_cost(self, spec_config=self._speculative_config, **extra_kwargs)) + def _get_one_model_draft_layer_mask(self) -> List[bool]: + """Return the same draft-only mask used by runtime construction.""" + num_draft_layers = self._get_num_draft_layers() + if self._speculative_config.spec_dec_mode.is_external_drafter(): + return [True] * num_draft_layers + target_num_layers = (self._model_engine.model.model_config. + pretrained_config.num_hidden_layers) + return [False] * target_num_layers + [True] * num_draft_layers + def _get_kv_size_per_token(self, kv_cache_config: Optional[KvCacheConfig] = None ) -> CacheCost: @@ -457,8 +545,13 @@ def _get_kv_size_per_token(self, kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) model_config = self._model_engine.model.model_config - total = self._per_manager_cache_cost(self._kv_cache_manager_cls, - model_config, kv_cache_config) + use_separate_draft_kv_cache = ( + self._should_create_separate_draft_kv_cache()) + total = self._per_manager_cache_cost( + self._kv_cache_manager_cls, + model_config, + kv_cache_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache) if self._is_encoder_decoder(): total += CacheCost.from_raw(self._get_cross_kv_size_per_token()) if self._draft_model_engine is not None: @@ -468,7 +561,7 @@ def _get_kv_size_per_token(self, total += self._per_manager_cache_cost(draft_kv_cache_manager_cls, draft_model_config, kv_cache_config) - elif self._should_create_separate_draft_kv_cache(): + elif use_separate_draft_kv_cache: # One-model draft with separate KV cache layout. # Pass num_layers explicitly since the HF config may report a # different layer count than what is actually used at runtime @@ -493,7 +586,8 @@ def _get_kv_size_per_token(self, self._kv_cache_manager_cls, effective_draft_config, kv_cache_config, - num_layers=self._get_num_draft_layers()) + num_layers=self._get_num_draft_layers(), + is_draft=True) return total def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction, @@ -1091,20 +1185,8 @@ def _create_one_model_draft_kv_cache_manager( Create a KV cache manager for draft model layers in one-model mode when target and draft have different KV cache layouts. """ - # Get target model's num_hidden_layers to compute correct layer indices. - # Draft model layers in one-model mode start at target_num_layers. - target_pretrained_config = self._model_engine.model.model_config.pretrained_config - target_num_layers = target_pretrained_config.num_hidden_layers - - # PARD, External Drafter: draft is a separate model, layers start from 0. - # Other methods (EAGLE3, MTP): draft layers are appended after target layers. num_draft_layers = self._get_num_draft_layers() - if self._speculative_config.spec_dec_mode.is_external_drafter(): - spec_dec_layer_mask = [True] * num_draft_layers - else: - spec_dec_layer_mask = [False] * target_num_layers + [ - True - ] * num_draft_layers + spec_dec_layer_mask = self._get_one_model_draft_layer_mask() # Get the effective draft config (explicit draft_config if available, # otherwise fall back to target model config for MTP). @@ -1191,9 +1273,13 @@ def _get_target_and_draft_cache_costs( target_kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) total_kv = self._get_kv_size_per_token(target_kv_cache_config) + use_separate_draft_kv_cache = ( + self._should_create_separate_draft_kv_cache()) target_kv = self._per_manager_cache_cost( - self._kv_cache_manager_cls, self._model_engine.model.model_config, - target_kv_cache_config) + self._kv_cache_manager_cls, + self._model_engine.model.model_config, + target_kv_cache_config, + use_separate_draft_kv_cache=use_separate_draft_kv_cache) # The draft contribution is whatever the aggregate has on top of the # target. Both pieces are CacheCost; subtraction is component-wise. draft_kv = CacheCost(slope=total_kv.slope - target_kv.slope, @@ -1696,6 +1782,21 @@ def _build_per_layer_num_kv_heads( ] * num_spec_layers +def _get_mamba_cache_layer_masks( + mamba_params: MambaKVCacheParams, + mapping: Mapping, + spec_config: Optional[SpeculativeConfig], + is_draft: bool, +) -> tuple[List[bool], List[bool]]: + use_separate_draft_kv_cache = ( + not mapping.enable_attention_dp + and should_use_separate_draft_kv_cache(spec_config)) + return mamba_params.get_layer_masks( + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + + def _create_kv_cache_manager( model_engine: Optional[PyTorchModelEngine], kv_cache_manager_cls, @@ -1839,6 +1940,8 @@ def _create_kv_cache_manager( manager_extra_kwargs = {} if issubclass(kv_cache_manager_cls, KVCacheManagerV2): manager_extra_kwargs["enable_stats"] = enable_kv_cache_stats + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + manager_extra_kwargs["is_disagg"] = is_disagg if is_mla(config): kv_cache_manager = kv_cache_manager_cls( @@ -1879,10 +1982,18 @@ def _create_kv_cache_manager( mamba_params = extract_mamba_kv_cache_params( config, - layer_mask=layer_mask, spec_config=spec_config, quant_config=quant_config, ) + mamba_layer_mask, full_attention_layer_mask = ( + _get_mamba_cache_layer_masks( + mamba_params, + mapping, + spec_config, + is_draft, + )) + num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 + else mamba_params.num_mamba_layers) # Replay state update kernel for MTP: default on for sm >= 80; gates # below disable it for incompatible feature combinations. Cpp cache @@ -1941,6 +2052,11 @@ def _create_kv_cache_manager( mamba_ssm_stochastic_rounding = (stochastic_rounding and mamba_params.mamba_ssm_cache_dtype == torch.float16) + mamba_manager_extra_kwargs = dict(manager_extra_kwargs) + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + mamba_manager_extra_kwargs["conv_state_layout"] = "x_b_c" + else: + mamba_manager_extra_kwargs["model_type"] = "nemotron_hybrid" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, @@ -1948,15 +2064,15 @@ def _create_kv_cache_manager( mamba_params.num_heads, mamba_params.n_groups, mamba_params.head_dim, - mamba_params.num_mamba_layers, - mamba_params.mamba_layer_mask, + num_mamba_layers, + mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=mamba_params.num_full_attention_layers, - layer_mask=mamba_params.full_attention_layer_mask, + num_layers=sum(full_attention_layer_mask), + layer_mask=full_attention_layer_mask, num_kv_heads=per_layer_num_kv_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, @@ -1968,10 +2084,9 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, - model_type="nemotron_hybrid", use_replay_state_update=use_replay, mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - **manager_extra_kwargs, + **mamba_manager_extra_kwargs, ) elif is_qwen3_hybrid(config): if max_beam_width > 1: @@ -1984,11 +2099,18 @@ def _create_kv_cache_manager( ) mamba_params = extract_mamba_kv_cache_params( config, - layer_mask=layer_mask, spec_config=spec_config, quant_config=quant_config, ) - + mamba_layer_mask, full_attention_layer_mask = ( + _get_mamba_cache_layer_masks( + mamba_params, + mapping, + spec_config, + is_draft, + )) + num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 + else mamba_params.num_mamba_layers) # Replay state update for GDN MTP: mirrors the nemotron_hybrid gating # above, minus the Mamba2-specific stochastic-rounding/Philox gate. # The GDN replay kernel does a plain cast on checkpoint commit, so @@ -2026,9 +2148,24 @@ def _create_kv_cache_manager( logger.info("GDN replay kernel is disabled; set " "TRTLLM_USE_GDN_REPLAY=1 to enable it") use_replay = False + + # Upstream GDN replay commits all local layer checkpoints through the + # contiguous C++ manager state view. V2 exposes per-layer state views, + # so enabling the same path there would fail for partitioned batches. + if (use_replay and issubclass(kv_cache_manager_cls, + MambaHybridCacheManagerV2)): + logger.info( + "GDN replay is not supported by MambaHybridCacheManagerV2; " + "using the legacy MTP path") + use_replay = False logger.info("GDN replay state update: " + ("ENABLED" if use_replay else "DISABLED")) + mamba_manager_extra_kwargs = dict(manager_extra_kwargs) + if issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + mamba_manager_extra_kwargs["conv_state_layout"] = "q_k_v" + else: + mamba_manager_extra_kwargs["model_type"] = "qwen3_next" kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, @@ -2036,15 +2173,15 @@ def _create_kv_cache_manager( mamba_params.num_heads, mamba_params.n_groups, mamba_params.head_dim, - mamba_params.num_mamba_layers, - mamba_params.mamba_layer_mask, + num_mamba_layers, + mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=mamba_params.num_full_attention_layers, - layer_mask=mamba_params.full_attention_layer_mask, + num_layers=sum(full_attention_layer_mask), + layer_mask=full_attention_layer_mask, num_kv_heads=per_layer_num_kv_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, @@ -2056,9 +2193,8 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, - model_type="qwen3_next", use_replay_state_update=use_replay, - **manager_extra_kwargs, + **mamba_manager_extra_kwargs, ) else: # NOTE: this is a workaround for VSWA to switch to calculate_max_num_blocks_for_vswa in KVCahceManager diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 601ab716dfcb..d8012453a40f 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import dataclasses from typing import List, Optional @@ -190,18 +193,49 @@ class MambaKVCacheParams: n_groups: int # n_groups | linear_num_key_heads head_dim: int # mamba_head_dim | linear_value_head_dim - # Per-layer masks and counts (trailing entries cover MTP/draft layers, - # which are attention-only and carry no Mamba state). + # Target-layer masks and counts. Appended MTP/draft layers need only a + # count because every draft layer is full attention. mamba_layer_mask: List[bool] - full_attention_layer_mask: List[bool] + target_full_attention_layer_mask: List[bool] num_mamba_layers: int - num_full_attention_layers: int + num_draft_layers: int # Dtypes dtype: torch.dtype # config.torch_dtype mamba_ssm_cache_dtype: Optional[ torch.dtype] # quant_config.mamba_ssm_cache_dtype + def get_layer_masks( + self, + *, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, + ) -> tuple[List[bool], List[bool]]: + """Return Mamba and attention masks for one cache manager. + + Target masks use target-model layer indices. Appended one-model draft + layers use indices immediately following the target layers, so a + draft-only manager receives target-sized false prefixes. A combined + manager receives the concatenated target and draft layouts. + """ + target_mamba_mask = list(self.mamba_layer_mask) + target_attention_mask = list(self.target_full_attention_layer_mask) + num_target_layers = len(target_mamba_mask) + num_draft_layers = self.num_draft_layers + draft_attention_mask = [True] * num_draft_layers + + if is_draft and num_draft_layers > 0: + return ( + [False] * (num_target_layers + num_draft_layers), + [False] * num_target_layers + draft_attention_mask, + ) + if use_separate_draft_kv_cache: + return target_mamba_mask, target_attention_mask + return ( + target_mamba_mask + [False] * num_draft_layers, + target_attention_mask + draft_attention_mask, + ) + def get_states_bytes_per_layer(self, mapping) -> int: """Return the total bytes of Mamba state per layer, used for budgeting.""" tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 @@ -219,48 +253,8 @@ def get_states_bytes_per_layer(self, mapping) -> int: return state_bytes_per_layer -def _nemotron_hybrid_layer_masks(config, layer_mask): - pattern = config.hybrid_override_pattern - if layer_mask is None: - return ([c == "*" for c in pattern], [c == "M" for c in pattern]) - - # One-model speculative decoding: layer_mask may extend past the hybrid - # pattern; treat trailing positions as attention-only draft layers. - full_attn, mamba = [], [] - for i, include in enumerate(layer_mask): - if i < len(pattern): - is_attn = pattern[i] == "*" - is_mamba = pattern[i] == "M" - else: - is_attn, is_mamba = True, False - full_attn.append(is_attn and include) - mamba.append(is_mamba and include) - return full_attn, mamba - - -def _qwen3_hybrid_layer_masks(config, layer_mask): - full_attn, mamba = get_qwen3_hybrid_layer_masks(config) - if layer_mask is None: - return full_attn, mamba - - if len(layer_mask) < len(full_attn): - raise ValueError( - "layer_mask is shorter than the Qwen3 hybrid layer pattern") - base_len = len(full_attn) - new_full_attn, new_mamba = [], [] - for i, include in enumerate(layer_mask): - if i < base_len: - new_full_attn.append(full_attn[i] and include) - new_mamba.append(mamba[i] and include) - else: - new_full_attn.append(include) - new_mamba.append(False) - return new_full_attn, new_mamba - - def extract_mamba_kv_cache_params( config, - layer_mask: Optional[List[bool]] = None, spec_config=None, quant_config=None, ) -> MambaKVCacheParams: @@ -270,13 +264,9 @@ def extract_mamba_kv_cache_params( Args: config: HuggingFace model config of a hybrid Mamba model. - layer_mask: Optional per-layer keep mask used by one-model speculative - decoding. Entries past the underlying hybrid pattern length are - treated as attention-only draft layers. When provided, the caller - is responsible for already including spec layers in the mask. - spec_config: When `layer_mask` is None, used to extend the masks with - MTP/draft attention layers (no Mamba state) so they receive KV - cache entries. + spec_config: Optional speculative-decoding config used to describe + appended attention-only MTP/draft layers separately from target + layers. quant_config: Optional, used only to surface `mamba_ssm_cache_dtype`. Returns: @@ -288,31 +278,26 @@ def extract_mamba_kv_cache_params( num_heads = config.mamba_num_heads n_groups = config.n_groups head_dim = config.mamba_head_dim - full_attn_mask, mamba_mask = _nemotron_hybrid_layer_masks( - config, layer_mask) + pattern = config.hybrid_override_pattern + target_full_attn_mask = [layer_type == "*" for layer_type in pattern] + mamba_mask = [layer_type == "M" for layer_type in pattern] elif is_qwen3_hybrid(config): state_size = config.linear_key_head_dim conv_kernel = config.linear_conv_kernel_dim num_heads = config.linear_num_value_heads n_groups = config.linear_num_key_heads head_dim = config.linear_value_head_dim - full_attn_mask, mamba_mask = _qwen3_hybrid_layer_masks( - config, layer_mask) + target_full_attn_mask, mamba_mask = get_qwen3_hybrid_layer_masks(config) else: raise ValueError( f"{type(config).__name__} is not a supported hybrid Mamba config") - # When no explicit layer_mask is given, extend the masks here so MTP/draft - # layers (attention-only, no Mamba state) get KV cache entries. With an - # explicit layer_mask, the caller already encoded those entries. - if layer_mask is None and spec_config is not None: + num_draft_layers = 0 + if spec_config is not None: # Imported lazily to avoid a circular dependency between # config_utils and tensorrt_llm._torch.speculative. from ..speculative.utils import get_num_spec_layers - num_spec_layers = get_num_spec_layers(spec_config) - if num_spec_layers > 0: - full_attn_mask.extend([True] * num_spec_layers) - mamba_mask.extend([False] * num_spec_layers) + num_draft_layers = get_num_spec_layers(spec_config) or 0 mamba_ssm_cache_dtype = None if quant_config is not None: @@ -330,9 +315,9 @@ def extract_mamba_kv_cache_params( n_groups=n_groups, head_dim=head_dim, mamba_layer_mask=mamba_mask, - full_attention_layer_mask=full_attn_mask, + target_full_attention_layer_mask=target_full_attn_mask, num_mamba_layers=sum(mamba_mask), - num_full_attention_layers=sum(full_attn_mask), + num_draft_layers=num_draft_layers, dtype=resolve_hf_torch_dtype(config) or torch.bfloat16, mamba_ssm_cache_dtype=mamba_ssm_cache_dtype, ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 24808e8b55b5..8048ee689fa5 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -86,6 +86,8 @@ KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, KVCacheV2PoolGroupIterationStats, + KVCacheV2SsmLifeCycleIterationStats, + KVCacheV2SsmSnapshotIterationStats, ) from .llm_request import LlmRequest, LlmRequestState, SamplingConfig, get_draft_token_length from .resource_manager import ( @@ -764,6 +766,7 @@ def __init__( execution_stream: Optional[torch.cuda.Stream] = None, is_disagg: bool = False, enable_stats: bool = False, + num_reserved_index_slots: int = 1, **kwargs, ) -> None: self.mapping = mapping @@ -1137,7 +1140,8 @@ def append_to_kv_heads_per_layer( # With pipeline parallelism, multiple microbatches can be in-flight # simultaneously, so we need slots for all concurrent sequences. - # Plus 1 for cuda graph dummy request. + # Reserve stable slots for persistent request IDs such as CUDA-graph + # padding requests. The default preserves the main-branch allocation. # In disaggregated mode, use a coefficient of 2: at any moment up to # `max_num_sequences` requests can be actively generating while another # up to `max_num_sequences` requests are still in KV transfer @@ -1145,10 +1149,15 @@ def append_to_kv_heads_per_layer( # capacity lets the next batch of active requests acquire slots without # waiting for the previous batch's transfers to finish. max_num_sequences = max_batch_size * mapping.pp_size - index_mapper_capacity = max_num_sequences * (2 if is_disagg else 1) + 1 + assert num_reserved_index_slots >= 0, "num_reserved_index_slots must be non-negative" + index_mapper_capacity = ( + max_num_sequences * (2 if is_disagg else 1) + num_reserved_index_slots + ) logger.info( f"KVCacheManagerV2: IndexMapper capacity={index_mapper_capacity} " - f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, max_beam_width={max_beam_width})" + f"(max_num_sequences={max_num_sequences}, is_disagg={is_disagg}, " + f"num_reserved_index_slots={num_reserved_index_slots}, " + f"max_beam_width={max_beam_width})" ) self.index_mapper = IndexMapper(index_mapper_capacity, max_beam_width) self._early_freed_index_requests: set[int] = set() @@ -1156,26 +1165,50 @@ def append_to_kv_heads_per_layer( self._log_kv_cache_pool_lifecycle_mapping() - def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: + """Return the roles represented by the two page-table index lanes. + + When present, role B must be addressable from role A using a constant + page-index offset. + """ + role_b = None if self.kv_cache_type == CacheTypeCpp.SELFKONLY else Role.VALUE + return Role.KEY, role_b + + def _get_block_scale_role(self, role_a: DataRole) -> Optional[DataRole]: + if self.dtype != DataType.NVFP4 or role_a != Role.KEY: + return None + return Role.KEY_BLOCK_SCALE + + def _build_pool_mapping_tensors(self): + """Build the (kv_cache_pool_pointers, kv_cache_pool_mapping) tensors. + + An overridable hook for subclasses whose pools coalesce extra + per-layer buffers alongside K/V. + """ kv_cache_pool_pointers_list = [] kv_cache_pool_mapping_list = [] block_scale_pool_pointers_list = [] if self.enable_swa_scratch_reuse: 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) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.PER_LAYER + layer_id, role_a, PageIndexMode.PER_LAYER ), 0, ] ) if self.dtype == DataType.NVFP4: + block_scale_role = self._get_block_scale_role(role_a) block_scale_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.PER_LAYER - ), + layer_id, block_scale_role, PageIndexMode.PER_LAYER + ) + if block_scale_role is not None + else 0, 0, ] ) @@ -1183,43 +1216,67 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: else: for pool_id in range(self.num_pools): layer_id = self.impl.layer_grouping[pool_id][0] + role_a, _ = self._get_pool_roles(pool_id) kv_cache_pool_pointers_list.append( [ - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED - ), + self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), 0, ] ) if self.dtype == DataType.NVFP4: + block_scale_role = self._get_block_scale_role(role_a) block_scale_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.SHARED - ), + layer_id, block_scale_role, PageIndexMode.SHARED + ) + if block_scale_role is not None + else 0, 0, ] ) for layer_id in typed_range(LayerId(self.num_local_layers)): layer_group_id = self.impl.get_layer_group_id(layer_id) - key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - offset = self._kv_pool_mapping_offset(layer_id, layer_group_id, key_base_addr) - - if self.dtype == DataType.NVFP4: - block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] - block_scale_addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.SHARED - ) - - block_scale_base_addr + role_a, role_b = self._get_pool_roles(layer_group_id) + index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] + if role_a == Role.KEY: + offset = self._kv_pool_mapping_offset(layer_id, layer_group_id, index_base_addr) + else: + addr_offset = ( + self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED) + - index_base_addr ) - block_scale_offset = exact_div( - block_scale_addr_offset, - self.get_layer_bytes_per_token(layer_id, Role.KEY_BLOCK_SCALE) - * self.kv_factor - * self.tokens_per_block, + offset_divisor = self.impl.get_page_stride(layer_id, role_a) + if role_b is not None: + offset_divisor *= self.kv_factor + offset = exact_div( + addr_offset, + offset_divisor, ) + + if self.dtype != DataType.NVFP4 or role_a != Role.KEY: + block_scale_offset = None + else: + block_scale_role = self._get_block_scale_role(role_a) + if block_scale_role is None: + block_scale_offset = None + else: + block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] + block_scale_addr_offset = ( + self.impl.get_mem_pool_base_address( + layer_id, block_scale_role, PageIndexMode.SHARED + ) + - block_scale_base_addr + ) + block_scale_offset = exact_div( + block_scale_addr_offset, + self.get_layer_bytes_per_token(layer_id, block_scale_role) + * self.kv_factor + * self.tokens_per_block, + ) + + if block_scale_offset is not None: assert block_scale_offset == offset, ( "Block scale offset and offset should be the same" ) @@ -1234,18 +1291,22 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: [pool_pointers[1], block_scale_pool_pointers[1]], ] - self.kv_cache_pool_pointers = torch.tensor( + kv_cache_pool_pointers = torch.tensor( kv_cache_pool_pointers_list, dtype=torch.int64, device="cpu", pin_memory=prefer_pinned(), ) - self.kv_cache_pool_mapping = torch.tensor( + kv_cache_pool_mapping = torch.tensor( kv_cache_pool_mapping_list, dtype=torch.int32, device="cpu", pin_memory=prefer_pinned(), ) + return kv_cache_pool_pointers, kv_cache_pool_mapping + + def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: + self.kv_cache_pool_pointers, self.kv_cache_pool_mapping = self._build_pool_mapping_tensors() self.index_scales = torch.empty( self.num_pools, dtype=torch.int32, pin_memory=prefer_pinned(), device="cpu" ) @@ -1254,12 +1315,13 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None: ) for pool_id in range(self.num_pools): layer_id = self.impl.layer_grouping[pool_id][0] - self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, Role.KEY) - if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + role_a, role_b = self._get_pool_roles(pool_id) + self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, role_a) + if role_b is not None: self.kv_offset[pool_id] = exact_div( - self.impl.get_mem_pool_base_address(layer_id, Role.VALUE, PageIndexMode.SHARED) - - self.impl.get_mem_pool_base_address(layer_id, Role.KEY, PageIndexMode.SHARED), - self.impl.get_page_stride(layer_id, Role.KEY), + self.impl.get_mem_pool_base_address(layer_id, role_b, PageIndexMode.SHARED) + - self.impl.get_mem_pool_base_address(layer_id, role_a, PageIndexMode.SHARED), + self.impl.get_page_stride(layer_id, role_a), ) else: self.kv_offset[pool_id] = 0 @@ -1400,7 +1462,8 @@ def _get_event_window_sizes_by_layer_group(self) -> Dict[int, int]: # mixed windows in one group, this needs to fan out per-layer. def get_event_window_size(layer_id: int) -> int: - window_size = self.kv_cache_manager_py_config.layers[layer_id].sliding_window_size + layer_config = self.kv_cache_manager_py_config.layers[layer_id] + window_size = getattr(layer_config, "sliding_window_size", None) return self.max_seq_len if window_size is None else int(window_size) return { @@ -1451,9 +1514,11 @@ 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] - roles = [Role.KEY, Role.VALUE] - if self.kv_cache_type == CacheTypeCpp.SELFKONLY: - roles[1] = Role.KEY + role_a, role_b = self._get_pool_roles(pool_id) + roles = [ + role_a, + role_a if role_b is None else role_b, + ] for role_idx, role in enumerate(roles): converter = self.impl.get_page_index_converter(layer_id, role) if converter.expansion != 1: @@ -2736,7 +2801,7 @@ def _build_pool_group_iteration_stats( ), ) - def _build_life_cycle_iteration_stats( + def _build_attention_life_cycle_iteration_stats( self, life_cycle_id: int, life_cycle_metadata, @@ -2747,6 +2812,7 @@ def _build_life_cycle_iteration_stats( reuse_delta, ) -> KVCacheV2LifeCycleIterationStats: pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] + assert kind == "attention" return KVCacheV2LifeCycleIterationStats( life_cycle_id=life_cycle_id, pool_group_id=pool_group_id, @@ -2763,6 +2829,29 @@ def _build_life_cycle_iteration_stats( ), ) + @staticmethod + def _build_ssm_life_cycle_iteration_stats( + life_cycle_id: int, + life_cycle_metadata, + snapshot_delta, + ) -> KVCacheV2SsmLifeCycleIterationStats: + pool_group_id, window_size, kind = life_cycle_metadata[life_cycle_id] + assert kind == "ssm" + assert window_size is None + return KVCacheV2SsmLifeCycleIterationStats( + life_cycle_id=life_cycle_id, + pool_group_id=pool_group_id, + snapshot_stats=KVCacheV2SsmSnapshotIterationStats( + iter_snapshot_lookups=snapshot_delta.iter_snapshot_lookups, + iter_snapshot_hits=snapshot_delta.iter_snapshot_hits, + iter_snapshot_misses=snapshot_delta.iter_snapshot_misses, + iter_reused_tokens=snapshot_delta.iter_reused_tokens, + iter_unreused_tokens=snapshot_delta.iter_unreused_tokens, + iter_aligned_snapshot_hits=snapshot_delta.iter_aligned_snapshot_hits, + iter_unaligned_snapshot_hits=snapshot_delta.iter_unaligned_snapshot_hits, + ), + ) + def get_kv_cache_stats(self): kv_cache_stats = KvCacheStats() pool_group_stats = self._get_storage_statistics(GPU_LEVEL) @@ -2811,6 +2900,7 @@ def get_iteration_stats(self): 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() + raw_ssm_snapshot_iteration_stats = self.impl.get_and_reset_ssm_snapshot_iteration_stats() 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 = [ @@ -2847,7 +2937,10 @@ def get_iteration_stats(self): for window_size in sorted(windows) } - pool_group_ids = sorted(set(windows_by_pool_group) | set(pool_group_deltas)) + all_pool_group_ids = set(range(len(primary_stats))) + pool_group_ids = sorted( + all_pool_group_ids | set(windows_by_pool_group) | set(pool_group_deltas) + ) stats_by_pool_group = { pool_group_id: self._build_pool_group_iteration_stats( pool_group_id, @@ -2862,7 +2955,7 @@ def get_iteration_stats(self): } stats_by_life_cycle = { - life_cycle_id: self._build_life_cycle_iteration_stats( + life_cycle_id: self._build_attention_life_cycle_iteration_stats( life_cycle_id, life_cycle_metadata, primary_stats, @@ -2873,6 +2966,14 @@ def get_iteration_stats(self): ) for life_cycle_id, reuse_delta in sorted(reuse_deltas_by_life_cycle.items()) } + for life_cycle_id, snapshot_delta in sorted(raw_ssm_snapshot_iteration_stats.items()): + assert int(life_cycle_id) not in stats_by_life_cycle + stats_by_life_cycle[int(life_cycle_id)] = self._build_ssm_life_cycle_iteration_stats( + int(life_cycle_id), + life_cycle_metadata, + snapshot_delta, + ) + stats_by_life_cycle = dict(sorted(stats_by_life_cycle.items())) return KVCacheV2IterationStatsReport( stats_by_window, stats_by_pool_group, stats_by_life_cycle @@ -3274,19 +3375,25 @@ def calculate_scaling_factor_size_bytes( ) return get_size_in_bytes(cache_size // quant_vector_size, scaling_factor_dtype) - def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: - some_checks_unavailable = False - has_invalid_values = torch.tensor( - [False], dtype=torch.bool, device=torch.cuda.current_device() - ) + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: pool_handled = set() - - # Handle each layer from start to end to traverse the whole KV cache. 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 + yield buffer + pool_handled.add(pool_id) + + def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool: + some_checks_unavailable = False + has_invalid_values = torch.tensor( + [False], dtype=torch.bool, device=torch.cuda.current_device() + ) + + for buffer in self._iter_cache_buffers_for_invalid_check(): # process in chunks of 256 pages to avoid OoM for i in range(0, buffer.shape[0], 256): buffer_slice = buffer[i : i + 256] @@ -3297,7 +3404,6 @@ def check_invalid_values_in_kv_cache(self, fill_with_zero: bool = False) -> bool some_checks_unavailable = True if fill_with_zero: buffer.zero_() - pool_handled.add(pool_id) torch.cuda.synchronize() if some_checks_unavailable: @@ -3425,10 +3531,19 @@ def update_resources( # will be resumed by the scheduler on the next iteration. if not kv_cache.is_active: continue + rewind_len = req.py_rewind_len + if self.is_draft: + runtime_draft_len = req.py_rewind_len + req.py_num_accepted_draft_tokens + # Dynamic-tree draft managers reserve K * max_draft_len slots, + # which can exceed the tree's runtime draft width. Reclaim that + # reserve slack together with rejected draft tokens; otherwise + # it accumulates in the draft KV cache after every generation + # step. Target managers do not allocate this reserve slack. + rewind_len += max(self._kv_reserve_draft_tokens - runtime_draft_len, 0) new_capacity = ( None if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) - else kv_cache.capacity - req.py_rewind_len + else kv_cache.capacity - rewind_len ) history_length = ( None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py index d9841ccbfe86..5ffa3742f3ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py @@ -14,7 +14,7 @@ # limitations under the License. from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal KV_CACHE_ITERATION_STATS_REUSE_KEYS = ( "iterReusedBlocks", @@ -70,11 +70,39 @@ class KVCacheV2LifeCycleIterationStats: stats: Any +@dataclass(slots=True) +class KVCacheV2SsmSnapshotIterationStats: + iter_snapshot_lookups: int + iter_snapshot_hits: int + iter_snapshot_misses: int + iter_reused_tokens: int + iter_unreused_tokens: int + iter_aligned_snapshot_hits: int + iter_unaligned_snapshot_hits: int + + @property + def iter_snapshot_hit_rate(self) -> float: + if self.iter_snapshot_hits == 0 or self.iter_snapshot_lookups == 0: + return 0.0 + return self.iter_snapshot_hits / self.iter_snapshot_lookups + + +@dataclass(slots=True) +class KVCacheV2SsmLifeCycleIterationStats: + life_cycle_id: int + pool_group_id: int + snapshot_stats: KVCacheV2SsmSnapshotIterationStats + window_size: None = field(default=None, init=False) + kind: Literal["ssm"] = field(default="ssm", init=False) + + @dataclass(slots=True) class KVCacheV2IterationStatsReport: by_window_size: dict[int, Any] by_pool_group: dict[int, KVCacheV2PoolGroupIterationStats] - by_life_cycle: dict[int, KVCacheV2LifeCycleIterationStats] = field(default_factory=dict) + by_life_cycle: dict[ + int, KVCacheV2LifeCycleIterationStats | KVCacheV2SsmLifeCycleIterationStats + ] = field(default_factory=dict) def serialize_kv_cache_iteration_stats(stats, keys: tuple[str, ...] | None = None) -> dict: @@ -115,6 +143,21 @@ def serialize_kv_cache_iteration_stats(stats, keys: tuple[str, ...] | None = Non return {key: fields[key] for key in keys} +def serialize_ssm_snapshot_iteration_stats( + stats: KVCacheV2SsmSnapshotIterationStats, +) -> dict: + return { + "iterSnapshotLookups": stats.iter_snapshot_lookups, + "iterSnapshotHits": stats.iter_snapshot_hits, + "iterSnapshotMisses": stats.iter_snapshot_misses, + "iterSnapshotHitRate": stats.iter_snapshot_hit_rate, + "iterReusedTokens": stats.iter_reused_tokens, + "iterUnreusedTokens": stats.iter_unreused_tokens, + "iterAlignedSnapshotHits": stats.iter_aligned_snapshot_hits, + "iterUnalignedSnapshotHits": stats.iter_unaligned_snapshot_hits, + } + + def append_kv_cache_iteration_stats(stats_dict: dict, kv_iter_stats) -> None: if kv_iter_stats is None: return @@ -147,13 +190,21 @@ def append_kv_cache_iteration_stats(stats_dict: dict, kv_iter_stats) -> None: if not kv_iter_stats.by_life_cycle: return - stats_dict["kvCacheIterationStatsByLifecycle"] = { - str(life_cycle_id): { + stats_by_life_cycle = {} + for life_cycle_id, stats in kv_iter_stats.by_life_cycle.items(): + serialized = { "lifeCycleId": stats.life_cycle_id, "poolGroupId": stats.pool_group_id, "windowSize": stats.window_size, "kind": stats.kind, - **serialize_kv_cache_iteration_stats(stats.stats, KV_CACHE_ITERATION_STATS_REUSE_KEYS), } - for life_cycle_id, stats in kv_iter_stats.by_life_cycle.items() - } + if isinstance(stats, KVCacheV2SsmLifeCycleIterationStats): + serialized["snapshotStats"] = serialize_ssm_snapshot_iteration_stats( + stats.snapshot_stats + ) + else: + serialized.update( + serialize_kv_cache_iteration_stats(stats.stats, KV_CACHE_ITERATION_STATS_REUSE_KEYS) + ) + stats_by_life_cycle[str(life_cycle_id)] = serialized + stats_dict["kvCacheIterationStatsByLifecycle"] = stats_by_life_cycle diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index a686adc97528..9b9ea2556e1e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -16,6 +16,7 @@ from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, MixedMambaHybridCacheManager) from .resource_manager import KVCacheManager @@ -157,14 +158,33 @@ def create_kv_cache_transceiver( "UCX_CUDA_IPC_ENABLE_MNNVL=n, UCX_RNDV_SCHEME=put_zcopy and/or unset UCX_NET_DEVICES upon server " "hangs or lower-than-expected performance.") - # Select transceiver implementation based on transceiver_runtime + # Select transceiver implementation based on transceiver_runtime. # transceiver_runtime == None or "CPP" -> use C++ transceiver (default) - # transceiver_runtime == "PYTHON" -> use Python transceiver - if cache_transceiver_config.transceiver_runtime == "PYTHON": - # Python transceiver currently only supports NIXL and DEFAULT backend - if cache_transceiver_config.backend not in ("DEFAULT", "NIXL"): + # transceiver_runtime == "PYTHON" -> use Python transceiver. + # + # MambaHybridCacheManagerV2 is backed by the Python KVCacheManagerV2 core, + # not the C++ BaseKVCacheManager binding required by CacheTransceiverCpp. + is_v2_mamba_hybrid = isinstance(mamba_cache_manager, + MambaHybridCacheManagerV2) + use_python_transceiver = ( + cache_transceiver_config.transceiver_runtime == "PYTHON") + + if is_v2_mamba_hybrid and not use_python_transceiver: + raise ValueError( + "MambaHybridCacheManagerV2 requires transceiver_runtime='PYTHON' " + "with backend='NIXL'; it cannot use the C++ transceiver.") + + if use_python_transceiver: + if isinstance(mamba_cache_manager, CppMambaHybridCacheManager): + raise ValueError( + "transceiver_runtime='PYTHON' cannot drive " + "CppMambaHybridCacheManager (C++ pool backed). Use " + "transceiver_runtime='CPP', or select the V2 manager " + "with use_kv_cache_manager_v2=True.") + # DEFAULT has already been resolved above, so Python must see NIXL. + if cache_transceiver_config.backend != "NIXL": raise ValueError( - f"Python transceiver currently only supports NIXL or DEFAULT backend, " + f"Python transceiver currently only supports the NIXL backend, " f"got {cache_transceiver_config.backend}. " f"Please use transceiver_runtime='CPP' for MPI, UCX, or MOONCAKE backends." ) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 16fe22d1c777..e7cab986c2ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -16,8 +16,9 @@ import math import os from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Union +from dataclasses import dataclass, replace +from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, + Optional, Tuple, Union) import torch import triton @@ -27,19 +28,29 @@ from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import ( + BlockReusePolicy, KVCacheManagerV2, Role) from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest) from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, CacheTypeCpp, DataType, KVCacheManager, PoolConfiguration, get_pp_layers) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests -from tensorrt_llm._utils import (nvtx_range, prefer_pinned, +from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, + nvtx_range, prefer_pinned, torch_dtype_to_binding) from tensorrt_llm.bindings.internal.batch_manager import ( LinearAttentionMetadata, LinearCacheType) from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, + BatchDesc, BufferConfig, + DataRole, KVCacheDesc) +from tensorrt_llm.runtime.kv_cache_manager_v2 import \ + KVCacheManagerConfig as KVCacheManagerConfigPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, + SsmLayerConfig) GB = 1 << 30 @@ -58,6 +69,63 @@ MIN_REPLAY_HISTORY_SIZE = 16 +def _get_num_cuda_graph_padding_dummy_slots( + spec_config: Optional["DecodingBaseConfig"], + max_batch_size: int, +) -> int: + """Return the number of persistent CUDA-graph padding dummy IDs. + + This is computed before ``ModelEngine`` exists and covers draft lengths + reachable at every batch size, including the zero-length acceptance-rate + fallback. ``ModelEngine._compute_dynamic_draft_len_mapping`` is created + later and covers only configured CUDA-graph batch sizes, so it cannot size + this persistent ID set. + """ + if spec_config is None: + return 1 + + draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + supports_dynamic_draft_len = (spec_dec_mode is not None and hasattr( + spec_dec_mode, "support_dynamic_draft_len") + and spec_dec_mode.support_dynamic_draft_len()) + if draft_len_schedule and supports_dynamic_draft_len: + runtime_draft_lengths = set() + first_uncovered_batch_size = 1 + for batch_size_threshold, draft_len in draft_len_schedule.items(): + if first_uncovered_batch_size > max_batch_size: + break + if batch_size_threshold >= first_uncovered_batch_size: + runtime_draft_lengths.add(draft_len) + first_uncovered_batch_size = batch_size_threshold + 1 + if first_uncovered_batch_size <= max_batch_size: + runtime_draft_lengths.add(0) + else: + max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 + max_total_draft_tokens = (getattr(spec_config, "max_total_draft_tokens", + 0) or 0) + is_linear_tree = getattr( + spec_config, + "is_linear_tree", + max_draft_len == max_total_draft_tokens, + ) + static_draft_len = (max_draft_len + if is_linear_tree else max_total_draft_tokens) + runtime_draft_lengths = {static_draft_len or 0} + + if ((getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and + (getattr(spec_config, "acceptance_rate_threshold", 0) or 0) > 0): + runtime_draft_lengths.add(0) + return len(runtime_draft_lengths) + + +class MambaRole: + """V2 buffer roles owned only by the hybrid Mamba manager.""" + + SSM_STATE = DataRole("ssm_state") + CONV_STATE = DataRole("conv_state") + + def get_tensor_size_bytes(tensor): """Calculate tensor size in bytes.""" if isinstance(tensor, torch.Tensor): @@ -146,9 +214,9 @@ def use_py_mamba_cache_manager() -> bool: Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager - with PythonMambaCacheManager inside instead of the default unified-pool - CppMambaHybridCacheManager. Disagg mode is unaffected — it already picks - PythonMambaCacheManager when transceiver_runtime='PYTHON'. + with PythonMambaCacheManager inside instead of the configured manager. + Disagg mode is unaffected — its compatibility routing selects Mixed or + Cpp based on the transceiver configuration. """ return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' @@ -161,6 +229,40 @@ class ReplayStateUpdateMetadata(NamedTuple): replay_history_size: int +def _advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: Optional[torch.Tensor] = None, +) -> None: + """Advance replay bookkeeping after a speculative generation step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to( + replay_metadata.prev_num_accepted_tokens.dtype) + prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = (prev_num_accepted_tokens + + replay_metadata.replay_step_width + > replay_metadata.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, + accepted_tokens, + prev_num_accepted_tokens + accepted_tokens, + ) + cache_buf_idx = replay_metadata.cache_buf_idx[slots] + next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx) + if is_dummy_request is not None: + next_num_accepted_tokens = torch.where( + is_dummy_request, + prev_num_accepted_tokens, + next_num_accepted_tokens, + ) + next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, + next_cache_buf_idx) + replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens + replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx + + class BaseMambaCacheManager(ABC): """Abstract interface for accessing mamba/recurrent state caches.""" @@ -802,33 +904,16 @@ def update_mamba_states( src_state_indices = self.intermediate_state_indices[:num_gens] if self._use_replay_state_update: - # SSM state is handled incrementally by the kernel. Mirror the - # kernel's per-slot checkpoint predicate from the previous PNAT and - # fixed replay step width: checkpoint steps write a fresh history - # buffer and flip, while no-checkpoint steps append to the active - # buffer and keep reading from it next step. - accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + - num_gens] - prev_num_accepted_tokens = \ - self.mamba_cache.prev_num_accepted_tokens[state_indices_d] - wrote_checkpoint = (prev_num_accepted_tokens + - self.replay_step_width - > self.replay_history_size) - next_num_accepted_tokens = torch.where( - wrote_checkpoint, accepted_tokens, - prev_num_accepted_tokens + accepted_tokens) - cache_buf_idx = self.mamba_cache.cache_buf_idx[state_indices_d] is_dummy_request = self._dummy_request_mask[ num_contexts:num_contexts + num_gens] - next_num_accepted_tokens = torch.where(is_dummy_request, - prev_num_accepted_tokens, - next_num_accepted_tokens) - self.mamba_cache.prev_num_accepted_tokens[state_indices_d] = \ - next_num_accepted_tokens - self.mamba_cache.cache_buf_idx[state_indices_d] = \ - torch.where(is_dummy_request, cache_buf_idx, - torch.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx)) + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) else: # Legacy: copy accepted SSM state from intermediate cache. ssm_states = self.mamba_cache.temporal @@ -1001,13 +1086,257 @@ def update_mamba_states( class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Marker base class for hybrid mamba cache manager implementations. + """Shared interface and state plumbing for hybrid Mamba managers. - Used purely for ``isinstance`` / type-hint purposes so callers can refer - to the family without caring about the concrete implementation. Concrete - selection (Mixed vs Cpp) lives in ``_util.py:_get_model_kv_cache_manager_cls``. + Concrete storage, state-index, and resource lifecycles remain owned by the + Cpp and V2 implementations. """ + _supports_additional_snapshot_offsets = False + + def _setup_mtp_intermediate_states(self, spec_config, + max_batch_size: int) -> None: + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if spec_config is None or self.local_num_mamba_layers == 0: + return + + tokens_per_gen_step = spec_config.tokens_per_gen_step + if not self._use_replay_state_update: + self.intermediate_ssm_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, + tokens_per_gen_step + ] + self.ssm_state_shape, + dtype=self.ssm_state_dtype, + device="cuda", + ) + self.intermediate_conv_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step + ] + self.conv_state_shape, + dtype=self.conv_state_dtype, + device="cuda", + ) + self.intermediate_state_indices = torch.arange(max_batch_size, + dtype=torch.int32, + device="cuda") + + def _allocate_pool_replay_buffers( + self, + spec_config, + cache_size: int, + device: Optional[torch.device], + ) -> bool: + """Allocate replay tensors shared by the Cpp and V2 state pools.""" + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + if (self.local_num_mamba_layers == 0 + or (not self._use_replay_state_update + and not self._mamba_ssm_stochastic_rounding)): + return False + + assert device is not None + self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + cache_size, self._seed_rank_offset, device) + if spec_config is None or not self._use_replay_state_update: + return False + + history_size = self.replay_history_size + assert history_size is not None + nheads, head_dim, d_state = self.ssm_state_shape + common_shape = [self.local_num_mamba_layers, cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.cache_buf_idx = torch.zeros(cache_size, + dtype=torch.int32, + device=device) + self.old_x = torch.zeros( + common_shape + [history_size, nheads, head_dim], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_shape + [history_size, self._n_groups_per_rank, d_state], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + return True + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + if self._dummy_request_mask is None: + return + + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, + non_blocking=True) + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + if num_contexts == 0: + return + + context_slots = self.cuda_state_indices[:num_contexts].long() + if (self._use_replay_state_update + and self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None): + self.prev_num_accepted_tokens[context_slots] = 0 + self.cache_buf_idx[context_slots] = 0 + if self.old_x is not None: + self.old_x[:, context_slots] = 0 + if self.old_B is not None: + self.old_B[:, context_slots] = 0 + if self.old_dt is not None: + self.old_dt[:, context_slots] = 0 + if self.old_dA_cumsum is not None: + self.old_dA_cumsum[:, context_slots] = 0 + + if self.mamba_ssm_rand_seed is None: + return + self._seed_request_counter += 1 + counter = self._seed_request_counter + rank_offset = self._seed_rank_offset + host_slots = self._host_state_indices[:num_contexts].tolist() + new_seeds = [ + _compute_deterministic_mamba_seed(counter, slot, rank_offset) + for slot in host_slots + ] + self.mamba_ssm_rand_seed[context_slots] = torch.tensor( + new_seeds, + dtype=torch.int64, + device=self.mamba_ssm_rand_seed.device, + ) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + """Set reusable Mamba snapshot boundaries before scheduling.""" + if not self.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + state_config = self.kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + for request in requests: + snapshot_points = set() + if interval is not None and interval > 0: + snapshot_points.update( + range(interval, request.prompt_len + 1, interval)) + if self._supports_additional_snapshot_offsets: + for offset in state_config.additional_snapshot_offsets_from_start: + if offset <= request.prompt_len: + snapshot_points.add(offset) + for offset in state_config.additional_snapshot_offsets_from_end: + point = request.prompt_len - offset + if point > 0: + snapshot_points.add(point) + request.expect_snapshot_points = sorted(snapshot_points) + + def is_speculative(self) -> bool: + return self.spec_config is not None + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_ssm_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_ssm_states is None: + return None + return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_conv_states(self, + layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_conv_states is None: + return None + return self.intermediate_conv_states[ + self.mamba_layer_offsets[layer_idx]] + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, + PythonMambaCacheManager.SpeculativeState, None]: + conv = self.get_conv_states(layer_idx) + ssm = self.get_ssm_states(layer_idx) + if self.spec_config is not None: + layer_offset = self.mamba_layer_offsets[layer_idx] + spec_kwargs = {} + if self.mamba_ssm_rand_seed is not None: + spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed + if self._use_replay_state_update: + spec_kwargs['old_x'] = self.old_x[layer_offset] + spec_kwargs['old_B'] = self.old_B[layer_offset] + spec_kwargs['old_dt'] = self.old_dt[layer_offset] + spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] + spec_kwargs['cache_buf_idx'] = self.cache_buf_idx + spec_kwargs['prev_num_accepted_tokens'] = ( + self.prev_num_accepted_tokens) + else: + spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ + layer_offset] + return PythonMambaCacheManager.SpeculativeState( + conv=conv, + temporal=ssm, + intermediate_conv_window=self. + intermediate_conv_states[layer_offset], + **spec_kwargs, + ) + return PythonMambaCacheManager.State(conv=conv, temporal=ssm) + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', + None) + cache_buf_idx = getattr(self, 'cache_buf_idx', None) + if (not self._use_replay_state_update + or prev_num_accepted_tokens is None or cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.ssm_state_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + return getattr(self, 'mamba_ssm_rand_seed', None) + def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" @@ -1236,12 +1565,169 @@ def _promote_mamba_state_triton(dst: torch.Tensor, ) +def _mamba_snapshot_rule_counts( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], + tokens_per_block: int, +) -> Tuple[int, int]: + """Return reachable fixed rules and their partial-block upper bound.""" + if not kv_cache_config.enable_block_reuse: + return 0, 0 + + num_rules = 0 + num_unaligned_rules = 0 + state_config = kv_cache_config.mamba_state_config + for offset in set(state_config.additional_snapshot_offsets_from_start): + if max_seq_len is not None and offset > max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += int(offset % tokens_per_block != 0) + for offset in set(state_config.additional_snapshot_offsets_from_end): + # A from-end offset is reachable iff some valid prompt is longer than + # the offset. Its absolute alignment depends on that prompt. + if max_seq_len is not None and offset >= max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += 1 + return num_rules, num_unaligned_rules + + +def _mamba_regular_snapshot_interval( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], +) -> Optional[int]: + if not kv_cache_config.enable_block_reuse: + return None + interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval + if interval is None or interval <= 0: + return None + if max_seq_len is not None and interval > max_seq_len: + return None + return interval + + +def _get_local_mamba_cache_layout( + model_config, + mapping: Mapping, + *, + spec_config=None, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, +): + """Return normalized params and local Mamba/attention layer counts. + + Cache construction and affine sizing must follow the model's PP layout: + partition base layers first, then place appended speculative layers on the + last PP rank. The normalized params retain target masks and the appended + draft-layer count so estimation selects the same combined or per-manager + layout as runtime. + """ + from tensorrt_llm._torch.pyexecutor.config_utils import \ + extract_mamba_kv_cache_params + + params = extract_mamba_kv_cache_params( + model_config.pretrained_config, + spec_config=spec_config, + quant_config=model_config.quant_config, + ) + mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + combined_layer_mask = [ + is_mamba or is_attention for is_mamba, is_attention in zip( + mamba_layer_mask, full_attention_layer_mask) + ] + local_layer_indices, _ = get_pp_layers( + sum(combined_layer_mask), + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + local_mamba_layers = sum(mamba_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + local_attention_layers = sum(full_attention_layer_mask[layer_idx] + for layer_idx in local_layer_indices) + return params, local_mamba_layers, local_attention_layers + + +def _estimate_mamba_hybrid_cache_cost( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + tokens_per_block: int, + max_seq_len: Optional[int], + num_reserved_dummy_slots: int, + include_explicit_snapshots: bool, + cap_partial_attention_snapshots: bool, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, + **kwargs, +) -> Tuple[int, int]: + spec_config = kwargs.get("spec_config") + params, local_mamba_layers, local_attention_layers = ( + _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + )) + attention_slope = (KVCacheManager.get_cache_size_per_token( + model_config, + mapping, + num_layers=local_attention_layers, + **kwargs, + ) if local_attention_layers > 0 else 0) + state_bytes_per_rank = (local_mamba_layers * + params.get_states_bytes_per_layer(mapping)) + max_resident_sequences = max_batch_size * mapping.pp_size + + if include_explicit_snapshots: + fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( + kv_cache_config, max_seq_len, tokens_per_block) + else: + fixed_rules = 0 + unaligned_fixed_rules = 0 + fixed_state_slots = (max_resident_sequences + num_reserved_dummy_slots + + max_resident_sequences * fixed_rules) + attention_block_bytes = attention_slope * tokens_per_block + + interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) + has_unaligned_periodic_snapshot = (interval is not None + and interval % tokens_per_block != 0) + if cap_partial_attention_snapshots: + # Snapshot alignment is unknown while estimating cache cost. Once a + # snapshot is possible, reserve one retained partial attention page + # per resident lineage. Dummy requests carry no attention capacity. + has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None + partial_attention_slots = (max_resident_sequences + if has_non_live_ssm_capacity else 0) + else: + partial_attention_slots = (max_resident_sequences * + unaligned_fixed_rules) + intercept = (fixed_state_slots * state_bytes_per_rank + + partial_attention_slots * attention_block_bytes) + + if interval is None: + regular_slope = 0 + else: + regular_slope = math.ceil(state_bytes_per_rank / interval) + if (has_unaligned_periodic_snapshot + and not cap_partial_attention_snapshots): + regular_slope += math.ceil(attention_block_bytes / interval) + return attention_slope + regular_slope, intercept + + class CppMambaHybridCacheManager(KVCacheManager, MambaHybridCacheManager): """Hybrid cache manager storing mamba states inside the KVCacheManager pool. Both KV cache blocks and recurrent state blocks are managed by the unified C++ KVCacheManager, enabling block reuse / prefix caching across attention - and mamba layers. This is the default hybrid manager. + and mamba layers. This compatibility manager remains available through the + manager preference override and legacy disaggregated routing. """ @@ -1334,6 +1820,12 @@ def __init__( # seed values without any torch.randint. self._seed_request_counter = 0 self.ssm_state_dtype = mamba_ssm_cache_dtype + # Keep the shared Mamba interface valid on PP ranks that do not own a + # local Mamba layer. + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None if self.local_num_mamba_layers == 0: logger.info( @@ -1360,7 +1852,7 @@ def __init__( self.kv_cache_config = kv_cache_config self.linear_attention_metadata = LinearAttentionMetadata() self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_cache_interval + kv_cache_config.mamba_state_config.periodic_snapshot_interval if kv_cache_config.enable_block_reuse else 0) return @@ -1398,10 +1890,10 @@ def __init__( # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; # we just need to tell C++ which ordering to use. self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" - self.ssm_count = math.prod(self.ssm_state_shape) - self.conv_count = math.prod(self.conv_state_shape) - self.ssm_bytes = self.ssm_count * self.ssm_state_dtype.itemsize - self.conv_bytes = self.conv_count * self.conv_state_dtype.itemsize + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) total_bytes = self.ssm_bytes + self.conv_bytes if total_bytes % self.ssm_state_dtype.itemsize != 0: @@ -1420,7 +1912,7 @@ def __init__( self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_cache_interval + kv_cache_config.mamba_state_config.periodic_snapshot_interval if kv_cache_config.enable_block_reuse else 0) # RNN model params for disagg TP-mismatch split/concat. conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} @@ -1525,10 +2017,6 @@ def __init__( device="cpu") self._request_id_to_state_index = {} self._request_id_to_is_dummy = {} - # Batch-order mask aligned with state_indices; duplicate dummy request - # IDs mark every batch row even when they share one cache slot. - self._dummy_request_mask = None - self._dummy_request_mask_host = None self.kv_cache_config = kv_cache_config self.is_estimating_kv_cache = is_estimating_kv_cache @@ -1549,70 +2037,35 @@ def get_cache_size_per_token( max_batch_size: int, kv_cache_config: KvCacheConfig, num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, **kwargs, ): """Affine memory model for the unified hybrid KV pool. Returns ``(slope_bytes_per_token, intercept_bytes)``: - * ``slope`` = attention KV bytes per token (parent's formula) plus - the amortized regular-snapshot bytes per token from mamba layers. - * ``intercept`` = ``max_batch_size * num_mamba_layers_per_rank * - state_bytes_per_layer * STATIC_SLOTS_PER_REQUEST``. + * ``slope`` = attention KV bytes per token plus conservative periodic + Mamba-state and partial-attention-snapshot costs. + * ``intercept`` = live and CUDA-graph dummy Mamba state. Memory budget -> max tokens then becomes ``T = (budget - intercept) // slope`` instead of plain ``T = budget // bytes_per_token``. """ - # Lazy import to avoid pulling config_utils into module import order. - from tensorrt_llm._torch.pyexecutor.config_utils import \ - extract_mamba_kv_cache_params - - # Attention slope from the parent's existing formula. - attention_slope = KVCacheManager.get_cache_size_per_token( - model_config, mapping, num_layers=num_layers, **kwargs) - - params = extract_mamba_kv_cache_params( - model_config.pretrained_config, - quant_config=model_config.quant_config, + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=1, + include_explicit_snapshots=False, + cap_partial_attention_snapshots=False, + **kwargs, ) - state_bytes_per_layer = params.get_states_bytes_per_layer(mapping) - - # This not precise since pp layers are sharded by their order in model, not by their types. - # e.g. the upper half are all mamba layers while the lower half are all attention layers. - # But that's close enough for real world models where mamba and attention layers are interleaved - # and we don't have access with layer_masks at this point. - num_mamba_layers_per_rank = len( - mapping.pp_layers(params.num_mamba_layers)) - state_bytes_per_rank = num_mamba_layers_per_rank * state_bytes_per_layer - - # Per-request fixed cost. STATIC_SLOTS_PER_REQUEST = 1 today (the - # live mamba state); fixed-position snapshots are not yet - # implemented and would simply increment this constant. With - # pipeline parallelism, multiple microbatches are in-flight - # concurrently on the same rank, so each rank holds Mamba state - # for up to ``max_batch_size * pp_size`` concurrent sequences. - STATIC_SLOTS_PER_REQUEST = 1 - pp_size = mapping.pp_size if mapping is not None else 1 - intercept = (max_batch_size * pp_size * state_bytes_per_rank * - STATIC_SLOTS_PER_REQUEST) - - # Regular-snapshot bytes per token. None / non-positive intervals - # mean "no regular snapshots", so the mamba contribution is zero. - interval = kv_cache_config.mamba_state_cache_interval if kv_cache_config.enable_block_reuse else 0 - if interval is None or interval <= 0: - mamba_slope = 0 - else: - mamba_slope = state_bytes_per_rank // interval - # heuristic: When block reuse is enabled, we assume the mamba snapshots are dominant instead of active states, - # otherwise we may run out of kv cache blocks prior to mamba blocks due to the large number of max_batch_size. - # So we ignore intercept and only calculate max_tokens based on slope - # This can be improved by a more accurate max_batch_size and ISL/OSL estimation in the future. - if mamba_slope > 0: - intercept = 0 - return attention_slope + mamba_slope, intercept - @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: return self._use_gdn_cached_replay_all_layer_commit @@ -1738,47 +2191,7 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( self.requests) self._setup_state_indices() - # Reset replay double-buffer state for fresh context blocks. A reused - # block (prefix-cache hit or block recycled across requests) may carry - # stale prev_num_accepted_tokens / cache_buf_idx values from a prior - # owner; the replay kernel reads these on the first decode step. - num_contexts = len(scheduled_batch.context_requests) - if num_contexts > 0: - ctx_slots = self.cuda_state_indices[:num_contexts].long() - if (self._use_replay_state_update - and self.prev_num_accepted_tokens is not None - and self.cache_buf_idx is not None): - self.prev_num_accepted_tokens[ctx_slots] = 0 - self.cache_buf_idx[ctx_slots] = 0 - if self.old_x is not None: - self.old_x[:, ctx_slots] = 0 - if self.old_B is not None: - self.old_B[:, ctx_slots] = 0 - if self.old_dt is not None: - self.old_dt[:, ctx_slots] = 0 - if self.old_dA_cumsum is not None: - self.old_dA_cumsum[:, ctx_slots] = 0 - # Deterministic per-context-slot seed rotation. Runs whenever - # the seed buffer exists, including the non-replay SR path. - # Bump the host counter once per batch and write one new seed - # per fresh context slot from a pure function of - # (counter, slot, rank). No torch.randint involved. - if self.mamba_ssm_rand_seed is not None: - self._seed_request_counter += 1 - counter = self._seed_request_counter - rank_offset = self._seed_rank_offset - host_slots = ctx_slots.cpu().tolist() - new_seeds = [ - _compute_deterministic_mamba_seed(counter, slot, - rank_offset) - for slot in host_slots - ] - seed_tensor = torch.tensor( - new_seeds, - dtype=torch.int64, - device=self.mamba_ssm_rand_seed.device, - ) - self.mamba_ssm_rand_seed[ctx_slots] = seed_tensor + self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) def prepare_resources(self, scheduled_batch: ScheduledRequests): super().prepare_resources(scheduled_batch) @@ -1796,9 +2209,6 @@ def flush_state_transfers(self) -> None: self.impl.refresh_blocks() self._pending_state_transfers = False - def is_speculative(self) -> bool: - return self.spec_config is not None - @nvtx_range("hybrid_update_mamba_states") def update_mamba_states( self, @@ -1840,32 +2250,17 @@ def update_mamba_states( # in one launch before PNAT and the active history buffer change. self._commit_gdn_cached_replay_history_layers( attn_metadata, num_gens) - # SSM state is handled incrementally by the kernel. Mirror the - # kernel's checkpoint predicate from the previous PNAT and fixed - # replay step width: checkpoint steps flip buffers, while no-write - # steps append to the active history. - slots = state_indices_d.long() - accepted = num_accepted_tokens[num_contexts:num_contexts + - num_gens].to( - self.prev_num_accepted_tokens. - dtype) - prev_num_accepted_tokens = self.prev_num_accepted_tokens[slots] - wrote_checkpoint = (prev_num_accepted_tokens + - self.replay_step_width - > self.replay_history_size) - next_num_accepted_tokens = torch.where( - wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) - cache_buf_idx = self.cache_buf_idx[slots] assert self._dummy_request_mask is not None is_dummy_request = self._dummy_request_mask[ num_contexts:num_contexts + num_gens] - next_num_accepted_tokens = torch.where(is_dummy_request, - prev_num_accepted_tokens, - next_num_accepted_tokens) - self.prev_num_accepted_tokens[slots] = next_num_accepted_tokens - self.cache_buf_idx[slots] = torch.where( - is_dummy_request, cache_buf_idx, - torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx)) + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) else: # Legacy: copy the accepted SSM state from the intermediate buffer. _promote_mamba_state_triton(self.all_ssm_states, @@ -1880,20 +2275,6 @@ def update_mamba_states( src_state_indices, accepted_positions, state_indices_d) - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - if self._dummy_request_mask is None: - return - - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.tensor(is_dummy, dtype=torch.bool)) - self._dummy_request_mask.copy_(self._dummy_request_mask_host, - non_blocking=True) - def get_num_available_tokens(self, token_num_upper_bound: int, max_num_draft_tokens: int = 0, @@ -1924,70 +2305,12 @@ def get_num_available_tokens(self, result = min(result, rs_token_cap) return max(result, 0) - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_ssm_states is None: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.intermediate_ssm_states[layer_offset] - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_conv_states is None: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.intermediate_conv_states[layer_offset] - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - conv = self.get_conv_states(layer_idx) - ssm = self.get_ssm_states(layer_idx) - if self.spec_config is not None: - layer_offset = self.mamba_layer_offsets[layer_idx] - spec_kwargs = {} - # Per-cache-slot Philox seed buffer is shared across replay and - # non-replay MTP paths. The mixer asserts non-None on both - # branches when SR is enabled, so pass it through whenever it - # exists — not just on the replay branch. - if self.mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed - if self._use_replay_state_update: - # Per-layer slices for the replay kernel; shared 1D tensors - # (cache_buf_idx, prev_num_accepted_tokens) are passed - # untouched via the SpeculativeState._SHARED_FIELDS contract. - spec_kwargs['old_x'] = self.old_x[layer_offset] - spec_kwargs['old_B'] = self.old_B[layer_offset] - spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] - spec_kwargs['cache_buf_idx'] = self.cache_buf_idx - spec_kwargs['prev_num_accepted_tokens'] = ( - self.prev_num_accepted_tokens) - else: - spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ - layer_offset] - return PythonMambaCacheManager.SpeculativeState( - conv=conv, - temporal=ssm, - intermediate_conv_window=self. - intermediate_conv_states[layer_offset], - **spec_kwargs, - ) - return PythonMambaCacheManager.State(conv=conv, temporal=ssm) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - if request in self.requests: - self.requests.remove(request) - self._request_id_to_state_index.pop(request.py_request_id, None) - self._request_id_to_is_dummy.pop(request.py_request_id, None) - super().free_resources(request, pin_on_release) + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + if request in self.requests: + self.requests.remove(request) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) def _setup_state_indices(self, requests=None) -> None: if self.local_num_mamba_layers == 0: @@ -2022,6 +2345,31 @@ def _setup_state_indices(self, requests=None) -> None: values = self.host_block_offsets[self.recurrent_states_pool_index, rows, 0, bi] invalid_mask = (values < 0) | (values >= max_blocks) + # The C++ recurrent-state manager uses null page-table entries for + # logical blocks that are not snapshot boundaries. Usually a + # context chunk ends exactly at an allocated snapshot (or at the + # final live-state block), but the scheduler may shorten a chunk + # further when KV capacity is tight. In that case the Mamba + # kernel must keep accumulating into the next allocated snapshot + # or final block, matching copyLinearAttentionBlock(), which also + # skips placeholders when it advances the live state. + for bad_i in invalid_mask.nonzero( + as_tuple=False).flatten().tolist(): + req = requests[bad_i] + if req.is_context_finished: + continue + last_prompt_block = (req.prompt_len - + 1) // self.tokens_per_block + row = self.host_block_offsets[self.recurrent_states_pool_index, + rows[bad_i], 0] + candidates = row[block_indices[bad_i]:last_prompt_block + 1] + valid_candidates = ((candidates >= 0) + & (candidates < max_blocks)).nonzero( + as_tuple=False) + if valid_candidates.numel() > 0: + values[bad_i] = candidates[valid_candidates[0, 0]] + + invalid_mask = (values < 0) | (values >= max_blocks) if invalid_mask.any(): bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) req = requests[bad_i] @@ -2058,6 +2406,11 @@ def _setup_state_indices(self, requests=None) -> None: def get_state_indices(self, request_ids: Optional[List[int]] = None, is_padding: Optional[List[bool]] = None) -> list: + if self.local_num_mamba_layers == 0: + # Mamba metadata is prepared on every PP rank even when this rank + # owns only attention layers. No local kernel consumes these + # indices, so avoid consulting state that is intentionally absent. + return [0] * len(request_ids) if request_ids is not None else [] if request_ids is not None: # Return indices in the order of the caller's request_ids, # not the internal self.requests order. This is critical when @@ -2077,24 +2430,6 @@ def get_state_indices(self, return indices return self.cuda_state_indices - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - """Set reusable Mamba snapshot boundaries before scheduling.""" - if not self.kv_cache_config.enable_block_reuse: - for request in requests: - request.expect_snapshot_points = [] - return - - interval = self.linear_attention_metadata.states_snapshot_interval - if interval is None or interval <= 0: - for request in requests: - request.expect_snapshot_points = [] - return - - for request in requests: - request.expect_snapshot_points = list( - range(interval, request.prompt_len + 1, interval)) - def _setup_states(self) -> None: # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) pool: torch.Tensor = self.impl.get_recurrent_states_pool().view( @@ -2114,164 +2449,851 @@ def _setup_states(self) -> None: self.all_ssm_states.zero_() self.all_conv_states.zero_() - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size) -> None: - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if self.spec_config is not None: - # DFlash/PARD use 2K query tokens per gen, so size by tokens_per_gen_step. - speculative_num_draft_tokens = self.spec_config.tokens_per_gen_step - 1 - num_local_mamba_layers = len(self.mamba_pp_layers) - - # Legacy SSM intermediate buffer is only needed when replay is - # disabled; replay reads from the per-block double-buffered cache - # set up in _setup_replay_buffers instead. - if not self._use_replay_state_update: - self.intermediate_ssm_states = torch.zeros( - size=[ - num_local_mamba_layers, max_batch_size, - speculative_num_draft_tokens + 1 - ] + self.ssm_state_shape, - dtype=self.ssm_state_dtype, - device="cuda", - ) + def _setup_replay_buffers(self, spec_config) -> None: + cache_size = self.all_ssm_states.shape[1] + device = self.all_ssm_states.device + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): + return - self.intermediate_conv_states = torch.zeros( - size=[ - num_local_mamba_layers, max_batch_size, - speculative_num_draft_tokens + 1 - ] + self.conv_state_shape, - dtype=self.conv_state_dtype, - device="cuda", + self._dummy_request_mask = torch.zeros(self.max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros( + self.max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) + + +class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): + """Hybrid Mamba cache manager backed by KVCacheManagerV2. + + Attention KV pages and Mamba recurrent-state pages are both owned by the + Python V2 cache manager. Mamba layers are represented as V2 SSM layers, + while this wrapper exposes the state tensors and slot indices expected by + the PyTorch Mamba kernels. + """ + + _supports_additional_snapshot_offsets = True + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + layer_mask: Optional[List[bool]] = None, + is_estimating_kv_cache: bool = False, + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", + **kwargs, + ) -> None: + if conv_state_layout not in ("x_b_c", "q_k_v"): + raise ValueError( + f"Unsupported convolution state layout: {conv_state_layout!r}") + total_layers = len(mamba_layer_mask) + if layer_mask is None: + full_attention_layer_mask = [False] * total_layers + elif len(layer_mask) != total_layers: + raise ValueError( + f"layer_mask length ({len(layer_mask)}) must match " + f"mamba_layer_mask length ({total_layers})") + else: + full_attention_layer_mask = list(layer_mask) + + combined_layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] + for i in range(total_layers) + ] + + self._mamba_layer_mask = list(mamba_layer_mask) + self._use_replay_state_update = use_replay_state_update + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None else None) + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + self._seed_request_counter = 0 + num_cuda_graph_padding_dummy_slots = ( + _get_num_cuda_graph_padding_dummy_slots(spec_config, + max_batch_size)) + self._num_reserved_dummy_slots = (num_cuda_graph_padding_dummy_slots + + int(mapping.enable_attention_dp)) + self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype + is not None else mamba_cache_dtype) + self.conv_state_dtype = mamba_cache_dtype + + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + self.mamba_pp_layers = [ + layer_idx for layer_idx in self.pp_layers + if mamba_layer_mask[layer_idx] + ] + self.local_num_mamba_layers = len(self.mamba_pp_layers) + + if self.local_num_mamba_layers > 0: + tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 + d_inner = mamba_head_dim * mamba_num_heads + grouped_state_dim = mamba_n_groups * mamba_d_state + conv_dim = d_inner + 2 * grouped_state_dim + nheads = mamba_num_heads + assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if kwargs.get("is_disagg", + False) and grouped_state_dim % tp_size != 0: + raise ValueError( + "Disaggregated Mamba transfer requires each convolution " + "state section to be divisible by tp_size") + if use_replay_state_update: + assert mamba_n_groups % tp_size == 0, \ + "replay state update requires mamba_n_groups divisible by tp_size" + self._n_groups_per_rank = mamba_n_groups // tp_size + d_inner_local = d_inner // tp_size + grouped_state_dim_local = grouped_state_dim // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + # TP-mismatch disaggregated transfers must split the flat + # convolution state at its true semantic boundaries. Mamba2 stores + # [x | B | C], while GDN stores [Q | K | V]. The large section is + # therefore first for Mamba2 and last for GDN. + if conv_state_layout == "x_b_c": + self.conv_section_dims = [ + d_inner_local, + grouped_state_dim_local, + grouped_state_dim_local, + ] + else: + self.conv_section_dims = [ + grouped_state_dim_local, + grouped_state_dim_local, + d_inner_local, + ] + self.ssm_bytes = (math.prod(self.ssm_state_shape) * + self.ssm_state_dtype.itemsize) + self.conv_bytes = (math.prod(self.conv_state_shape) * + self.conv_state_dtype.itemsize) + else: + logger.info( + "No local mamba layers for this rank, skipping mamba state views" ) + self._n_groups_per_rank = 0 + self.conv_state_shape = [] + self.ssm_state_shape = [] + self.conv_section_dims = [] + self.ssm_bytes = 0 + self.conv_bytes = 0 - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") + if isinstance(num_kv_heads, int): + per_layer_kv_heads = [num_kv_heads] * total_layers + else: + if len(num_kv_heads) != total_layers: + raise ValueError( + f"num_kv_heads list length ({len(num_kv_heads)}) does not " + f"match total layers ({total_layers})") + per_layer_kv_heads = list(num_kv_heads) + for i, is_mamba in enumerate(mamba_layer_mask): + if is_mamba: + per_layer_kv_heads[i] = 0 - def _setup_replay_buffers(self, spec_config) -> None: - """Allocate per-pool-block replay buffers used by replay_selective_state_update. + self._setup_mtp_intermediate_states(spec_config, max_batch_size) - Unlike the Mixed cache manager (where slots are 0..max_batch_size-1), - the unified C++ KV pool assigns recurrent-state block indices up to - ``num_blocks_in_pool``. The replay kernel indexes ``cache_buf_idx`` and - ``prev_num_accepted_tokens`` by these block indices, so the buffers - must match the pool extent rather than ``max_batch_size``. - """ - # Replay tensors require spec_config + replay path enabled. The - # rand_seed buffer is separable from replay and must also be - # allocated for non-replay SR so the flashinfer path has a - # persistent deterministic seed source. - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None + kv_cache_config = kv_cache_config.model_copy(deep=True) + if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: + block_reuse_policy = BlockReusePolicy( + kv_cache_config.block_reuse_policy) + if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: + # SSM reuse is valid only at explicit snapshot boundaries. + kv_cache_config.block_reuse_policy = ( + BlockReusePolicy.PER_REQUEST.value) + self.kv_cache_config = kv_cache_config - if (not self._use_replay_state_update - and not self._mamba_ssm_stochastic_rounding): - return + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=mamba_num_layers + num_layers, + num_kv_heads=per_layer_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=combined_layer_mask, + is_draft=is_draft, + is_estimating_kv_cache=is_estimating_kv_cache, + num_reserved_index_slots=self._num_reserved_dummy_slots, + **kwargs, + ) - cache_size = self.all_ssm_states.shape[1] - device = self.all_ssm_states.device - # Always-available deterministic seed buffer when SR (or replay) - # is on. Works for non-MTP runs because we don't depend on - # spec_config to allocate it. - self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - cache_size, self._seed_rank_offset, device) + self.mamba_layer_offsets = { + layer_id: offset + for offset, layer_id in enumerate(self.mamba_pp_layers) + } + self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} - if spec_config is None or not self._use_replay_state_update: - # Without spec_config or replay we still keep the seed buffer - # (above) so the non-MTP flashinfer SR path has a persistent - # rand_seed source. - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None - return + state_index_capacity = (self.max_batch_size + + self._num_reserved_dummy_slots) + self.cuda_state_indices = torch.zeros([state_index_capacity], + dtype=torch.int32, + device="cuda") + self._host_state_indices = torch.zeros([state_index_capacity], + dtype=torch.int32, + pin_memory=prefer_pinned()) - history_size = self.replay_history_size - num_local_mamba_layers = self.local_num_mamba_layers - nheads, head_dim, d_state = self.ssm_state_shape - n_groups_per_rank = self._n_groups_per_rank + if self.local_num_mamba_layers > 0: + first_mamba_local_layer = self.layer_offsets[ + self.mamba_pp_layers[0]] + self.ssm_layer_group_id = self.impl.get_layer_group_id( + LayerId(first_mamba_local_layer)) + self._ssm_page_index_scale = self.impl.get_page_index_scale( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_pages = self.impl.get_page_index_upper_bound( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) + num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // + self._ssm_page_index_scale) + required_live_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + if num_ssm_slots < required_live_slots: + KVCacheManagerV2.shutdown(self) + raise ValueError( + "The V2 Mamba state pool has only " + f"{num_ssm_slots} slots but needs at least " + f"{required_live_slots} live/dummy slots. Increase the " + "KV cache budget or allocate a larger Mamba pool_ratio.") + self._setup_states() + self._setup_replay_buffers(spec_config) + else: + self.ssm_layer_group_id = None + self._ssm_page_index_scale = 1 + self.all_ssm_states = [] + self.all_conv_states = [] + self._setup_replay_buffers(spec_config) - # Shared across layers (consumed by the replay kernel via slot index). - self.prev_num_accepted_tokens = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.cache_buf_idx = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self._dummy_request_mask = torch.zeros(self.max_batch_size, + @staticmethod + def get_cache_size_per_token(model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs): + spec_config = kwargs.get("spec_config") + num_reserved_dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( + spec_config, max_batch_size) + int(mapping.enable_attention_dp)) + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=num_reserved_dummy_slots, + include_explicit_snapshots=True, + cap_partial_attention_snapshots=True, + **kwargs, + ) + + def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: + return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] + + def _get_pool_roles(self, + pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_local_mamba_layer(layer_id): + return MambaRole.SSM_STATE, None + return super()._get_pool_roles(pool_id) + + def _max_resident_sequences(self) -> int: + return self.max_batch_size * self.mapping.pp_size + + def _mamba_state_bytes_per_slot(self) -> int: + return self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) + + def _num_ssm_snapshots_for_capacity( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + if capacity <= 0 or not kv_cache_config.enable_block_reuse: + return 0 + + fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, + self.max_seq_len, + self.tokens_per_block) + interval = _mamba_regular_snapshot_interval(kv_cache_config, + self.max_seq_len) + regular_snapshots = capacity // interval if interval is not None else 0 + return (self._max_resident_sequences() * fixed_rules + + regular_snapshots) + + def _num_ssm_states_per_typical_request( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + fixed_rules, _ = _mamba_snapshot_rule_counts( + kv_cache_config, + capacity, + self.tokens_per_block, + ) + # Additional snapshots are stable boundaries that must remain alive. + # Periodic snapshots are evictable cache entries and therefore do not + # increase the guaranteed state count represented by BatchDesc. + return 1 + fixed_rules + + def _typical_request_descs( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> List[KVCacheDesc]: + """Model one request with one descriptor per live SSM state.""" + num_states = self._num_ssm_states_per_typical_request( + capacity, kv_cache_config) + capacity_per_state, capacity_remainder = divmod(capacity, num_states) + capacities = [ + capacity_per_state + int(i < capacity_remainder) + for i in range(num_states) + ] + return [ + KVCacheDesc( + capacity=state_capacity, + history_length=max(0, state_capacity - 1), + ) for state_capacity in capacities + ] + + def _get_typical_request_capacity( + self, + kv_cache_config: KvCacheConfig, + ) -> int: + if kv_cache_config.avg_seq_len is not None: + return kv_cache_config.avg_seq_len + + fallback_capacity = max(1, self.max_seq_len // 2) + logger.warning( + "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " + "model using KV cache manager V2. Falling back to " + f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " + "'kv_cache_config.avg_seq_len' in the YAML configuration to the " + "workload's average total sequence length for an accurate KV/SSM " + "pool ratio.") + return fallback_capacity + + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + attention_quota = super()._get_quota_from_max_tokens(max_tokens) + num_request_lineages = self._max_resident_sequences() + snapshot_slots = self._num_ssm_snapshots_for_capacity( + max_tokens, self.kv_cache_config) + state_slots = (num_request_lineages + self._num_reserved_dummy_slots + + snapshot_slots) + state_quota = state_slots * self._mamba_state_bytes_per_slot() + # Once the plan contains any non-live SSM capacity, reserve one partial + # attention page per request lineage. This remains conservative when + # the plan contains fewer than one non-live slot per lineage. + extra_attention_quota = (num_request_lineages * + self._attention_cache_bytes_per_token() * + self.tokens_per_block + if snapshot_slots > 0 else 0) + return attention_quota + state_quota + extra_attention_quota + + def _get_max_tokens_from_quota(self, quota: int) -> float: + if self._get_quota_from_max_tokens(0) > quota: + return 0 + + low = 0 + high = 1 + while self._get_quota_from_max_tokens(high) <= quota: + low = high + high *= 2 + if high >= 1 << 62: + return float("inf") + + while low + 1 < high: + mid = (low + high) // 2 + if self._get_quota_from_max_tokens(mid) <= quota: + low = mid + else: + high = mid + return low + + def _minimum_live_gpu_quota(self) -> int: + """Return the minimum quota for live states and one attention page.""" + attention_block_quota = (self._attention_cache_bytes_per_token() * + self.tokens_per_block) + num_state_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + return max( + self._get_quota_from_max_tokens(0), + state_quota + attention_block_quota, + ) + + def _build_cache_config( + self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + kv_cache_config = self.kv_cache_config + cache_tiers = config.cache_tiers + gpu_quota = cache_tiers[0].quota + minimum_live_quota = self._minimum_live_gpu_quota() + if minimum_live_quota > gpu_quota: + raise ValueError( + "The V2 Mamba GPU cache quota is too small for live recurrent " + f"states and attention pages: got {gpu_quota} bytes, need at " + f"least {minimum_live_quota} bytes.") + # _build_base_config already constructed every attention layer, + # including dtype-specific scale and subclass-provided side buffers. + # Preserve those configs and replace only the local Mamba layers. + layers = list(config.layers) + for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): + if self._mamba_layer_mask[global_layer_idx]: + layer_id = LayerId(local_layer_idx) + layers[local_layer_idx] = SsmLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig(role=MambaRole.SSM_STATE, + size=self.ssm_bytes), + BufferConfig(role=MambaRole.CONV_STATE, + size=self.conv_bytes), + ], + ) + + dummy_requests = [ + KVCacheDesc(capacity=0, history_length=0) + for _ in range(self._num_reserved_dummy_slots) + ] + constraints = [ + replace( + batch, + kv_caches=[*batch.kv_caches, *dummy_requests], + ) for batch in config.constraints + ] + + typical_step = config.typical_step + if config.initial_pool_ratio is None: + typical_capacity = self._get_typical_request_capacity( + kv_cache_config) + request_descs = self._typical_request_descs(typical_capacity, + kv_cache_config) + typical_step = BatchDesc(request_descs * + self._max_resident_sequences() + + dummy_requests) + return replace( + config, + layers=layers, + typical_step=typical_step, + constraints=constraints, + # SSM lifecycles require minimum-snapshot commit semantics. The + # flag is harmless when reuse is disabled because no commits are + # attempted, while the runtime config still needs the invariant. + commit_min_snapshot=True, + ) + + def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, + state_shape: List[int]) -> torch.Tensor: + addr = self.impl.get_mem_pool_base_address(LayerId(local_layer_idx), + role, PageIndexMode.SHARED) + num_pages = self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), role) + raw = convert_to_torch_tensor( + TensorWrapper(addr, dtype, [num_pages] + state_shape)) + page_index_scale = self.impl.get_page_index_scale( + LayerId(local_layer_idx), role) + num_slots = (num_pages + page_index_scale - 1) // page_index_scale + # V2 coalesces same-size per-layer buffers inside each slot. Kernels + # index Mamba states by logical slot id, so expose only this layer's + # sub-page from each coalesced slot instead of the raw page-index view. + return raw.as_strided( + [num_slots] + state_shape, + [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), + ) + + def _setup_states(self) -> None: + local_layer_ids = [ + self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers + ] + self.all_ssm_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, + self.ssm_state_dtype, self.ssm_state_shape) + for local_layer_idx in local_layer_ids + ] + self.all_conv_states = [ + self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, + self.conv_state_dtype, self.conv_state_shape) + for local_layer_idx in local_layer_ids + ] + + def _setup_replay_buffers(self, spec_config) -> None: + cache_size = 0 + device = None + if self.local_num_mamba_layers > 0: + cache_size = self.all_ssm_states[0].shape[0] + assert all(t.shape[0] == cache_size for t in self.all_ssm_states) + device = self.all_ssm_states[0].device + if not self._allocate_pool_replay_buffers(spec_config, cache_size, + device): + return + + mask_capacity = self._host_state_indices.shape[0] + self._dummy_request_mask = torch.zeros(mask_capacity, dtype=torch.bool, device=device) - self._dummy_request_mask_host = torch.zeros(self.max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned()) - self.old_x = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - history_size, - nheads, - head_dim, - dtype=self.conv_state_dtype, - device=device) - # Per-layer double-buffered caches. - self.old_B = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - history_size, - n_groups_per_rank, - d_state, - dtype=self.conv_state_dtype, - device=device) - self.old_dt = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - nheads, - history_size, - dtype=torch.float32, - device=device) - self.old_dA_cumsum = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - nheads, - history_size, - dtype=torch.float32, - device=device) + self._dummy_request_mask_host = torch.zeros( + mask_capacity, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None + def _attention_cache_bytes_per_token(self) -> int: + # Mamba layers have zero KV heads, so the generic calculation naturally + # returns only bytes owned by local attention layers. + return super().get_cache_bytes_per_token() + + def get_cache_bytes_per_token(self) -> int: + cache_bytes = self._attention_cache_bytes_per_token() + + interval = ( + self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) + if (self.kv_cache_config.enable_block_reuse and interval is not None + and interval > 0): + cache_bytes += self._mamba_state_bytes_per_slot() // interval + if cache_bytes == 0 and self.local_num_mamba_layers > 0: + cache_bytes = self._mamba_state_bytes_per_slot() + return max(1, cache_bytes) + + def get_num_free_blocks(self) -> int: + assert len(self.kv_cache_map) == 0, ( + "get_num_free_blocks is only used when the kv cache manager is empty" + ) + attention_pages = [] + ssm_pages = [] + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + if self._is_local_mamba_layer(local_layer_idx): + ssm_pages.append( + self.impl.get_page_index_upper_bound( + layer_id, MambaRole.SSM_STATE) // + self._ssm_page_index_scale) + else: + attention_pages.append( + self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // + self.kv_factor) + if attention_pages: + return max(attention_pages) + return max(ssm_pages) if ssm_pages else 0 - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', - None) - cache_buf_idx = getattr(self, 'cache_buf_idx', None) - if (not self._use_replay_state_update - or prev_num_accepted_tokens is None or cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): + @property + def blocks_in_primary_pool(self) -> int: + for local_layer_idx in range(self.num_local_layers): + if self._is_local_mamba_layer(local_layer_idx): + continue + return self.impl.get_page_index_upper_bound( + LayerId(local_layer_idx), Role.KEY) + return 0 + + def get_buffers(self, + layer_idx: int, + kv_layout: str = "NHD") -> Optional[torch.Tensor]: + local_layer_idx = self.layer_offsets[layer_idx] + if self._is_local_mamba_layer(local_layer_idx): return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=prev_num_accepted_tokens, - cache_buf_idx=cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) + return super().get_buffers(layer_idx, kv_layout) - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.ssm_state_dtype + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: + for global_layer_id, local_layer_id in self.layer_offsets.items(): + if self._is_local_mamba_layer(local_layer_id): + continue + # A layer group is a lifecycle, not a physical memory pool. + # Differently sized attention buffers can share one lifecycle, + # so scan every attention layer in this diagnostic path. + yield KVCacheManagerV2.get_buffers(self, global_layer_id) - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - """Return the persistent (cache_size,) int64 Philox seed buffer or - None when stochastic rounding is not active for this manager.""" - return getattr(self, 'mamba_ssm_rand_seed', None) + yield from self.all_ssm_states + yield from self.all_conv_states + + def add_dummy_requests( + self, + request_ids: List[int], + token_nums: Optional[List[int]] = None, + is_gen: bool = False, + prepare_resource: bool = True, + max_num_draft_tokens: int = 0, + kv_reserve_draft_tokens: Optional[int] = None, + use_mrope: bool = False, + max_beam_width: int = 1, + encoder_output_lens: Optional[List[int]] = None, + num_extra_decoding_steps: int = 0, + draft_kv_cache_manager: Optional[BaseResourceManager] = None, + ) -> List[LlmRequest]: + requests = super().add_dummy_requests( + request_ids=request_ids, + token_nums=token_nums, + is_gen=is_gen, + prepare_resource=prepare_resource, + max_num_draft_tokens=max_num_draft_tokens, + kv_reserve_draft_tokens=kv_reserve_draft_tokens, + use_mrope=use_mrope, + max_beam_width=max_beam_width, + encoder_output_lens=encoder_output_lens, + num_extra_decoding_steps=num_extra_decoding_steps, + draft_kv_cache_manager=draft_kv_cache_manager, + ) + if requests and prepare_resource: + self._setup_state_indices(requests) + return requests + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is not None and kv_cache.is_active: + self.try_commit_blocks(request, kv_cache) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + self._setup_state_indices(requests) + num_contexts = len(scheduled_batch.context_requests) + self._reset_context_mamba_slots(num_contexts) + + def _setup_state_indices(self, requests: List[LlmRequest]) -> None: + if self.local_num_mamba_layers == 0: + return + n = len(requests) + assert n <= self._host_state_indices.shape[0], ( + f"State-index batch size {n} exceeds max_batch_size " + f"{self._host_state_indices.shape[0]}") + self._host_state_indices.zero_() + if n > 0: + for i, req in enumerate(requests): + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + raise RuntimeError( + f"Missing V2 KV cache for request {req.py_request_id}") + base_index = kv_cache.get_ssm_block_base_index( + self.ssm_layer_group_id) + if base_index < 0: + raise RuntimeError( + f"Invalid SSM state block index {base_index} for " + f"request {req.py_request_id}") + self._host_state_indices[i] = base_index + + self.cuda_state_indices.copy_(self._host_state_indices, + non_blocking=True) + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask(is_dummy) + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy + + def get_state_indices(self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None): + if self.local_num_mamba_layers == 0: + # Mamba metadata is still prepared on attention-only PP ranks, + # but no local kernel consumes these indices. Return harmless + # placeholders instead of consulting an intentionally empty map. + if request_ids is not None: + return [0] * len(request_ids) + return self.cuda_state_indices + if request_ids is not None: + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + return self.cuda_state_indices + + def get_max_resource_count(self) -> int: + return self.max_batch_size + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None): + if self.local_num_mamba_layers == 0: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + num_accepted_draft_tokens = ( + num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( + torch.int32) + # Dynamic tree selects a tree node rather than a linear draft depth. + accepted_positions = (accepted_leaf_positions.to(torch.int32) + if accepted_leaf_positions is not None else + num_accepted_draft_tokens) + if state_indices is None: + state_indices = self.get_state_indices() + state_indices_d = state_indices[num_contexts:num_contexts + + num_gens].to(torch.int32) + src_state_indices = self.intermediate_state_indices[:num_gens] + + if self._use_replay_state_update: + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts:num_contexts + num_gens], + is_dummy_request, + ) + else: + for layer_offset, dst in enumerate(self.all_ssm_states): + _promote_mamba_state_triton( + dst.unsqueeze(0), + self.intermediate_ssm_states[layer_offset:layer_offset + 1], + src_state_indices, + accepted_positions, + state_indices_d, + ) + + for layer_offset, dst in enumerate(self.all_conv_states): + _promote_mamba_state_triton( + dst.unsqueeze(0), + self.intermediate_conv_states[layer_offset:layer_offset + 1], + src_state_indices, + accepted_positions, + state_indices_d, + ) + + def _mark_context_position_as_history(self, request: LlmRequest, + kv_cache) -> None: + """Advance history without making later recurrent state reusable.""" + history_length = request.context_current_position + if history_length <= kv_cache.history_length: + return + capacity = max(kv_cache.capacity, history_length) + if not kv_cache.resize(capacity, history_length=history_length): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to {history_length} tokens") + + def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + if not should_block_reuse: + return + + if kv_cache is None: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return + + snapshot_points = request.expect_snapshot_points + commit_limit = (min(max(snapshot_points), request.prompt_len) + if snapshot_points else request.prompt_len) + commit_end = min(request.context_current_position, commit_limit) + if (request.context_current_position in request.expect_snapshot_points + and commit_end > kv_cache.num_committed_tokens): + tokens = self._augment_tokens_for_block_reuse( + request.get_tokens(DEFAULT_BEAM_INDEX), + request, + start=kv_cache.num_committed_tokens, + end=commit_end, + ) + kv_cache.commit(tokens) + if request.context_current_position >= commit_limit: + self._mark_context_position_as_history(request, kv_cache) + if request.context_remaining_length == 0: + kv_cache.stop_committing() + + def update_context_resources(self, + scheduled_batch: ScheduledRequests) -> None: + for request in scheduled_batch.context_requests: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None or not kv_cache.is_active: + continue + + should_block_reuse = (self.enable_block_reuse and not self.is_draft + and not request.is_dummy_request) + is_all_reusable = ( + self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE) + is_snapshot_boundary = (request.context_current_position + in request.expect_snapshot_points) + has_pending_snapshot = any( + point > request.context_current_position + for point in request.expect_snapshot_points) + should_resize = (not should_block_reuse or + (not is_all_reusable and not has_pending_snapshot)) + should_commit = (is_all_reusable or is_snapshot_boundary + or request.context_remaining_length == 0) + + if should_resize and not kv_cache.resize( + None, request.context_current_position): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to " + f"{request.context_current_position} tokens at context " + "update") + if should_commit: + self.try_commit_blocks(request, kv_cache) + if request.context_remaining_length == 0: + if self.conversation_manager is not None: + self.conversation_manager.save_drop_plan(request, kv_cache) + kv_cache.enable_swa_scratch_reuse = False + + def shutdown(self): + self.all_ssm_states = [] + self.all_conv_states = [] + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + super().shutdown() diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index f2284e06bb46..56e2e71d265c 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -42,7 +42,8 @@ MoeLoadBalancer, maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode from ..virtual_memory import scope as virtual_memory_scope -from .config_utils import resolve_hf_torch_dtype, resolve_ssm_cache_dtype +from .config_utils import (is_hybrid_linear, resolve_hf_torch_dtype, + resolve_ssm_cache_dtype) _KV_CACHE_MAP = { "fp8": QuantAlgo.FP8.value, @@ -52,6 +53,36 @@ _VALID_KV_CACHE_DTYPES = ("fp8", "nvfp4", "auto") +def _validate_and_adjust_mamba_snapshot_config(config: ModelConfig, + llm_args: TorchLlmArgs) -> None: + """Validate snapshot reuse after the model and V2 setting are resolved.""" + if not is_hybrid_linear(config.pretrained_config): + return + + kv_cache_config = llm_args.kv_cache_config + state_config = kv_cache_config.mamba_state_config + has_additional_snapshots = bool( + state_config.additional_snapshot_offsets_from_start + or state_config.additional_snapshot_offsets_from_end) + if (has_additional_snapshots + and kv_cache_config.use_kv_cache_manager_v2 is not True): + raise ValueError( + "Mamba additional snapshot offsets require " + "kv_cache_config.use_kv_cache_manager_v2=True after resolving " + "the model configuration.") + + has_periodic_snapshots = state_config.periodic_snapshot_interval > 0 + if (kv_cache_config.enable_block_reuse and not has_periodic_snapshots + and not has_additional_snapshots): + logger.warning( + "Disabling KV cache block reuse for the hybrid Mamba model " + "because no Mamba state snapshot policy is configured. Set " + "kv_cache_config.mamba_state_config.periodic_snapshot_interval " + "to a positive value or provide additional snapshot offsets to " + "enable block reuse.") + kv_cache_config.enable_block_reuse = False + + def validate_and_set_mamba_ssm_cache_dtype( config: ModelConfig, mamba_ssm_cache_dtype: str, @@ -391,6 +422,8 @@ def load_config_and_apply_defaults( config.pretrained_config) model_cls = AutoModelForCausalLM._resolve_class(config) + use_kv_cache_manager_v2 = ( + llm_args.kv_cache_config.use_kv_cache_manager_v2) # model_cls is None when the architecture is unknown/unsupported. model_defaults = {} @@ -404,15 +437,6 @@ def load_config_and_apply_defaults( f"Applied model defaults for {model_cls.__name__}: {applied_defaults}" ) - use_kv_cache_manager_v2 = llm_args.kv_cache_config.use_kv_cache_manager_v2 - _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) - if use_kv_cache_manager_v2 == "auto": - logger.info( - "Resolved use_kv_cache_manager_v2='auto' to %s for %s", - llm_args.kv_cache_config.use_kv_cache_manager_v2, - model_cls.__name__ - if model_cls is not None else "unknown model") - # The transceiver preference follows the checkpoint's original # architecture: _resolve_class may rewrite it to an execution class # (e.g. MTPDraftModelForCausalLM), which must not drop the target @@ -426,6 +450,15 @@ def load_config_and_apply_defaults( # Resolve "auto" sentinel values after model defaults are applied. _resolve_transceiver_runtime_auto(llm_args, preference_cls, config.pretrained_config) + _resolve_kv_cache_manager_v2_auto( + llm_args, model_defaults, original_setting=use_kv_cache_manager_v2) + _validate_and_adjust_mamba_snapshot_config(config, llm_args) + if use_kv_cache_manager_v2 == "auto": + logger.info( + "Resolved use_kv_cache_manager_v2='auto' to %s for %s", + llm_args.kv_cache_config.use_kv_cache_manager_v2, + model_cls.__name__ + if model_cls is not None else "unknown model") return llm_args diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7920c8942e71..7c8384696c47 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -687,16 +687,6 @@ def drafting_loop_wrapper(model): config = model_engine.model.model_config.pretrained_config max_num_seq_slots = getattr(model_engine, "max_num_seq_slots", max_batch_size * getattr(mapping, "pp_size", 1)) - if is_hybrid_linear(config) and kv_cache_config.enable_block_reuse and ( - cache_transceiver_config is not None - and cache_transceiver_config.backend is not None - and cache_transceiver_config.transceiver_runtime == "PYTHON"): - logger.warning( - "Disabling block reuse for MambaHybridCacheManager-based models when disagg + Python transceiver enabled" - ) - kv_cache_config.enable_block_reuse = False - _set_model_engines_cache_reuse([model_engine, draft_model_engine], - False) if is_mla(config): if model_engine.model.model_config.enable_flash_mla: tokens_per_block = 64 @@ -777,8 +767,9 @@ def drafting_loop_wrapper(model): ctx_chunk_config = None if kv_cache_config.enable_block_reuse and is_hybrid_linear(config): - ctx_chunk_config = (ContextChunkingPolicy.FORCE_CHUNK, - kv_cache_config.mamba_state_cache_interval) + # Snapshot boundaries come from expect_snapshot_points. The unit is + # only used to align chunks shortened by the scheduling budget. + ctx_chunk_config = (ContextChunkingPolicy.FORCE_CHUNK, tokens_per_block) guided_decoder: Optional[GuidedDecoder] = None if guided_decoding_config is not None: @@ -913,17 +904,16 @@ def drafting_loop_wrapper(model): if is_disagg and is_hybrid: # NOTE: TRTLLM_USE_PY_MAMBA is an agg-mode-only override and has - # no effect in disagg. The disagg manager choice is driven solely - # by transceiver_runtime: PYTHON => PythonMambaCacheManager, - # otherwise CppMambaHybridCacheManager (unified pool, default). + # no effect in disagg. The disagg manager choice is driven by + # get_kv_cache_manager_cls and cache_transceiver_config. if os.environ.get("TRTLLM_USE_PY_MAMBA", "0") == "1": logger.warning( "TRTLLM_USE_PY_MAMBA is ignored in disaggregated serving; " - "use cache_transceiver_config.transceiver_runtime='PYTHON' " - "to select PythonMambaCacheManager.") + "configure transceiver_runtime='PYTHON' with backend='NIXL' " + "to select MixedMambaHybridCacheManager.") else: logger.info("Disaggregated serving with hybrid model detected. " - "Using CppMambaHybridCacheManager.") + "Using the configured Mamba cache manager.") # Get draft config for one-engine speculative decoding if available draft_config = getattr(model_engine.model, 'draft_config', None) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index e42c8e249dc6..3a23a3b33e52 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -467,10 +467,13 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int], pp_size = self.mapping.pp_size if self.mapping is not None else 1 live_state_slots = self.max_batch_size * pp_size max_snapshots = live_state_slots - if kv_cache_config.enable_block_reuse: - max_snapshots += ( - kv_cache_config.max_tokens // - linear_attention_metadata.states_snapshot_interval) + snapshot_interval = ( + linear_attention_metadata.states_snapshot_interval) + if (kv_cache_config.enable_block_reuse + and snapshot_interval is not None + and snapshot_interval > 0): + max_snapshots += (kv_cache_config.max_tokens // + snapshot_interval) blocks_per_window[LinearCacheType.RECURRENT_STATES.value] = ( int(max_snapshots), 0) @@ -1932,12 +1935,23 @@ def _calculate_max_num_blocks_for_linear_attention( pp_size = self.mapping.pp_size if self.mapping is not None else 1 intercept = self.max_batch_size * pp_size * state_bytes_local - max_tokens = max((primary_budget - intercept) // slope, 0) + if slope > 0: + max_tokens = max((primary_budget - intercept) // slope, 0) + elif primary_budget >= intercept: + # With snapshots disabled, a rank containing only recurrent-state + # layers has no per-token cache cost after its live slots are + # allocated. Bound the otherwise-unlimited token count by the + # configured capacity or by all resident sequences at max length. + max_tokens = (kv_cache_config.max_tokens + if kv_cache_config.max_tokens is not None else + self.max_seq_len * self.max_batch_size * pp_size) + else: + max_tokens = 0 if kv_cache_config.max_tokens is not None: max_tokens = min(kv_cache_config.max_tokens, max_tokens) if max_tokens < kv_cache_config.max_tokens: logger.warning( - f'The memory budget for Mamba + KV cache cannot fit the user-specified max_tokens of {kv_cache_config.max_tokens}. The calculated max_tokens based on the memory budget is {max_tokens}. Please consider adjusting max_batch_size/max_tokens/mamba_state_cache_interval.' + f'The memory budget for Mamba + KV cache cannot fit the user-specified max_tokens of {kv_cache_config.max_tokens}. The calculated max_tokens based on the memory budget is {max_tokens}. Please consider adjusting max_batch_size/max_tokens/mamba_state_config.periodic_snapshot_interval.' ) kv_blocks_in_primary_pool = int(max_tokens // self.tokens_per_block) diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py index 053d23362f98..20791f971f81 100644 --- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py +++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py @@ -624,7 +624,11 @@ def _relocate_kv_eagerly(self, attn_metadata, batch_size): attn_num_heads, self._kv_head_dim_bytes, cache_mgr.max_total_draft_tokens, - cache_mgr.max_attention_window_vec[0], + # Dynamic-tree MTP currently supports full-attention KV layers. + # Hybrid managers may store a recurrent-state sentinel first, + # while V2 represents full attention as None, so neither form is + # suitable for the integer maxKVCacheLen operator argument. + cache_mgr.max_seq_len, pool_pointers, block_offsets, cache_mgr.max_blocks_per_seq, diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index ffc560b2d6b6..2dfe1e91ffe3 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -1386,7 +1386,11 @@ def _serve_llm(): llm_args_extra_dict = {} if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: - llm_args_extra_dict = yaml.safe_load(f) or {} + llm_args_extra_dict = yaml.safe_load(f) + if llm_args_extra_dict is None: + llm_args_extra_dict = {} + elif not isinstance(llm_args_extra_dict, dict): + raise ValueError("Configuration file root must be a mapping.") extra_allow_request_chat_template = _pop_bool_config_option( llm_args_extra_dict, "allow_request_chat_template") allow_request_chat_template = (allow_request_chat_template @@ -1621,7 +1625,11 @@ def serve_encoder(model: str, host: str, port: int, log_level: str, encoder_args_extra_dict = {} if extra_encoder_options is not None: with open(extra_encoder_options, 'r') as f: - encoder_args_extra_dict = yaml.safe_load(f) or {} + encoder_args_extra_dict = yaml.safe_load(f) + if encoder_args_extra_dict is None: + encoder_args_extra_dict = {} + elif not isinstance(encoder_args_extra_dict, dict): + raise ValueError("Configuration file root must be a mapping.") extra_allow_request_chat_template = _pop_bool_config_option( encoder_args_extra_dict, "allow_request_chat_template") allow_request_chat_template = (allow_request_chat_template @@ -1739,6 +1747,10 @@ def serve_embedding( if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: extra_dict = yaml.safe_load(f) + if extra_dict is None: + extra_dict = {} + elif not isinstance(extra_dict, dict): + raise ValueError("Configuration file root must be a mapping.") llm_args = update_llm_args_with_extra_dict( llm_args, extra_dict, explicit_cli_keys=explicit_cli_keys) diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 8aa94fcb114a..1bd895dbd59b 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/tensorrt_llm/llmapi/__init__.py @@ -16,9 +16,9 @@ DynamicBatchConfig, Eagle3DecodingConfig, EagleDecodingConfig, EncodeCudaGraphConfig, ExtendedRuntimePerfKnobConfig, KvCacheConfig, LlmArgs, - LookaheadDecodingConfig, MedusaDecodingConfig, - MiniMaxM3SparseAttentionConfig, MoeConfig, - MTPDecodingConfig, NGramDecodingConfig, + LookaheadDecodingConfig, MambaStateConfig, + MedusaDecodingConfig, MiniMaxM3SparseAttentionConfig, + MoeConfig, MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, PrometheusMetricsConfig, ReorderRequestPolicyConfig, RocketSparseAttentionConfig, SADecodingConfig, SAEnhancerConfig, @@ -43,6 +43,7 @@ 'ConversationParams', 'DisaggScheduleStyle', 'KvCacheConfig', + 'MambaStateConfig', 'KvCacheRetentionConfig', 'CudaGraphConfig', 'DecodeCudaGraphConfig', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index d09e0b576b5b..394753181235 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -34,7 +34,7 @@ from pydantic import AliasChoices, BaseModel, ConfigDict from pydantic import Field as PydanticField from pydantic import (NonNegativeFloat, NonNegativeInt, PositiveInt, - PrivateAttr, field_validator, model_validator) + PrivateAttr, StrictInt, field_validator, model_validator) from strenum import StrEnum from transformers import PreTrainedTokenizerBase @@ -58,9 +58,7 @@ CapacitySchedulerPolicy as _CapacitySchedulerPolicy, ContextChunkingPolicy as _ContextChunkingPolicy, DecodingConfig, - DecodingMode, DynamicBatchConfig as _DynamicBatchConfig, - EagleConfig as _EagleConfig, ExecutorConfig as _ExecutorConfig, ExtendedRuntimePerfKnobConfig as _ExtendedRuntimePerfKnobConfig, KvCacheConfig as _KvCacheConfig, @@ -3583,6 +3581,42 @@ class ReorderRequestPolicyConfig(StrictBaseModel): description="The arguments of the request reordering policy.") +_StrictPositiveInt = Annotated[StrictInt, PydanticField(gt=0)] +_StrictNonNegativeInt = Annotated[StrictInt, PydanticField(ge=0)] + + +class MambaStateConfig(StrictBaseModel): + """Configuration for reusable Mamba recurrent-state snapshots.""" + + periodic_snapshot_interval: NonNegativeInt = Field( + default=0, + status="prototype", + telemetry=True, + description= + "The number of tokens between periodic snapshots in the Mamba " + "prefix cache. Periodic snapshots are disabled by default; set this " + "to a positive value to enable them.") + + additional_snapshot_offsets_from_start: List[_StrictPositiveInt] = Field( + default_factory=list, + status="prototype", + telemetry=False, + description= + "Additional Mamba state snapshot offsets measured from the start " + "of each prompt. Offsets beyond the prompt length are ignored. " + "These snapshots require KV cache manager V2.") + + additional_snapshot_offsets_from_end: List[_StrictNonNegativeInt] = Field( + default_factory=list, + status="prototype", + telemetry=False, + description= + "Additional Mamba state snapshot offsets measured backward from " + "the end of each prompt. An offset of 0 selects the prompt end. " + "Offsets that do not resolve inside the prompt are ignored. These " + "snapshots require KV cache manager V2.") + + @PybindMirror.mirror_pybind_fields(_KvCacheConfig) class KvCacheConfig(StrictBaseModel, PybindMirror): """Configuration for the KV cache.""" @@ -3716,10 +3750,18 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description="The number of tokens per block.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. - mamba_state_cache_interval: PositiveInt = Field( - default=256, + mamba_state_cache_interval: Optional[NonNegativeInt] = Field( + default=None, + status="deprecated", + telemetry=False, + exclude=True, description= - "The number of tokens between cache steps in the Mamba prefix cache.") + "Deprecated alias for mamba_state_config.periodic_snapshot_interval.") + + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. + mamba_state_config: MambaStateConfig = Field( + default_factory=MambaStateConfig, + description="Configuration for reusable Mamba state snapshots.") use_kv_cache_manager_v2: bool | Literal["auto"] = Field( default="auto", @@ -3781,20 +3823,20 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): default=None, min_length=1, status="prototype", - description= - "Initial pool ratios for KV cache manager v2. When used by DeepSeek-V4, " - "values map to KVCacheManagerV2 pool_group_id order and must sum to 1.0. " - "When set, DeepSeek-V4 uses this directly and avg_seq_len does not take effect." - ) + description="Initial pool ratios for KV cache manager v2. Values map to " + "KVCacheManagerV2 pool_group_id order and must sum to 1.0. Hybrid Mamba " + "models and DeepSeek-V4 use this directly, so avg_seq_len does not take " + "effect when this is set.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. avg_seq_len: Optional[PositiveInt] = Field( default=None, status="prototype", description= - "Average sequence length used by DeepSeek-V4 to build the KV cache manager v2 " - "typical step. If unset, max_seq_len is used. This does not take effect when " - "pool_ratio is set.") + "Average total sequence length of the serving workload, used to build the " + "KV cache manager v2 typical step for hybrid Mamba models and DeepSeek-V4. " + "Hybrid Mamba models warn and fall back to half of max_seq_len when this is " + "unset. This does not take effect when pool_ratio is set.") # This is a pure python field, not a pybind field. It is only for the Pytorch backend. block_reuse_policy: Literal[ @@ -3805,8 +3847,10 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): "'all_reusable' commits reusable blocks after every context chunk; " "'per_request' commits them only after the final context chunk; " "'per_conversation' uses 'per_request' commits and drops the previous " - "turn's committed SWA-window blocks after the current turn's final context " - "chunk. All reusable blocks remain subject to normal cache eviction. " + "turn's committed SWA-window blocks and Mamba stable-boundary state " + "after the current turn's final context chunk. Periodic Mamba state " + "snapshots are disabled with 'per_conversation'. All reusable blocks " + "remain subject to normal cache eviction. " "Requests without conversation params use 'per_request' behavior. When " "'all_reusable' and SWA scratch reuse are both enabled, only non-scratch " "blocks are committed for reuse.") @@ -3870,6 +3914,43 @@ def validate_max_gpu_total_bytes(cls, v: int): "kv_cache_config.max_gpu_total_bytes must be non-negative") return v + @model_validator(mode='after') + def migrate_legacy_mamba_interval(self) -> 'KvCacheConfig': + """Copy the deprecated Mamba interval into its nested replacement.""" + if self.mamba_state_cache_interval is None: + return self + if ("periodic_snapshot_interval" + in self.mamba_state_config.model_fields_set): + raise ValueError("Cannot set both " + "'kv_cache_config.mamba_state_cache_interval' and " + "'kv_cache_config.mamba_state_config." + "periodic_snapshot_interval'.") + logger.warning( + "'kv_cache_config.mamba_state_cache_interval' is deprecated; use " + "'kv_cache_config.mamba_state_config." + "periodic_snapshot_interval' instead.") + self.mamba_state_config = self.mamba_state_config.model_copy( + update={ + "periodic_snapshot_interval": self.mamba_state_cache_interval + }) + return self + + @model_validator(mode='after') + def disable_periodic_mamba_snapshots_for_conversations( + self) -> 'KvCacheConfig': + """Use only explicit stable boundaries for conversation reuse.""" + if (self.block_reuse_policy == "per_conversation" + and self.mamba_state_config.periodic_snapshot_interval != 0): + interval = self.mamba_state_config.periodic_snapshot_interval + logger.warning( + f"'kv_cache_config.mamba_state_config.periodic_snapshot_interval={interval}' " + "is ignored because " + "'kv_cache_config.block_reuse_policy=per_conversation' disables " + "periodic Mamba snapshots; setting it to 0.") + self.mamba_state_config = self.mamba_state_config.model_copy( + update={"periodic_snapshot_interval": 0}) + return self + @model_validator(mode='after') def validate_disk_cache_config(self): if self.disk_cache_size is not None and self.disk_cache_size > 0: @@ -3883,6 +3964,18 @@ def validate_disk_cache_config(self): ) return self + @model_validator(mode='after') + def validate_mamba_snapshot_offsets(self) -> 'KvCacheConfig': + state_config = self.mamba_state_config + has_additional_snapshots = bool( + state_config.additional_snapshot_offsets_from_start + or state_config.additional_snapshot_offsets_from_end) + if (has_additional_snapshots and self.use_kv_cache_manager_v2 is False): + raise ValueError( + "kv_cache_config.mamba_state_config additional snapshot " + "offsets require kv_cache_config.use_kv_cache_manager_v2=True.") + return self + @field_validator('max_attention_window') @classmethod def validate_max_attention_window(cls, v: Optional[List[int]]): @@ -4454,6 +4547,10 @@ def speculative_model(self) -> Optional[Union[str, Path]]: def from_yaml(cls, yaml_path: Union[str, Path]): with open(yaml_path, "r") as f: config_dict = yaml.safe_load(f) + if config_dict is None: + config_dict = {} + elif not isinstance(config_dict, dict): + raise ValueError("Configuration file root must be a mapping.") return cls(**config_dict) @field_validator("dtype") @@ -5884,6 +5981,8 @@ def update_llm_args_with_extra_dict( If `explicit_cli_keys` is None, YAML wins on conflicts. """ + llm_args_dict = dict(llm_args_dict) + # CLI scalar -> nested KvCacheConfig field. Callers add the CLI scalar # name to `explicit_cli_keys` to make it win over YAML's same-named # field inside `kv_cache_config:`. @@ -5995,11 +6094,15 @@ def update_llm_args_with_extra_options( if extra_llm_api_options is not None: with open(extra_llm_api_options, 'r') as f: llm_args_dict = yaml.safe_load(f) - llm_args = update_llm_args_with_extra_dict( - llm_args, - llm_args_dict, - extra_llm_api_options, - explicit_cli_keys=explicit_cli_keys) + if llm_args_dict is None: + llm_args_dict = {} + elif not isinstance(llm_args_dict, dict): + raise ValueError("Configuration file root must be a mapping.") + llm_args = update_llm_args_with_extra_dict( + llm_args, + llm_args_dict, + extra_llm_api_options, + explicit_cli_keys=explicit_cli_keys) return llm_args diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 159ebbd7315a..63b360f5584b 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -558,9 +558,18 @@ def _compute_applied(defaults: Dict[str, Any], def _resolve_kv_cache_manager_v2_auto( - llm_args: 'TorchLlmArgs', model_defaults_dict: Dict[str, Any]) -> bool: - """Resolve the KV cache manager auto setting after model defaults are applied.""" - setting = llm_args.kv_cache_config.use_kv_cache_manager_v2 + llm_args: 'TorchLlmArgs', + model_defaults_dict: Dict[str, Any], + original_setting: Optional[Union[bool, str]] = None) -> bool: + """Resolve the KV cache manager auto setting after model defaults are applied. + + The transceiver runtime auto setting must be resolved first. In + disaggregated serving, hybrid Mamba V2 requires the Python transceiver with + NIXL, so an incompatible route falls back to V1 unless the user explicitly + selected V2. + """ + setting = (llm_args.kv_cache_config.use_kv_cache_manager_v2 + if original_setting is None else original_setting) if setting != "auto": return setting @@ -574,6 +583,18 @@ def _resolve_kv_cache_manager_v2_auto( "Model default kv_cache_config.use_kv_cache_manager_v2 must be " f"True, False, or 'auto', got {model_default!r}.") + transceiver_config = llm_args.cache_transceiver_config + if (model_default and transceiver_config is not None + and transceiver_config.backend is not None): + effective_backend, _ = transceiver_config._resolve_default_backend() + runtime = transceiver_config.transceiver_runtime + if effective_backend != "NIXL" or runtime != "PYTHON": + logger.info( + "KV cache manager V2 is the model default, but disaggregated " + "serving uses transceiver_runtime=%r with backend=%r; " + "falling back to V1.", runtime, effective_backend) + model_default = False + llm_args.kv_cache_config.use_kv_cache_manager_v2 = model_default return model_default diff --git a/tensorrt_llm/metrics/collector.py b/tensorrt_llm/metrics/collector.py index cd5a91123ddf..0586b46173ab 100644 --- a/tensorrt_llm/metrics/collector.py +++ b/tensorrt_llm/metrics/collector.py @@ -805,7 +805,15 @@ def log_iteration_stats(self, iteration_stats: dict) -> None: kv_iter_by_pool_group = iteration_stats.get( "kvCacheIterationStatsByPoolGroup") if kv_iter or kv_iter_by_lifecycle or kv_iter_by_pool_group: - reuse_stats = kv_iter_by_lifecycle or kv_iter or {} + # Prefer lifecycle-level attention stats when present. An SSM-only + # lifecycle report must not hide the legacy/window-level attention + # aggregate. Missing kind remains attention-compatible. + attention_lifecycle_stats = { + key: stats + for key, stats in (kv_iter_by_lifecycle or {}).items() + if stats.get("kind", "attention") == "attention" + } + reuse_stats = attention_lifecycle_stats or kv_iter or {} pool_group_stats = kv_iter_by_pool_group or kv_iter or {} total_secondary_max = 0 total_secondary_used = 0 diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 5557a1d2f32e..e96c62b5501c 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -78,6 +78,7 @@ _KV_CACHE_ITERATION_STATS_DELTA_FIELDS, KVCacheIterationStatsDelta, KVCacheStatsDelta, + SsmSnapshotIterationStatsDelta, ) from ._storage import BufferId # noqa: F401 from ._storage._config import CoalescedBuffer, SlotDesc, SlotDescVariant # noqa: F401 @@ -142,6 +143,7 @@ "SlotDesc", "SlotDescVariant", "SsmLayerConfig", + "SsmSnapshotIterationStatsDelta", "SwaScratchReuseConfig", "TokenId", "TokenIdExt", diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 0945db4d2394..1a03c885b73a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -93,6 +93,18 @@ class KVCacheIterationStatsDelta: iter_host_dropped_blocks: int = 0 iter_host_dropped_bytes: int = 0 +@dataclass(slots=True) +class SsmSnapshotIterationStatsDelta: + iter_snapshot_lookups: int = 0 + iter_snapshot_hits: int = 0 + iter_snapshot_misses: int = 0 + iter_reused_tokens: int = 0 + iter_unreused_tokens: int = 0 + iter_aligned_snapshot_hits: int = 0 + iter_unaligned_snapshot_hits: int = 0 + @property + def iter_snapshot_hit_rate(self) -> float: ... + @dataclass(slots=True, frozen=True) class PoolGroupPeakBlockStats: available: int @@ -485,6 +497,9 @@ class KVCacheManager: def get_quota(self, cache_level: CacheLevel) -> int: ... def get_committed_stats(self) -> KVCacheStatsDelta: ... def get_and_reset_iteration_stats(self) -> dict[LifeCycleId, KVCacheIterationStatsDelta]: ... + def get_and_reset_ssm_snapshot_iteration_stats( + self, + ) -> dict[LifeCycleId, SsmSnapshotIterationStatsDelta]: ... def get_and_reset_iteration_peak_block_stats( self, cache_level: CacheLevel ) -> Sequence[PoolGroupPeakBlockStats]: ... 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 38798e17b03a..110912a07ddb 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 @@ -67,6 +67,7 @@ class ReuseMatch(NamedTuple): blocks: list["Block"] num_tokens: int + num_lookup_tokens: int # id_offset is usually vocab_size @@ -361,6 +362,10 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N assert NDEBUG or (not b.is_full and b is not self and b.key == k and not b.next) to_remove.append(k) event_manager = get_tree(prev).event_manager if to_remove else None + # Keep RootBlock attached while covered children are replaced. Adding + # the replacement first prevents detach_next() from pruning an + # otherwise-empty root before this block becomes its new child. + prev.next[self.key] = self for k in to_remove: b = detach_next(prev, k) assert isinstance(b, Block) @@ -368,7 +373,6 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N 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 _release_pages(self) -> None: """Reclaim every page held by this block. @@ -675,7 +679,11 @@ def match( matched = self._prune_match( list(self._match_token_path(reuse_scope, tokens, enable_partial_match)) ) - return ReuseMatch([block for block, _ in matched], self._num_matched_tokens(matched)) + return ReuseMatch( + [block for block, _ in matched], + self._num_matched_tokens(matched), + len(tokens), + ) def _check_sanity(self) -> bool: raise NotImplementedError( 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 5494449c0bd6..303d79e9ee13 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 @@ -412,6 +412,9 @@ def commit_pending_stats(self) -> KVCacheStatsDelta: self.manager.commit_stats( self._pending_stats.global_stats, self._pending_stats.iteration_stats_by_life_cycle ) + self.manager._commit_ssm_snapshot_iteration_stats( + self._pending_stats.ssm_snapshot_iteration_stats_by_life_cycle + ) request_stats = self._pending_stats.request_stats.copy() self._pending_stats.clear() self.manager.clear_stats_dirty(self.id) @@ -1047,22 +1050,41 @@ def reuse_scope(self) -> ReuseScope: return self._reuse_scope def plan_committed_block_drop(self) -> PlannedDropHandle | None: - """Plan dropping SWA blocks needed only by the next conversation turn. + """Plan dropping pages needed only by the next conversation turn. The plan covers committed pages in each SWA life cycle's current - attention window. Full-attention and attention-sink blocks are excluded - because later turns may still need them. SSM state is not yet supported. - This must be called after stop_committing(). Returns None without - creating a plan if any required SWA page is unavailable. + attention window and the exact SSM snapshot for the committed prefix. + Full-attention and attention-sink blocks are excluded because later + turns may still need them. This must be called after stop_committing(). + Returns None without creating a plan if any required page is unavailable. """ if self._commit_state != self.CommitState.USER_STOP: raise LogicError("plan_committed_block_drop() requires stop_committing()") - end = self._num_committed_blocks + ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id + matched_blocks: list[Block] = [] + if self.num_committed_tokens > 0: + # Locate pages through the radix tree rather than + # SeqBlock.tree_block: the latter is not guaranteed to identify a + # partial snapshot after reuse. Requiring an exact match keeps the + # preceding conversation plan intact if this turn no longer has a + # complete reusable endpoint. All PP ranks use the same lookup so + # attention-only ranks include the final partial SWA block too. + match = self.manager._match_reuse(self.reuse_scope, self._committed_tokens) + if match.num_tokens != self.num_committed_tokens or not match.blocks: + return None + matched_blocks = match.blocks + elif ssm_lc_id is not None: + return None + + def get_committed_page(block: Block, life_cycle_id: LifeCycleId) -> CommittedPage | None: + page_ref = block.storage[life_cycle_id] + return page_ref() if page_ref is not None else None + + end = BlockOrdinal(len(matched_blocks)) pages_to_drop: list[CommittedPage] = [] for lc_idx, lc in self.manager._life_cycles.items(): if isinstance(lc, SsmLifeCycle): - # TODO: Support recording reusable SSM state pages. continue if lc.window_size is None: continue @@ -1071,16 +1093,17 @@ def plan_committed_block_drop(self) -> PlannedDropHandle | None: ) window_start = min(stale_range.end, end) for ordinal in typed_range(window_start, end): - tree_block = self._blocks[ordinal].tree_block - if tree_block is None: - return None - page_ref = tree_block.storage[lc_idx] - if page_ref is None: - return None - page = page_ref() + tree_block = matched_blocks[ordinal] + page = get_committed_page(tree_block, lc_idx) if page is None: return None pages_to_drop.append(page) + + if ssm_lc_id is not None: + page = get_committed_page(matched_blocks[-1], ssm_lc_id) + if not isinstance(page, SsmCommittedPage): + return None + pages_to_drop.append(page) return PlannedDropHandle(pages_to_drop) # Users promise to not commit any more tokens. For cases where we shouldn't reuse generated tokens @@ -2059,6 +2082,15 @@ def _setup_for_reuse(self, match: ReuseMatch) -> None: ) snapshot_holder = unwrap_rawref(snapshot_ref).hold() self._ssm_blocks[DEFAULT_BEAM_INDEX][ssm_lc_id] = snapshot_holder + if should_record_stats and ssm_lc_id is not None: + changed = self._pending_stats.record_ssm_snapshot_lookup( + ssm_lc_id, + lookup_tokens=match.num_lookup_tokens, + reused_tokens=num_tokens, + tokens_per_block=tokens_per_block, + ) + if changed: + self.manager.mark_stats_dirty(self.id) self._num_committed_blocks = BlockOrdinal(len(self._committed_tokens) // tokens_per_block) for beam_indices in self._base_page_indices: for indices in beam_indices: 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 2d991b617db3..2492278bc9df 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 @@ -39,7 +39,7 @@ from .._config import DataRole, KVCacheManagerConfig from .._life_cycle_registry import LayerGroupId, LifeCycle, LifeCycleId, LifeCycleRegistry from .._page import Page, _PageHolder -from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta +from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta, SsmSnapshotIterationStatsDelta from .._storage._config import BufferId, SlotDesc, create_storage_config from .._storage._core import PoolGroupIndex, PoolIndex, SlotId from .._storage_manager import StorageManager @@ -212,6 +212,7 @@ class KVCacheManager: "_stats_enabled", "_committed_stats", "_iteration_stats_by_life_cycle", + "_ssm_snapshot_iteration_stats_by_life_cycle", "_iteration_peak_num_blocks_by_cache_level", "_dirty_stats_kv_cache_ids", "_stats_excluded_kv_cache_ids", @@ -241,6 +242,7 @@ class KVCacheManager: _stats_enabled: bool _committed_stats: KVCacheStatsDelta _iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] + _ssm_snapshot_iteration_stats_by_life_cycle: dict[LifeCycleId, SsmSnapshotIterationStatsDelta] _iteration_peak_num_blocks_by_cache_level: TypedIndexList[ CacheLevel, TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats] ] @@ -284,6 +286,7 @@ def __init__( self._stats_enabled = config.enable_stats self._committed_stats = KVCacheStatsDelta() self._iteration_stats_by_life_cycle = {} + self._ssm_snapshot_iteration_stats_by_life_cycle = {} self._reset_iteration_peak_num_blocks() self._dirty_stats_kv_cache_ids = set() self._stats_excluded_kv_cache_ids = set() @@ -539,6 +542,31 @@ def get_and_reset_iteration_stats(self) -> dict[LifeCycleId, KVCacheIterationSta self._iteration_stats_by_life_cycle.clear() return stats + def _commit_ssm_snapshot_iteration_stats( + self, + iteration_stats_by_life_cycle: dict[LifeCycleId, SsmSnapshotIterationStatsDelta], + ) -> None: + if not self._stats_enabled: + return + for life_cycle, iteration_stats in iteration_stats_by_life_cycle.items(): + if iteration_stats.empty: + continue + destination = self._ssm_snapshot_iteration_stats_by_life_cycle.setdefault( + life_cycle, SsmSnapshotIterationStatsDelta() + ) + destination.add(iteration_stats) + + def get_and_reset_ssm_snapshot_iteration_stats( + self, + ) -> dict[LifeCycleId, SsmSnapshotIterationStatsDelta]: + stats = { + life_cycle: delta.copy() + for life_cycle, delta in self._ssm_snapshot_iteration_stats_by_life_cycle.items() + if not delta.empty + } + self._ssm_snapshot_iteration_stats_by_life_cycle.clear() + return stats + def get_and_reset_iteration_peak_block_stats( self, cache_level: CacheLevel ) -> TypedIndexList[PoolGroupIndex, PoolGroupPeakBlockStats]: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py index 94d0957537be..bf8318e1c726 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py @@ -17,7 +17,7 @@ from .._common import BlockOrdinal from .._life_cycle_registry import LifeCycleId -from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta +from .._stats import KVCacheIterationStatsDelta, KVCacheStatsDelta, SsmSnapshotIterationStatsDelta @dataclass(slots=True) @@ -49,6 +49,9 @@ class _PendingStats: iteration_stats_by_life_cycle: dict[LifeCycleId, KVCacheIterationStatsDelta] = field( default_factory=dict ) + ssm_snapshot_iteration_stats_by_life_cycle: dict[ + LifeCycleId, SsmSnapshotIterationStatsDelta + ] = field(default_factory=dict) allocation_segments: list[_PendingAllocationSegment] = field(default_factory=list) @property @@ -57,12 +60,14 @@ def empty(self) -> bool: self.request_stats.empty and self.global_stats.empty and not self.iteration_stats_by_life_cycle + and not self.ssm_snapshot_iteration_stats_by_life_cycle ) def clear(self) -> None: self.request_stats.clear() self.global_stats.clear() self.iteration_stats_by_life_cycle.clear() + self.ssm_snapshot_iteration_stats_by_life_cycle.clear() self.allocation_segments.clear() def add(self, delta: _PendingStatsDelta) -> bool: @@ -165,6 +170,39 @@ def record_reuse( ) ) + def record_ssm_snapshot_lookup( + self, + life_cycle: LifeCycleId, + *, + lookup_tokens: int, + reused_tokens: int, + tokens_per_block: int, + ) -> bool: + if lookup_tokens == 0: + return False + assert lookup_tokens > 0 + assert 0 <= reused_tokens <= lookup_tokens + assert tokens_per_block > 0 + + is_hit = reused_tokens > 0 + # Alignment describes the reusable snapshot boundary, not whether the + # state itself is complete. Every hit represents one complete SSM + # snapshot; token counters carry the benefit of that single lookup. + delta = SsmSnapshotIterationStatsDelta( + iter_snapshot_lookups=1, + iter_snapshot_hits=int(is_hit), + iter_snapshot_misses=int(not is_hit), + iter_reused_tokens=reused_tokens, + iter_unreused_tokens=lookup_tokens - reused_tokens, + iter_aligned_snapshot_hits=int(is_hit and reused_tokens % tokens_per_block == 0), + iter_unaligned_snapshot_hits=int(is_hit and reused_tokens % tokens_per_block != 0), + ) + pending = self.ssm_snapshot_iteration_stats_by_life_cycle.setdefault( + life_cycle, SsmSnapshotIterationStatsDelta() + ) + pending.add(delta) + return True + def subtract_allocation_range(self, block_begin: BlockOrdinal, block_end: BlockOrdinal) -> bool: if block_begin >= block_end or not self.allocation_segments: return False diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py index 292876ff4985..03d95d745df0 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py @@ -78,6 +78,23 @@ def iter_cache_hit_rate(self) -> float: return self.iter_reused_blocks / total +@dataclass(slots=True) +class SsmSnapshotIterationStatsDelta(_StatsDeltaMixin): + iter_snapshot_lookups: int = 0 + iter_snapshot_hits: int = 0 + iter_snapshot_misses: int = 0 + iter_reused_tokens: int = 0 + iter_unreused_tokens: int = 0 + iter_aligned_snapshot_hits: int = 0 + iter_unaligned_snapshot_hits: int = 0 + + @property + def iter_snapshot_hit_rate(self) -> float: + if self.iter_snapshot_hits == 0 or self.iter_snapshot_lookups == 0: + return 0.0 + return self.iter_snapshot_hits / self.iter_snapshot_lookups + + _KV_CACHE_ITERATION_STATS_DELTA_FIELDS = tuple( field.name for field in fields(KVCacheIterationStatsDelta) ) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index fb87dbc2b432..fb0df8ea4b87 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -720,7 +720,7 @@ "annotation": "", "converter": "", "kind": "value", - "path": "kv_cache_config.mamba_state_cache_interval" + "path": "kv_cache_config.mamba_state_config.periodic_snapshot_interval" }, { "allowed_values": [], diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 378aa74751db..7bb1457f2bed 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -31,11 +31,11 @@ from tensorrt_llm.llmapi import ( AttentionDpConfig, CudaGraphConfig, DeepSeekSparseAttentionConfig, DFlashDecodingConfig, DSparkDecodingConfig, DraftTargetDecodingConfig, - Eagle3DecodingConfig, KvCacheConfig, MiniMaxM3SparseAttentionConfig, - MoeConfig, MTPDecodingConfig, NGramDecodingConfig, PARDDecodingConfig, - RocketSparseAttentionConfig, SADecodingConfig, SamplingParams, - SchedulerConfig, SkipSoftmaxAttentionConfig, SAEnhancerConfig, - TorchCompileConfig) + Eagle3DecodingConfig, KvCacheConfig, MambaStateConfig, + MiniMaxM3SparseAttentionConfig, MoeConfig, MTPDecodingConfig, + NGramDecodingConfig, PARDDecodingConfig, RocketSparseAttentionConfig, + SADecodingConfig, SamplingParams, SchedulerConfig, + SkipSoftmaxAttentionConfig, SAEnhancerConfig, TorchCompileConfig) # isort: on from tensorrt_llm.quantization import QuantAlgo @@ -5925,7 +5925,7 @@ def test_nvfp4(self, moe_backend, tp_size, ep_size, attention_dp, kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, enable_block_reuse=enable_block_reuse) if enable_block_reuse: - kv_cache_config.mamba_state_cache_interval = 256 + kv_cache_config.mamba_state_config.periodic_snapshot_interval = 256 pytorch_config = dict(disable_overlap_scheduler=False, cuda_graph_config=CudaGraphConfig( max_batch_size=512, enable_padding=True)) @@ -6291,6 +6291,9 @@ def test_fp8(self, enable_block_reuse, mocker): kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8, enable_block_reuse=enable_block_reuse) + if enable_block_reuse: + kv_cache_config.avg_seq_len = 2048 + kv_cache_config.mamba_state_config.periodic_snapshot_interval = 256 # DeepGEMM MoE kernels only support datacenter Blackwell (SM100/SM103). # Fall back to the CUTLASS MoE backend (which supports FP8 block scales) # on other architectures such as Hopper (SM90) and consumer Blackwell @@ -6444,6 +6447,9 @@ def test_nvfp4(self, tp_size, ep_size, attention_dp, moe_backend, kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9, enable_block_reuse=enable_block_reuse) + if enable_block_reuse: + kv_cache_config.avg_seq_len = 2048 + kv_cache_config.mamba_state_config.periodic_snapshot_interval = 256 cuda_graph_config = CudaGraphConfig(max_batch_size=256, enable_padding=True) @@ -7128,7 +7134,7 @@ def test_fp8_4gpus(self, attention_dp, use_cpp_mamba, monkeypatch): @skip_pre_blackwell @pytest.mark.parametrize( - "tp_size, ep_size, mamba_state_cache_interval, attention_dp, use_mtp", + "tp_size, ep_size, periodic_snapshot_interval, attention_dp, use_mtp", [ (1, 1, 256, False, False), (4, 1, 256, False, True), @@ -7139,7 +7145,7 @@ def test_fp8_4gpus(self, attention_dp, use_cpp_mamba, monkeypatch): ids=["TP1", "TP4_MTP", "TEP4", "DEP4_MTP_OFF", "DEP4_MTP_ON"], ) def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, - mamba_state_cache_interval, attention_dp, + periodic_snapshot_interval, attention_dp, use_mtp): gpu_needed = max(tp_size, ep_size) if get_device_count() < gpu_needed: @@ -7154,8 +7160,10 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, f"{llm_models_root()}/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4", kv_cache_config=KvCacheConfig( enable_block_reuse=True, + avg_seq_len=2048, mamba_ssm_cache_dtype="float16", - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=periodic_snapshot_interval), free_gpu_memory_fraction=0.8, ), max_batch_size=32, @@ -7228,7 +7236,10 @@ def test_nvfp4_8gpus_mtp(self): model_path, kv_cache_config=KvCacheConfig( enable_block_reuse=True, + avg_seq_len=2048, mamba_ssm_cache_dtype="float16", + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=256), free_gpu_memory_fraction=0.5, ), max_batch_size=32, @@ -7561,7 +7572,7 @@ def test_nvfp4_8gpus(self, attention_dp, moe_backend): @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(4) @pytest.mark.parametrize( - "tp_size, ep_size, mamba_state_cache_interval, attention_dp, use_mtp", + "tp_size, ep_size, periodic_snapshot_interval, attention_dp, use_mtp", [ (4, 1, 256, False, True), (4, 4, 256, False, False), @@ -7571,7 +7582,7 @@ def test_nvfp4_8gpus(self, attention_dp, moe_backend): ids=["TP4_MTP", "TEP4", "ADP4", "ADP4_MTP"], ) def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, - mamba_state_cache_interval, attention_dp, + periodic_snapshot_interval, attention_dp, use_mtp): mtp_config = MTPDecodingConfig( num_nextn_predict_layers=3, @@ -7582,8 +7593,10 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, f"{llm_models_root()}/NVIDIA-Nemotron-3-Ultra-550B-A55B-NVFP4", kv_cache_config=KvCacheConfig( enable_block_reuse=True, + avg_seq_len=2048, mamba_ssm_cache_dtype="float16", - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=periodic_snapshot_interval), free_gpu_memory_fraction=0.5, ), max_batch_size=max_batch_size, diff --git a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py index 3d9a3f83dcc3..08a6f4701761 100644 --- a/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py +++ b/tests/unittest/_torch/executor/test_disagg_inflight_cancel_gate.py @@ -23,6 +23,10 @@ from tensorrt_llm._torch.pyexecutor import py_executor as executor_module from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import BindKvCacheTransceiver from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, +) from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig @@ -519,6 +523,59 @@ def test_flag_unset_preserves_python_transceiver(monkeypatch): constructor.assert_called_once() +def test_python_nixl_transceiver_accepts_v2_mamba_manager(monkeypatch): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + expected = object() + constructor = Mock(return_value=expected) + fake_module = SimpleNamespace(KvCacheTransceiverV2=constructor) + monkeypatch.setitem(sys.modules, "tensorrt_llm._torch.disaggregation.transceiver", fake_module) + manager = object.__new__(MambaHybridCacheManagerV2) + + result = transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + assert result is expected + constructor.assert_called_once() + + +@pytest.mark.parametrize("runtime", [None, "CPP", "auto"]) +def test_cpp_runtime_rejects_v2_mamba_manager(runtime): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(MambaHybridCacheManagerV2) + + with pytest.raises(ValueError, match="requires transceiver_runtime='PYTHON'"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + +def test_python_runtime_rejects_cpp_mamba_manager(): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime="PYTHON") + manager = object.__new__(CppMambaHybridCacheManager) + + with pytest.raises(ValueError, match="cannot drive CppMambaHybridCacheManager"): + transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + +@pytest.mark.parametrize("runtime", [None, "CPP"]) +def test_cpp_runtime_keeps_cpp_mamba_manager(monkeypatch, runtime): + config = CacheTransceiverConfig(backend="NIXL", transceiver_runtime=runtime) + manager = object.__new__(CppMambaHybridCacheManager) + expected = object() + constructor = Mock(return_value=expected) + monkeypatch.setattr(transceiver_module, "BindKvCacheTransceiver", constructor) + + result = transceiver_module.create_kv_cache_transceiver( + Mock(), Mock(), manager, Mock(), config, manager + ) + + assert result is expected + constructor.assert_called_once() + + def test_flag_unset_preserves_libfabric_selection(monkeypatch): monkeypatch.setenv(transceiver_module._NIXL_KVCACHE_BACKEND_ENV, "LIBFABRIC") config = CacheTransceiverConfig(backend="NIXL") diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index e5245db947fc..e45be30dcab0 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -510,6 +510,39 @@ def test_per_conversation_policy_ignores_overlapping_request( _free_if_active(manager, request_a) +def test_iteration_stats_reports_physical_pool_groups_without_window_metadata() -> None: + manager = object.__new__(KVCacheManagerV2) + manager.enable_stats = True + snapshot_delta = SimpleNamespace( + iter_snapshot_lookups=2, + iter_snapshot_hits=1, + iter_snapshot_misses=1, + iter_reused_tokens=32, + iter_unreused_tokens=16, + iter_aligned_snapshot_hits=1, + iter_unaligned_snapshot_hits=0, + ) + manager.impl = SimpleNamespace( + cache_tier_list=[object()], + get_and_reset_iteration_stats=lambda: {}, + get_and_reset_ssm_snapshot_iteration_stats=lambda: {3: snapshot_delta}, + ) + manager._stats_life_cycle_metadata = lambda: {3: (1, None, "ssm")} + manager._storage_pool_groups_by_window = lambda: {} + manager._get_and_reset_iteration_peak_block_stats = lambda _level: [None, None] + manager._get_storage_statistics = lambda _level: [object(), object()] + manager._build_pool_group_iteration_stats = lambda pool_group_id, *_args: pool_group_id + + stats = manager.get_iteration_stats() + + assert stats.by_pool_group == {0: 0, 1: 1} + ssm_stats = stats.by_life_cycle[3] + assert ssm_stats.kind == "ssm" + assert ssm_stats.pool_group_id == 1 + assert ssm_stats.snapshot_stats.iter_snapshot_hit_rate == 0.5 + assert ssm_stats.snapshot_stats.iter_reused_tokens == 32 + + def test_disagg_role_mapper_kinds_default_to_indexed(): from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py index df38d7bba8fc..a37288a80895 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -16,18 +16,31 @@ CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType -def _manager(*, is_draft: bool, kv_compression_manages_history: bool = False) -> KVCacheManagerV2: +def _manager( + *, + is_draft: bool, + kv_compression_manages_history: bool = False, + kv_reserve_draft_tokens: int = 0, +) -> KVCacheManagerV2: manager = KVCacheManagerV2.__new__(KVCacheManagerV2) manager.is_draft = is_draft manager.kv_compression_manages_history = kv_compression_manages_history + manager._kv_reserve_draft_tokens = kv_reserve_draft_tokens manager.kv_cache_map = {} return manager -def _request(request_id: int, *, rewind: int = 0, complete: bool = False) -> SimpleNamespace: +def _request( + request_id: int, + *, + rewind: int = 0, + accepted_draft_tokens: int = 0, + complete: bool = False, +) -> SimpleNamespace: return SimpleNamespace( py_request_id=request_id, py_rewind_len=rewind, + py_num_accepted_draft_tokens=accepted_draft_tokens, max_beam_num_tokens=201, state=LlmRequestState.GENERATION_COMPLETE if complete @@ -107,6 +120,23 @@ def test_capacity_only_is_scoped_to_target_manager() -> None: target_cache.resize.assert_called_once_with(253, None) +@pytest.mark.parametrize( + ("is_draft", "expected_capacity"), + [(True, 201), (False, 230)], + ids=["draft-reclaims-reserve", "target-has-no-reserve"], +) +def test_dynamic_tree_reserved_capacity(is_draft: bool, expected_capacity: int) -> None: + manager = _manager(is_draft=is_draft, kv_reserve_draft_tokens=60) + # The runtime tree used 31 draft positions: 26 rejected and 5 accepted. + request = _request(1, rewind=26, accepted_draft_tokens=5) + cache = _cache() + manager.kv_cache_map[request.py_request_id] = cache + + manager.update_resources(SimpleNamespace(generation_requests=[request])) + + cache.resize.assert_called_once_with(expected_capacity, 200) + + def test_capacity_only_completion_preserves_history() -> None: manager = _manager(is_draft=False, kv_compression_manages_history=True) request = _request(1, complete=True) diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 1e545ac5cb79..0a7c67e9e702 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -1542,6 +1542,21 @@ def test_force_chunk_allows_snapshot_point_shorter_than_unit_size(self): assert ids(out.context_requests) == [0] assert req.context_chunk_size == 150 + def test_force_chunk_makes_progress_before_budget_limited_snapshot(self): + mgr = make_kv_cache_manager(tokens_per_block=32) + sched = make_scheduler( + mgr, + max_num_tokens=128, + ctx_chunk_config=(ContextChunkingPolicy.FORCE_CHUNK, 32), + ) + req = make_ctx_request(0, context_remaining_length=512) + req.expect_snapshot_points = [256, 512] + + out = sched.schedule_request([req], set()) + + assert ids(out.context_requests) == [0] + assert req.context_chunk_size == 128 + def test_multiple_ctx_share_budget(self): """Two ctx requests share the budget.""" mgr = make_kv_cache_manager(tokens_per_block=64) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 0291ef3e74dd..e7d61248e7b4 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1,15 +1,29 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Regression tests for MambaCacheManager padding-slot behavior and -CppMambaHybridCacheManager PP-sharding edge cases.""" +"""Regression tests for Python, Cpp, and V2 Mamba cache managers.""" import os from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch +from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager +from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MambaLayerGroup +from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.pyexecutor._util import ( + KvCacheCreator, + _create_kv_cache_manager, + get_kv_cache_manager_cls, +) +from tensorrt_llm._torch.pyexecutor.config_utils import ( + MambaKVCacheParams, + extract_mamba_kv_cache_params, +) from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import BlockReusePolicy, KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, @@ -18,19 +32,586 @@ from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, + MambaRole, + MixedMambaHybridCacheManager, PythonMambaCacheManager, + ReplayStateUpdateMetadata, + _advance_replay_state, + _get_local_mamba_cache_layout, _get_mamba_hybrid_pool_size, + _get_num_cuda_graph_padding_dummy_slots, + _mamba_snapshot_rule_counts, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import ( + CacheTypeCpp, + DataType, + KVCacheManager, + get_pp_layers, ) -from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig +from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, + KvCacheConfig, + MambaStateConfig, + MTPDecodingConfig, + TorchLlmArgs, +) +from tensorrt_llm.llmapi.llm_utils import ( + _resolve_kv_cache_manager_v2_auto, + _resolve_transceiver_runtime_auto, + apply_model_defaults_to_llm_args, +) from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BatchDesc, + BufferConfig, + GpuCacheTierConfig, + KVCacheDesc, + KVCacheManagerConfig, + LayerId, + SsmLayerConfig, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as RuntimeKVCacheManager skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_advance_replay_state_uses_checkpoint_predicate_and_skips_dummies(): + metadata = ReplayStateUpdateMetadata( + prev_num_accepted_tokens=torch.tensor([11, 12, 13], dtype=torch.int32), + cache_buf_idx=torch.tensor([0, 1, 1], dtype=torch.int32), + replay_step_width=5, + replay_history_size=16, + ) + + _advance_replay_state( + metadata, + state_indices=torch.tensor([0, 1, 2], dtype=torch.int32), + accepted_tokens=torch.tensor([2, 3, 4], dtype=torch.int32), + is_dummy_request=torch.tensor([False, False, True]), + ) + + # Equality does not write a checkpoint; overflow does. Dummy slots do not + # advance either piece of replay bookkeeping. + assert metadata.prev_num_accepted_tokens.tolist() == [13, 3, 13] + assert metadata.cache_buf_idx.tolist() == [0, 0, 1] + + +def test_cuda_graph_padding_dummy_slot_count_tracks_reachable_draft_lengths(): + assert _get_num_cuda_graph_padding_dummy_slots(None, 64) == 1 + + static = MTPDecodingConfig(max_draft_len=4) + assert _get_num_cuda_graph_padding_dummy_slots(static, 64) == 1 + + gated = MTPDecodingConfig( + max_draft_len=4, + acceptance_rate_window_size=8, + acceptance_rate_threshold=0.5, + ) + assert _get_num_cuda_graph_padding_dummy_slots(gated, 64) == 2 + + dynamic = MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={4: 4, 8: 2, 32: 1}, + ) + assert _get_num_cuda_graph_padding_dummy_slots(dynamic, 5) == 2 + assert _get_num_cuda_graph_padding_dummy_slots(dynamic, 32) == 3 + assert _get_num_cuda_graph_padding_dummy_slots(dynamic, 64) == 4 + + repeated = MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={4: 4, 8: 4, 32: 2}, + ) + assert _get_num_cuda_graph_padding_dummy_slots(repeated, 64) == 3 + + +def _hybrid_model_config(): + config = SimpleNamespace( + architectures=["Qwen3_5MoeForCausalLM"], + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + ) + return SimpleNamespace( + pretrained_config=config, + sparse_attention_config=None, + get_num_mamba_layers=lambda: 1, + ) + + +def _hybrid_cache_sizing_model_config(layer_types): + config = SimpleNamespace( + architectures=["Qwen3_5ForCausalLM"], + num_hidden_layers=len(layer_types), + layer_types=layer_types, + linear_key_head_dim=8, + linear_conv_kernel_dim=4, + linear_num_value_heads=4, + linear_num_key_heads=1, + linear_value_head_dim=8, + num_key_value_heads=2, + num_attention_heads=4, + hidden_size=32, + torch_dtype=torch.float16, + ) + return SimpleNamespace(pretrained_config=config, quant_config=None) + + +def test_mamba_kv_cache_params_separate_target_and_draft_masks(): + model_config = _hybrid_cache_sizing_model_config( + [ + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + ] + ) + params = extract_mamba_kv_cache_params( + model_config.pretrained_config, + spec_config=MTPDecodingConfig(max_draft_len=1), + ) + + assert params.mamba_layer_mask == [True, False, True, False] + assert params.target_full_attention_layer_mask == [False, True, False, True] + assert params.num_draft_layers == 1 + + assert params.get_layer_masks() == ( + [True, False, True, False, False], + [False, True, False, True, True], + ) + assert params.get_layer_masks(use_separate_draft_kv_cache=True) == ( + [True, False, True, False], + [False, True, False, True], + ) + assert params.get_layer_masks(is_draft=True) == ( + [False, False, False, False, False], + [False, False, False, False, True], + ) + + +@pytest.mark.parametrize( + ("use_v2", "enable_block_reuse", "expected"), + [ + (True, False, MambaHybridCacheManagerV2), + (True, True, MambaHybridCacheManagerV2), + (False, True, CppMambaHybridCacheManager), + ("auto", True, CppMambaHybridCacheManager), + ], +) +def test_hybrid_cache_manager_factory_honors_v2_setting( + monkeypatch, use_v2, enable_block_reuse, expected +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + kv_cache_config = KvCacheConfig( + enable_block_reuse=enable_block_reuse, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=use_v2, + ) + + assert get_kv_cache_manager_cls(_hybrid_model_config(), kv_cache_config) is expected + + +def test_qwen3_gdn_replay_falls_back_for_v2_manager(monkeypatch): + """GDN's all-layer replay commit requires the contiguous C++ state view.""" + captured_cpp = {} + captured_v2 = {} + + class RecordingCppManager(CppMambaHybridCacheManager): + def __init__(self, *args, **kwargs): + captured_cpp.update(kwargs) + + class RecordingV2Manager(MambaHybridCacheManagerV2): + def __init__(self, *args, **kwargs): + captured_v2.update(kwargs) + + pretrained_config = SimpleNamespace( + architectures=["Qwen3_5MoeForCausalLM"], + hidden_size=32, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + ) + model_config = SimpleNamespace( + pretrained_config=pretrained_config, + quant_config=None, + ) + mamba_params = MambaKVCacheParams( + state_size=8, + conv_kernel=4, + num_heads=4, + n_groups=1, + head_dim=8, + mamba_layer_mask=[True, False], + target_full_attention_layer_mask=[False, True], + num_mamba_layers=1, + num_draft_layers=1, + dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.bfloat16, + ) + monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "1") + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 90) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor._util.extract_mamba_kv_cache_params", + lambda *args, **kwargs: mamba_params, + ) + + common_kwargs = dict( + model_engine=None, + mapping=Mapping(world_size=1, tp_size=1, pp_size=1), + tokens_per_block=32, + max_seq_len=2048, + # Cover the batch size at which upstream selects partitioned replay. + max_batch_size=16, + spec_config=MTPDecodingConfig(max_draft_len=3), + sparse_attention_config=None, + max_num_tokens=256, + max_beam_width=1, + kv_connector_manager=None, + model_config=model_config, + dtype=torch.bfloat16, + is_draft=False, + ) + _create_kv_cache_manager( + kv_cache_manager_cls=RecordingCppManager, + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=False), + **common_kwargs, + ) + _create_kv_cache_manager( + kv_cache_manager_cls=RecordingV2Manager, + kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=True), + **common_kwargs, + ) + + assert captured_cpp["use_replay_state_update"] is True + assert captured_cpp["model_type"] == "qwen3_next" + assert captured_v2["use_replay_state_update"] is False + assert "model_type" not in captured_v2 + assert captured_v2["conv_state_layout"] == "q_k_v" + + +def test_hybrid_cache_manager_factory_rejects_cpp_preference_with_explicit_v2( + monkeypatch, +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "CPP") + + with pytest.raises(ValueError, match="conflicts with explicit"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=True, + ), + ) + + +def test_hybrid_cache_manager_factory_v2_preference_does_not_select_v2( + monkeypatch, +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "V2") + + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=False, + ), + ) + is CppMambaHybridCacheManager + ) + + +@pytest.mark.parametrize( + ("field", "offsets"), + [ + ("additional_snapshot_offsets_from_start", [128]), + ("additional_snapshot_offsets_from_end", [0]), + ], +) +def test_hybrid_cache_manager_factory_requires_v2_for_explicit_snapshots( + monkeypatch, + field, + offsets, +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + mamba_state_config=MambaStateConfig(**{field: offsets}), + use_kv_cache_manager_v2="auto", + ) + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_config=kv_cache_config, + ) + + assert _resolve_kv_cache_manager_v2_auto(llm_args, {}) is False + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + with pytest.raises(ValueError, match="use_kv_cache_manager_v2=True"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + llm_args.kv_cache_config, + ) + + +@pytest.mark.parametrize("backend", ["NIXL", "DEFAULT"]) +def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, backend): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=True, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime="PYTHON" + ), + ) + is MambaHybridCacheManagerV2 + ) + + +def test_hybrid_cache_manager_factory_rejects_python_v1_disagg_reuse(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + with pytest.raises(ValueError, match="requires use_kv_cache_manager_v2=True"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=False, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + ), + ) + + +@pytest.mark.parametrize( + ("backend", "runtime", "backend_env"), + [ + ("DEFAULT", "PYTHON", "TRTLLM_USE_UCX_KVCACHE"), + ("UCX", "PYTHON", None), + ("NIXL", "auto", None), + ("NIXL", None, None), + ("NIXL", "CPP", None), + ("UCX", None, None), + ], +) +def test_hybrid_cache_manager_factory_rejects_unsupported_v2_disagg_route( + monkeypatch, backend, runtime, backend_env +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + if backend_env is not None: + monkeypatch.setenv(backend_env, "1") + + with pytest.raises(ValueError, match="requires transceiver_runtime='PYTHON'"): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + use_kv_cache_manager_v2=True, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime=runtime + ), + ) + + +@pytest.mark.parametrize( + ("env_name", "env_value", "expected_error"), + [ + ( + "TLLM_MAMBA_MANAGER_PREFERENCE", + "MIXED", + "does not support block reuse", + ), + ( + "TRTLLM_USE_PY_MAMBA", + "1", + "does not support block reuse", + ), + ], +) +def test_hybrid_cache_manager_factory_rejects_mixed_override_with_reuse( + monkeypatch, env_name, env_value, expected_error +): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + monkeypatch.setenv(env_name, env_value) + + with pytest.raises(ValueError, match=expected_error): + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + ), + ) + + +@pytest.mark.parametrize("use_v2", [False, "auto"]) +def test_hybrid_cache_manager_factory_keeps_v1_disagg_route(monkeypatch, use_v2): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=False, + use_kv_cache_manager_v2=use_v2, + ), + is_disagg=True, + ) + is CppMambaHybridCacheManager + ) + assert ( + get_kv_cache_manager_cls( + _hybrid_model_config(), + KvCacheConfig( + enable_block_reuse=False, + use_kv_cache_manager_v2=use_v2, + ), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON" + ), + ) + is MixedMambaHybridCacheManager + ) + + +def test_hybrid_models_default_to_v2_and_python_transceiver(monkeypatch): + from tensorrt_llm._torch.models.modeling_nemotron_h import NemotronHForCausalLM + from tensorrt_llm._torch.models.modeling_qwen3_5 import Qwen3_5VLModel + from tensorrt_llm._torch.models.modeling_qwen3_next import Qwen3NextForCausalLM + + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + + for model_cls in (NemotronHForCausalLM, Qwen3NextForCausalLM, Qwen3_5VLModel): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig(backend="DEFAULT"), + ) + model_defaults = model_cls.get_model_defaults(llm_args) + apply_model_defaults_to_llm_args(llm_args, model_defaults) + _resolve_transceiver_runtime_auto(llm_args, model_cls) + _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults, original_setting="auto") + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + assert llm_args.kv_cache_config.enable_block_reuse is False + assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +def test_v2_disagg_slice_skips_state_index_on_mamba_free_pp_rank(): + manager = object.__new__(MambaHybridCacheManagerV2) + manager.local_num_mamba_layers = 0 + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._kv_cache_manager = manager + transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) + transceiver._page_table = SimpleNamespace(layer_groups=[]) + request = SimpleNamespace( + is_generation_only_request=lambda: False, + prompt_len=0, + py_request_id=123, + ) + + kv_slice = transceiver._create_kv_slice(request) + + assert kv_slice.mamba_state_index is None + + +def test_v2_disagg_slice_reads_state_index_without_refreshing_batch_mask(): + manager = object.__new__(MambaHybridCacheManagerV2) + manager.local_num_mamba_layers = 1 + manager._request_id_to_state_index = {123: 7} + manager.get_state_indices = MagicMock( + side_effect=AssertionError("state-index lookup must not refresh the dummy mask") + ) + transceiver = object.__new__(KvCacheTransceiverV2) + transceiver._kv_cache_manager = manager + transceiver._reuse_adapter = SimpleNamespace(tokens_per_block=32) + transceiver._page_table = SimpleNamespace(layer_groups=[]) + request = SimpleNamespace( + is_generation_only_request=lambda: False, + prompt_len=0, + py_request_id=123, + ) + + kv_slice = transceiver._create_kv_slice(request) + + assert kv_slice.mamba_state_index == 7 + manager.get_state_indices.assert_not_called() + + +@pytest.mark.parametrize( + "max_beam_width, has_connector, expected", + [ + (2, False, "max_beam_width > 1"), + (1, True, "kv_connector_manager"), + (2, True, "kv_connector_manager, max_beam_width > 1"), + ], +) +def test_v2_hybrid_incompatibility_fails_without_cpp_fallback( + max_beam_width, has_connector, expected +): + config = SimpleNamespace( + architectures=["Qwen3_5MoeForCausalLM"], + num_hidden_layers=2, + layer_types=["linear_attention", "full_attention"], + ) + model_config = SimpleNamespace( + pretrained_config=config, + sparse_attention_config=None, + ) + creator = object.__new__(KvCacheCreator) + creator._kv_connector_manager = object() if has_connector else None + creator._max_beam_width = max_beam_width + + with pytest.raises(NotImplementedError, match=expected): + creator._fallback_if_unsupported_kv_cache_manager_v2( + MambaHybridCacheManagerV2, model_config, KvCacheConfig() + ) + + def _make_mgr( max_batch_size=4, max_draft_len=2, enable_attention_dp=False, use_replay_state_update=False ): @@ -346,11 +927,436 @@ def test_non_mtp_pytorch_prepare_and_get_state_indices_flow(): ] +def test_v2_hybrid_prepare_expect_snapshot_points(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.enable_block_reuse = True + mgr.kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64, + additional_snapshot_offsets_from_start=[128, 999], + additional_snapshot_offsets_from_end=[0, 22, 999], + ), + ) + requests = [ + SimpleNamespace(prompt_len=150, expect_snapshot_points=[999]), + SimpleNamespace(prompt_len=128, expect_snapshot_points=[]), + SimpleNamespace(prompt_len=32, expect_snapshot_points=[]), + ] + + mgr.prepare_expect_snapshot_points(requests) + + assert [request.expect_snapshot_points for request in requests] == [ + [64, 128, 150], + [64, 106, 128], + [10, 32], + ] + + +def test_v2_hybrid_prepare_expect_snapshot_points_without_periodic_snapshots(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.enable_block_reuse = True + mgr.kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[128], + additional_snapshot_offsets_from_end=[0, 13], + ), + ) + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[]) + + mgr.prepare_expect_snapshot_points([request]) + + assert request.expect_snapshot_points == [128, 137, 150] + + +def test_mamba_snapshot_rule_count_deduplicates_and_filters_unreachable_points(): + config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[64, 65, 64], + additional_snapshot_offsets_from_end=[0, 32, 4096], + ), + ) + + assert _mamba_snapshot_rule_counts(config, 128, 32) == (4, 3) + + +def test_v2_hybrid_snapshot_sizing_scales_with_pp_and_explicit_rules(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.max_batch_size = 4 + mgr.mapping = Mapping(world_size=2, rank=0, tp_size=1, pp_size=2) + mgr.max_seq_len = 128 + mgr.tokens_per_block = 32 + mgr._num_reserved_dummy_slots = 0 + config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[64], + additional_snapshot_offsets_from_end=[0, 32], + ), + ) + + assert mgr._num_ssm_snapshots_for_capacity(512, config) == 24 + assert mgr._num_ssm_states_per_typical_request(128, config) == 4 + assert [desc.capacity for desc in mgr._typical_request_descs(128, config)] == [ + 32, + 32, + 32, + 32, + ] + + periodic_config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=48), + ) + assert mgr._num_ssm_states_per_typical_request(47, periodic_config) == 1 + assert mgr._num_ssm_states_per_typical_request(48, periodic_config) == 1 + + +def test_cpp_mamba_estimator_handles_disabled_snapshots_without_attention(): + manager = object.__new__(KVCacheManager) + manager._primary_pool_memory_bytes = 4096 + manager._secondary_pool_memory_bytes = 0 + manager.linear_attention_metadata = SimpleNamespace( + all_recurrent_states_bytes=64, + states_snapshot_interval=0, + ) + manager.max_attention_window_vec = [LinearCacheType.RECURRENT_STATES.value] + manager.get_cache_bytes_per_token = lambda: 0 + manager.mapping = SimpleNamespace(pp_size=1) + manager.max_batch_size = 4 + manager.max_seq_len = 128 + manager.tokens_per_block = 32 + manager.spec_config = None + + blocks = manager._calculate_max_num_blocks_for_linear_attention( + KvCacheConfig( + max_tokens=512, + enable_block_reuse=False, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=0), + ) + ) + + assert blocks[128] == (16, 0) + assert blocks[LinearCacheType.RECURRENT_STATES.value] == (5, 0) + + +def test_hybrid_mtp_layout_honors_explicit_base_partition(): + model_config = _hybrid_cache_sizing_model_config( + [ + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + ] + ) + spec_config = MTPDecodingConfig(max_draft_len=1) + expected_pp_layers = ([0, 1], [2, 3, 4]) + + for rank in range(2): + mapping = Mapping( + world_size=2, + rank=rank, + tp_size=1, + pp_size=2, + pp_partition=[2, 2], + ) + pp_layers, total_layers = get_pp_layers( + 5, + mapping, + spec_config=spec_config, + layer_mask=[True] * 5, + ) + assert total_layers == 5 + assert pp_layers == expected_pp_layers[rank] + + _, local_mamba_layers, local_attention_layers = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + ) + assert local_mamba_layers == 1 + assert local_attention_layers == rank + 1 + + for manager_cls in ( + CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, + ): + cache_cost = manager_cls.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=1, + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + spec_config=spec_config, + ) + assert cache_cost == (64 * (rank + 1), 2400) + + +def test_hybrid_separate_mtp_draft_estimator_has_no_mamba_state(): + model_config = _hybrid_cache_sizing_model_config( + [ + "linear_attention", + "full_attention", + "linear_attention", + "full_attention", + ] + ) + spec_config = MTPDecodingConfig(max_draft_len=1) + + for rank in range(2): + mapping = Mapping(world_size=2, rank=rank, tp_size=1, pp_size=2) + + target_cost = MambaHybridCacheManagerV2.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=1, + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + spec_config=spec_config, + use_separate_draft_kv_cache=True, + ) + assert target_cost == (64, 2400) + + _, local_mamba_layers, local_attention_layers = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=True, + ) + assert local_mamba_layers == 0 + assert local_attention_layers == rank + + draft_cost = MambaHybridCacheManagerV2.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=1, + kv_cache_config=KvCacheConfig(enable_block_reuse=False), + num_layers=1, + spec_config=spec_config, + is_draft=True, + ) + assert draft_cost == (64 * rank, 0) + + +@pytest.mark.parametrize( + ("spec_config", "enable_attention_dp", "expected_intercept"), + [ + (None, False, 320), + (MTPDecodingConfig(max_draft_len=4), True, 384), + ( + MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={1: 4, 2: 2, 3: 1}, + ), + True, + 576, + ), + ], +) +def test_v2_hybrid_estimator_counts_dummy_states_without_attention_capacity( + monkeypatch, spec_config, enable_attention_dp, expected_intercept +): + monkeypatch.setattr( + KVCacheManager, + "get_cache_size_per_token", + lambda *args, **kwargs: 11, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._get_local_mamba_cache_layout", + lambda *args, **kwargs: ( + SimpleNamespace(get_states_bytes_per_layer=lambda mapping: 64), + 1, + 1, + ), + ) + mapping = Mapping( + world_size=1, + tp_size=1, + pp_size=1, + enable_attention_dp=enable_attention_dp, + ) + + assert MambaHybridCacheManagerV2.get_cache_size_per_token( + object(), + mapping, + max_batch_size=4, + kv_cache_config=KvCacheConfig(), + spec_config=spec_config, + ) == (11, expected_intercept) + + +def test_v2_hybrid_attention_bound_is_snapshot_alignment_agnostic(): + model_config = _hybrid_cache_sizing_model_config(["linear_attention", "full_attention"]) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + + def estimate(interval): + return MambaHybridCacheManagerV2.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=2, + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=interval), + ), + tokens_per_block=32, + max_seq_len=128, + ) + + aligned = estimate(32) + unaligned = estimate(48) + + assert aligned[1] == unaligned[1] + + +def _base_attention_layer_configs(num_layers): + return [ + AttentionLayerConfig( + layer_id=LayerId(layer_idx), + buffers=[BufferConfig(role="key", size=256)], + ) + for layer_idx in range(num_layers) + ] + + +def test_v2_hybrid_typical_batch_splits_capacity_across_ssm_states_and_dummies(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 1 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + mgr._minimum_live_gpu_quota = lambda: 0 + kv_cache_config = KvCacheConfig( + enable_partial_reuse=True, + avg_seq_len=96, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=48, + additional_snapshot_offsets_from_start=[32], + additional_snapshot_offsets_from_end=[0], + ), + ) + mgr.kv_cache_config = kv_cache_config + constraints = [BatchDesc([KVCacheDesc(capacity=64, history_length=0)])] + base_layers = _base_attention_layer_configs(2) + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=base_layers, + constraints=constraints, + ) + + config = mgr._build_cache_config(base_config) + + assert isinstance(config.layers[0], SsmLayerConfig) + assert config.layers[1] is base_layers[1] + assert config.typical_step == BatchDesc( + [KVCacheDesc(capacity=32, history_length=31)] * 6 + + [KVCacheDesc(capacity=0, history_length=0)] + ) + assert config.constraints == [ + BatchDesc( + [ + KVCacheDesc(capacity=64, history_length=0), + KVCacheDesc(capacity=0, history_length=0), + ] + ) + ] + assert sum(kv.capacity for kv in config.typical_step.kv_caches) == 2 * 96 + assert not hasattr(config.typical_step.kv_caches[0], "num_ssm_slots") + + +def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.max_seq_len = 4096 + warnings_seen = [] + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager.logger.warning", + lambda message: warnings_seen.append(message), + ) + + capacity = mgr._get_typical_request_capacity(KvCacheConfig()) + + assert capacity == 2048 + assert len(warnings_seen) == 1 + assert "kv_cache_config.avg_seq_len" in warnings_seen[0] + assert "max_seq_len / 2=2048" in warnings_seen[0] + assert "workload's average total sequence length" in warnings_seen[0] + + +def test_v2_hybrid_rejects_quota_below_live_state_floor(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.local_num_mamba_layers = 1 + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr._num_reserved_dummy_slots = 1 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.pp_layers = [0, 1] + mgr.max_attention_window_vec = [128, 128] + mgr.max_num_tokens = 128 + mgr.enable_swa_scratch_reuse = False + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + mgr._attention_cache_bytes_per_token = lambda: 16 + mgr.kv_cache_config = KvCacheConfig(enable_partial_reuse=False) + minimum_quota = mgr._minimum_live_gpu_quota() + + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=minimum_quota - 1)], + layers=[], + ) + + with pytest.raises(ValueError, match="too small for live recurrent states"): + mgr._build_cache_config(base_config) + + +def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.local_num_mamba_layers = 1 + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr._num_reserved_dummy_slots = 1 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 1 + mgr.pp_layers = [0] + mgr.max_attention_window_vec = [128] + mgr.max_num_tokens = 128 + mgr.enable_swa_scratch_reuse = False + mgr.get_layer_bytes_per_token = lambda **kwargs: 0 + mgr._attention_cache_bytes_per_token = lambda: 0 + mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) + + assert mgr._minimum_live_gpu_quota() == 3 * (64 + 32) + + def test_cpp_hybrid_prepare_expect_snapshot_points(): mgr = object.__new__(CppMambaHybridCacheManager) + mgr.enable_block_reuse = True mgr.kv_cache_config = KvCacheConfig( enable_block_reuse=True, - mamba_state_cache_interval=64, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), ) mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=64) requests = [ @@ -368,21 +1374,137 @@ def test_cpp_hybrid_prepare_expect_snapshot_points(): ] -def test_cpp_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): +@pytest.mark.parametrize( + ("allocated_offsets", "context_current_position", "expected_offset"), + [ + ({9: 25}, 0, 25), + ({7: 17, 9: 25}, 0, 17), + ({7: 17, 9: 25}, 256, 25), + ], +) +def test_cpp_hybrid_state_indices_skip_context_placeholders( + allocated_offsets, context_current_position, expected_offset +): + """Capacity-limited chunks use the next real snapshot/final block.""" + null_index = torch.iinfo(torch.int32).max + block_offsets = torch.full((1, 1, 2, 10), null_index, dtype=torch.int32) + for logical_index, pool_offset in allocated_offsets.items(): + block_offsets[0, 0, 0, logical_index] = pool_offset + + request = SimpleNamespace( + py_request_id=0, + prompt_len=314, + is_context_finished=False, + context_current_position=context_current_position, + context_chunk_size=32, + prepopulated_prompt_len=0, + is_dummy=False, + ) mgr = object.__new__(CppMambaHybridCacheManager) + mgr.local_num_mamba_layers = 1 + mgr.requests = [request] + mgr.tokens_per_block = 32 + mgr.kv_cache_config = SimpleNamespace(enable_block_reuse=True) + mgr.impl = SimpleNamespace( + copy_batch_block_offsets=lambda *args: None, + get_cache_block_ids=lambda *args: [], + ) + mgr.host_block_offsets = block_offsets + mgr.recurrent_states_pool_index = 0 + mgr.blocks_per_window = {LinearCacheType.RECURRENT_STATES.value: (1264, 0)} + mgr._host_state_indices = torch.zeros(1, dtype=torch.int32) + mgr.cuda_state_indices = torch.zeros(1, dtype=torch.int32) + mgr._row_indices = torch.arange(1, dtype=torch.long) + mgr._request_id_to_state_index = {} + mgr._request_id_to_is_dummy = {} + mgr._dummy_request_mask = None + + mgr._setup_state_indices() + + assert mgr._host_state_indices.tolist() == [expected_offset] + assert mgr.cuda_state_indices.tolist() == [expected_offset] + assert mgr.get_state_indices([request.py_request_id], [False]) == [expected_offset] + + +def test_v2_block_reuse_commit_saves_ssm_snapshot_at_snapshot_point(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.enable_block_reuse = True + mgr.is_draft = False + mgr._augment_tokens_for_block_reuse = lambda tokens, request, start, end: tokens[start:end] + mgr._mark_context_position_as_history = MagicMock() + + token_ids = list(range(150)) + request = SimpleNamespace( + prompt_len=150, + context_current_position=137, + context_remaining_length=13, + expect_snapshot_points=[137], + is_dummy_request=False, + is_dummy=False, + py_request_id=0, + get_tokens=lambda beam_idx: token_ids, + ) + kv_cache = SimpleNamespace( + num_committed_tokens=0, + commit=MagicMock(), + stop_committing=MagicMock(), + ) + + mgr.try_commit_blocks(request, kv_cache) + + kv_cache.commit.assert_called_once_with(token_ids[:137]) + kv_cache.stop_committing.assert_not_called() + mgr._mark_context_position_as_history.assert_called_once_with(request, kv_cache) + + # The remaining suffix advances request history but must not publish a + # second attention/SSM snapshot beyond the configured boundary. + kv_cache.num_committed_tokens = 137 + request.context_current_position = 150 + request.context_remaining_length = 0 + mgr.try_commit_blocks(request, kv_cache) + + kv_cache.commit.assert_called_once_with(token_ids[:137]) + kv_cache.stop_committing.assert_called_once_with() + + +def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): + mgr = object.__new__(MambaHybridCacheManagerV2) + base_add_dummy_requests = mocker.patch.object( + KVCacheManagerV2, "add_dummy_requests", return_value=[] + ) + + mgr.add_dummy_requests([123], encoder_output_lens=[17]) + + assert base_add_dummy_requests.call_args.kwargs["encoder_output_lens"] == [17] + + +@pytest.mark.parametrize( + "manager_cls", + [ + MambaHybridCacheManagerV2, + CppMambaHybridCacheManager, + MixedMambaHybridCacheManager, + ], + ids=["v2", "cpp", "mixed"], +) +def test_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(manager_cls): + mgr = object.__new__(manager_cls) + mgr.enable_block_reuse = False mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) - request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) + request = SimpleNamespace(prompt_len=64, expect_snapshot_points=[64]) mgr.prepare_expect_snapshot_points([request]) assert request.expect_snapshot_points == [] -@pytest.mark.parametrize("interval", [0, -64, None]) -def test_cpp_hybrid_prepare_expect_snapshot_points_clears_for_invalid_interval(interval): +def test_cpp_hybrid_prepare_expect_snapshot_points_clears_for_disabled_interval(): mgr = object.__new__(CppMambaHybridCacheManager) - mgr.kv_cache_config = SimpleNamespace(enable_block_reuse=True) - mgr.linear_attention_metadata = SimpleNamespace(states_snapshot_interval=interval) + mgr.enable_block_reuse = True + mgr.kv_cache_config = SimpleNamespace( + enable_block_reuse=True, + mamba_state_config=SimpleNamespace(periodic_snapshot_interval=0), + ) request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[64]) mgr.prepare_expect_snapshot_points([request]) @@ -404,16 +1526,72 @@ def test_expect_snapshot_points_binding_round_trip(): assert request.expect_snapshot_points == [64, 128] +@skip_no_cuda +def test_v2_hybrid_pool_ratio_controls_allocated_memory(): + def allocated_memory(pool_ratio): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + mgr.ssm_bytes = 64 + mgr.conv_bytes = 32 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 1 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + + kv_cache_config = KvCacheConfig( + pool_ratio=pool_ratio, + enable_partial_reuse=False, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=64), + ) + mgr.kv_cache_config = kv_cache_config + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], + layers=_base_attention_layer_configs(2), + initial_pool_ratio=pool_ratio, + ) + config = mgr._build_cache_config(base_config) + runtime_manager = RuntimeKVCacheManager(config) + try: + statistics = runtime_manager._storage.get_statistics() + allocated_bytes = [ + int(stats.total) * sum(int(size) for size in stats.slot_size) + for stats in statistics + ] + return allocated_bytes, list(runtime_manager._current_gpu_ratio) + finally: + runtime_manager.shutdown() + + low_mamba_allocation, low_actual_ratio = allocated_memory([0.25, 0.75]) + high_mamba_allocation, high_actual_ratio = allocated_memory([0.75, 0.25]) + + assert low_actual_ratio == pytest.approx([0.25, 0.75]) + assert high_actual_ratio == pytest.approx([0.75, 0.25]) + assert high_mamba_allocation[0] > low_mamba_allocation[0] + assert high_mamba_allocation[1] < low_mamba_allocation[1] + + # --------------------------------------------------------------------------- -# CppMambaHybridCacheManager: recurrent-state snapshot pool sizing +# Cpp/V2 Mamba hybrid managers: recurrent-state allocation and reuse # -# Sized in KVCacheManager._calculate_max_num_blocks_for_linear_attention. -# Mirrors the MixedMambaCacheManager fix where each kind of padding sentinel -# (CUDA-graph dummy, plus one per draft length under spec decoding) must not -# evict live recurrent state. Wanli's fix made all sentinels share one slot -# in the Python manager (#13489); the C++ hybrid path instead reserves a -# dedicated slot per sentinel kind in the underlying pool — same invariant, -# different mechanism. These tests guard the pool sizing. +# The Cpp pool is sized in +# KVCacheManager._calculate_max_num_blocks_for_linear_attention. It reserves a +# dedicated slot for each padding sentinel kind so dummy requests cannot evict +# live recurrent state. The V2 tests cover unified-pool state views, slot +# bookkeeping, replay, and snapshot reuse. # --------------------------------------------------------------------------- @@ -421,12 +1599,13 @@ def _build_hybrid_with_mamba_layer( spec_config=None, max_batch_size=4, enable_block_reuse=False, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, is_estimating_kv_cache=False, dtype=DataType.HALF, mamba_layer_mask=None, attention_layer_mask=None, mamba_ssm_cache_dtype=torch.float16, + use_replay_state_update=False, ): """Construct a real CppMambaHybridCacheManager with one mamba layer + one full-attention layer so the parent KVCacheManager goes through the @@ -439,7 +1618,7 @@ def _build_hybrid_with_mamba_layer( kv_cache_config = KvCacheConfig( max_tokens=512, enable_block_reuse=enable_block_reuse, - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=periodic_snapshot_interval), ) return CppMambaHybridCacheManager( mamba_d_state=8, @@ -464,6 +1643,681 @@ def _build_hybrid_with_mamba_layer( layer_mask=attn_mask, is_estimating_kv_cache=is_estimating_kv_cache, dtype=dtype, + use_replay_state_update=use_replay_state_update, + ) + + +def _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + num_mamba_layers=1, + num_attention_layers=1, + num_kv_heads=4, + mapping=None, + spec_config=None, + use_replay_state_update=False, + enable_block_reuse=False, + enable_partial_reuse=True, + block_reuse_policy="all_reusable", + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_end=None, + enable_attention_dp=False, + enable_swa_scratch_reuse=False, + dtype=DataType.HALF, + conv_state_layout="x_b_c", +): + """Construct a real MambaHybridCacheManagerV2.""" + mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers + attn_mask = [False] * num_mamba_layers + [True] * num_attention_layers + if mapping is None: + mapping = Mapping( + world_size=1, + rank=0, + tp_size=1, + pp_size=1, + enable_attention_dp=enable_attention_dp, + ) + kv_cache_config = KvCacheConfig( + max_tokens=512, + enable_block_reuse=enable_block_reuse, + enable_partial_reuse=enable_partial_reuse, + block_reuse_policy=block_reuse_policy, + enable_swa_scratch_reuse=enable_swa_scratch_reuse, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=periodic_snapshot_interval, + additional_snapshot_offsets_from_end=list(additional_snapshot_offsets_from_end or []), + ), + dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", + ) + return MambaHybridCacheManagerV2( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=num_mamba_layers, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=num_attention_layers, + num_kv_heads=num_kv_heads, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=max_batch_size, + mapping=mapping, + spec_config=spec_config, + layer_mask=attn_mask, + vocab_size=1024, + use_replay_state_update=use_replay_state_update, + dtype=dtype, + conv_state_layout=conv_state_layout, + ) + + +def _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5): + """Spec config whose per-step token width is wider than draft depth. + + This mirrors parallel-draft style metadata closely enough for cache-manager + sizing without constructing the full speculative worker stack. + """ + return SimpleNamespace( + max_draft_len=max_draft_len, + max_total_draft_tokens=tokens_per_gen_step - 1, + tokens_per_gen_step=tokens_per_gen_step, + spec_dec_mode=SimpleNamespace(use_one_engine=lambda: False), + ) + + +def _assert_replay_layer_cache_uses_history_size(layer_cache, history_size): + assert layer_cache.old_x is not None + assert layer_cache.old_B is not None + assert layer_cache.old_dt is not None + assert layer_cache.old_dA_cumsum is not None + assert layer_cache.cache_buf_idx is not None + assert layer_cache.prev_num_accepted_tokens is not None + assert layer_cache.old_x.dim() == 5 + cache_size = layer_cache.temporal.shape[0] + assert layer_cache.old_x.shape[0] == cache_size + assert layer_cache.old_B.shape[0] == cache_size + assert layer_cache.old_dt.shape[0] == cache_size + assert layer_cache.old_dA_cumsum.shape[0] == cache_size + assert layer_cache.cache_buf_idx.shape[0] == cache_size + assert layer_cache.prev_num_accepted_tokens.shape[0] == cache_size + assert layer_cache.old_x.shape[1] == 2 + assert layer_cache.old_B.shape[1] == 2 + assert layer_cache.old_dt.shape[1] == 2 + assert layer_cache.old_dA_cumsum.shape[1] == 2 + assert layer_cache.old_x.shape[2] == history_size + assert layer_cache.old_B.shape[2] == history_size + assert layer_cache.old_dt.shape[-1] == history_size + assert layer_cache.old_dA_cumsum.shape[-1] == history_size + + +@skip_no_cuda +def test_v2_hybrid_allocates_mamba_state_and_dummy_indices(): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4) + try: + assert mgr.local_num_mamba_layers == 1 + assert len(mgr.all_ssm_states) == 1 + assert len(mgr.all_conv_states) == 1 + assert mgr.all_ssm_states[0].shape[1:] == torch.Size([4, 8, 8]) + assert mgr.all_conv_states[0].shape[1:] == torch.Size([48, 3]) + assert mgr.get_max_resource_count() == 4 + assert mgr.blocks_in_primary_pool > 0 + assert isinstance(mgr.check_invalid_values_in_kv_cache(), bool) + + requests = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False) + + assert len(requests) == 1 + indices = mgr.get_state_indices([123], [False]) + assert len(indices) == 1 + assert indices[0] >= 0 + assert mgr._request_id_to_is_dummy[123] + assert mgr.cuda_state_indices[0].item() == indices[0] + assert mgr.get_ssm_states(0).data_ptr() == mgr.all_ssm_states[0].data_ptr() + assert mgr.get_conv_states(0).data_ptr() == mgr.all_conv_states[0].data_ptr() + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_nvfp4_page_table_omits_ssm_block_scales(): + mgr = _build_v2_hybrid_with_mamba_layer(dtype=DataType.NVFP4) + try: + ssm_pool_id = mgr.impl.get_layer_group_id(LayerId(0)) + attention_pool_id = mgr.impl.get_layer_group_id(LayerId(1)) + + assert torch.count_nonzero(mgr.kv_cache_pool_pointers[ssm_pool_id, :, 1]) == 0 + assert torch.count_nonzero(mgr.kv_cache_pool_pointers[attention_pool_id, :, 1]) > 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_supports_pure_mamba_pp_rank(): + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=2, + num_attention_layers=0, + ) + try: + assert mgr.local_num_mamba_layers == 1 + assert mgr.blocks_in_primary_pool == 0 + assert mgr._attention_cache_bytes_per_token() == 0 + assert len(mgr.all_ssm_states) == 1 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_metadata_supports_attention_only_pp_rank(): + mapping = Mapping(world_size=2, rank=1, tp_size=1, pp_size=2) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=1, + mapping=mapping, + ) + try: + assert mgr.local_num_mamba_layers == 0 + mgr.add_dummy_requests([123], token_nums=[8], is_gen=False) + + metadata = Mamba2Metadata(max_batch_size=1, chunk_size=8) + seq_lens = torch.tensor([8], dtype=torch.int32) + metadata.prepare( + SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=1, + num_ctx_tokens=8, + kv_cache_manager=mgr, + request_ids=[123], + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int32) + ), + ) + ) + + assert metadata.state_indices[0].item() == 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_invalid_check_scans_distinct_attention_pools(): + mgr = _build_v2_hybrid_with_mamba_layer( + num_attention_layers=2, + num_kv_heads=[1, 2, 4], + ) + try: + mgr.check_invalid_values_in_kv_cache(fill_with_zero=True) + first_attention_buffer = mgr.get_buffers(1) + second_attention_buffer = mgr.get_buffers(2) + assert first_attention_buffer.data_ptr() != second_attention_buffer.data_ptr() + + for buffer in ( + second_attention_buffer, + mgr.all_ssm_states[0], + mgr.all_conv_states[0], + ): + buffer.flatten()[0] = torch.nan + assert mgr.check_invalid_values_in_kv_cache() + assert mgr.check_invalid_values_in_kv_cache(fill_with_zero=True) + assert not torch.isnan(buffer).any() + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_dummy_indices_keep_cuda_buffer_address(): + max_batch_size = 1 + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=max_batch_size, enable_attention_dp=True) + try: + request_ids = list(range(100, 100 + max_batch_size)) + mgr.add_dummy_requests( + request_ids, + token_nums=[8] * max_batch_size, + is_gen=False, + ) + state_indices_ptr = mgr.cuda_state_indices.data_ptr() + + new_requests = mgr.add_dummy_requests( + [ATTENTION_DP_DUMMY_REQUEST_ID], token_nums=[8], is_gen=False + ) + + assert len(new_requests) == 1 + expected_capacity = max_batch_size + mgr._num_reserved_dummy_slots + assert mgr.cuda_state_indices.shape[0] == expected_capacity + assert mgr._host_state_indices.shape[0] == expected_capacity + assert mgr.cuda_state_indices.data_ptr() == state_indices_ptr + assert ( + mgr.cuda_state_indices[0].item() + == mgr.get_state_indices([ATTENTION_DP_DUMMY_REQUEST_ID], [False])[0] + ) + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_reserves_every_persistent_dummy_slot(): + spec_config = MTPDecodingConfig( + max_draft_len=4, + draft_len_schedule={1: 4, 2: 2, 3: 1}, + ) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + enable_attention_dp=True, + ) + try: + runtime_draft_lengths = [4, 2, 1, 0] + cuda_graph_dummy_ids = [ + CUDA_GRAPH_DUMMY_REQUEST_ID - draft_len for draft_len in runtime_draft_lengths + ] + request_ids = [101, 102, 103, 104] + + assert mgr._num_reserved_dummy_slots == 5 + assert mgr.index_mapper.num_free_slots() == len(request_ids) + 5 + + assert ( + mgr.add_dummy_requests(request_ids, token_nums=[1] * len(request_ids), is_gen=False) + is not None + ) + for request_id, draft_len in zip(cuda_graph_dummy_ids, runtime_draft_lengths): + assert ( + mgr.add_dummy_requests( + [request_id], + is_gen=True, + max_num_draft_tokens=draft_len, + ) + is not None + ) + assert ( + mgr.add_dummy_requests([ATTENTION_DP_DUMMY_REQUEST_ID], token_nums=[1], is_gen=False) + is not None + ) + + all_request_ids = request_ids + cuda_graph_dummy_ids + [ATTENTION_DP_DUMMY_REQUEST_ID] + state_indices = mgr.get_state_indices(all_request_ids, [False] * len(all_request_ids)) + assert len(set(state_indices)) == len(all_request_ids) + assert mgr.index_mapper.num_free_slots() == 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_free_resources_drops_stale_state_index_mapping(): + mgr = _build_v2_hybrid_with_mamba_layer() + try: + request = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False)[0] + request_id = request.py_request_id + assert request_id in mgr._request_id_to_state_index + assert request_id in mgr._request_id_to_is_dummy + + # Move state-index preparation to another request before freeing the + # older one, as happens when an asynchronous transfer finishes late. + mgr.add_dummy_requests([456], token_nums=[8], is_gen=False) + mgr.free_resources(request) + + assert request_id not in mgr._request_id_to_state_index + assert request_id not in mgr._request_id_to_is_dummy + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_uses_upstream_min_snapshot_policy(): + mgr = _build_v2_hybrid_with_mamba_layer( + enable_block_reuse=True, + enable_partial_reuse=True, + ) + try: + assert mgr.block_reuse_policy is BlockReusePolicy.PER_REQUEST + assert mgr.kv_cache_config.enable_partial_reuse + assert mgr.kv_cache_manager_py_config.commit_min_snapshot + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_preserves_per_conversation_and_disables_periodic_snapshots(): + mgr = _build_v2_hybrid_with_mamba_layer( + enable_block_reuse=True, + block_reuse_policy="per_conversation", + periodic_snapshot_interval=64, + additional_snapshot_offsets_from_end=[0], + ) + try: + assert mgr.block_reuse_policy is BlockReusePolicy.PER_CONVERSATION + assert mgr.conversation_manager is not None + assert mgr.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 0 + assert mgr.kv_cache_manager_py_config.commit_min_snapshot + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[]) + mgr.prepare_expect_snapshot_points([request]) + assert request.expect_snapshot_points == [150] + finally: + mgr.shutdown() + + +def test_v2_hybrid_saves_conversation_plan_only_after_final_context_chunk(): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.enable_block_reuse = True + mgr.is_draft = False + mgr.block_reuse_policy = BlockReusePolicy.PER_CONVERSATION + events = [] + mgr._augment_tokens_for_block_reuse = lambda tokens, request, start, end: tokens[start:end] + mgr._mark_context_position_as_history = MagicMock() + mgr.conversation_manager = MagicMock() + mgr.conversation_manager.save_drop_plan.side_effect = lambda request, kv_cache: events.append( + "save" + ) + + request = SimpleNamespace( + py_request_id=7, + is_dummy_request=False, + context_current_position=128, + context_remaining_length=0, + expect_snapshot_points=[128], + prompt_len=128, + is_last_context_chunk=True, + get_tokens=lambda beam_idx: list(range(128)), + ) + kv_cache = SimpleNamespace( + is_active=True, + num_committed_tokens=0, + resize=MagicMock(return_value=True), + enable_swa_scratch_reuse=True, + ) + + def commit(tokens): + events.append("commit") + kv_cache.num_committed_tokens += len(tokens) + + kv_cache.commit = MagicMock(side_effect=commit) + kv_cache.stop_committing = MagicMock(side_effect=lambda: events.append("stop")) + mgr.kv_cache_map = {request.py_request_id: kv_cache} + batch = ScheduledRequests() + batch.append_context_request(request) + + mgr.update_context_resources(batch) + + kv_cache.commit.assert_called_once_with(list(range(128))) + kv_cache.stop_committing.assert_called_once_with() + mgr.conversation_manager.save_drop_plan.assert_called_once_with(request, kv_cache) + assert events == ["commit", "stop", "save"] + assert not kv_cache.enable_swa_scratch_reuse + + +@skip_no_cuda +def test_v2_hybrid_mamba_state_views_use_logical_slots(): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4, num_mamba_layers=2) + try: + assert len(mgr.all_ssm_states) == 2 + assert len(mgr.all_conv_states) == 2 + + ssm_slots = mgr.all_ssm_states[0].shape[0] + conv_slots = mgr.all_conv_states[0].shape[0] + assert all(t.shape[0] == ssm_slots for t in mgr.all_ssm_states) + assert all(t.shape[0] == conv_slots for t in mgr.all_conv_states) + assert ssm_slots == conv_slots + + local_layer_ids = [mgr.layer_offsets[layer_id] for layer_id in mgr.mamba_pp_layers] + for local_layer_idx, ssm_state, conv_state in zip( + local_layer_ids, mgr.all_ssm_states, mgr.all_conv_states + ): + layer_id = LayerId(local_layer_idx) + ssm_scale = mgr.impl.get_page_index_scale(layer_id, MambaRole.SSM_STATE) + conv_scale = mgr.impl.get_page_index_scale(layer_id, MambaRole.CONV_STATE) + assert ssm_state.stride(0) == ssm_state[0].numel() * ssm_scale + assert conv_state.stride(0) == conv_state[0].numel() * conv_scale + assert ( + ssm_state.shape[0] + == ( + mgr.impl.get_page_index_upper_bound(layer_id, MambaRole.SSM_STATE) + + ssm_scale + - 1 + ) + // ssm_scale + ) + assert ( + conv_state.shape[0] + == ( + mgr.impl.get_page_index_upper_bound(layer_id, MambaRole.CONV_STATE) + + conv_scale + - 1 + ) + // conv_scale + ) + + mgr.add_dummy_requests([123, 456], token_nums=[8, 8], is_gen=False) + indices = mgr.get_state_indices([123, 456], [False, False]) + assert all(0 <= index < ssm_slots for index in indices) + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_swa_scratch_keeps_ssm_placeholder_rows(): + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + enable_swa_scratch_reuse=True, + ) + try: + request_id = 123 + mgr.add_dummy_requests([request_id], token_nums=[8], is_gen=False) + block_offsets = torch.zeros( + mgr.num_attention_op_pools, + 1, + 2, + mgr.max_blocks_per_seq, + dtype=torch.int32, + device="cuda", + ) + + mgr.copy_batch_block_offsets( + block_offsets, + [request_id], + beam_width=1, + num_contexts=1, + num_seqs=1, + ) + torch.cuda.synchronize() + + assert mgr.num_attention_op_pools == mgr.num_local_layers + assert mgr.kv_cache_pool_mapping.shape[0] == mgr.num_local_layers + assert mgr.kv_cache_pool_pointers[0, 0].item() != 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_disagg_page_table_preserves_lifecycle_indices(): + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=4, num_mamba_layers=2) + try: + page_table = build_page_table_from_manager(mgr) + + assert len(page_table.layer_groups) == mgr.impl._storage.num_life_cycles + assert isinstance(page_table.layer_groups[0], MambaLayerGroup) + assert isinstance(page_table.layer_groups[1], AttentionLayerGroup) + + requests = mgr.add_dummy_requests([123], token_nums=[64], is_gen=False) + assert len(requests) == 1 + attention_blocks = list( + mgr.kv_cache_map[123].get_aggregated_page_indices(1, valid_only=True) + ) + assert attention_blocks + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_disagg_page_table_uses_qwen3_next_conv_sections(): + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + conv_state_layout="q_k_v", + ) + try: + page_table = build_page_table_from_manager(mgr) + mamba_group = page_table.layer_groups[0] + + assert isinstance(mamba_group, MambaLayerGroup) + d_conv_m1 = mgr.conv_state_shape[1] + conv_elem_size = mgr.all_conv_states[0].element_size() + assert mamba_group.conv_section_bytes == [ + dim * d_conv_m1 * conv_elem_size for dim in mgr.conv_section_dims + ] + assert mgr.conv_section_dims == [8, 8, 32] + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_intermediate_states_size_by_tokens_per_gen_step(): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + ) + try: + assert mgr.intermediate_ssm_states.shape[2] == 5 + assert mgr.intermediate_conv_states.shape[2] == 5 + layer_cache = mgr.mamba_layer_cache(0) + assert layer_cache.intermediate_ssm.data_ptr() == ( + mgr.intermediate_ssm_states[0].data_ptr() + ) + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_hybrid_static_dynamic_tree_capacity(): + spec_config = MTPDecodingConfig( + max_draft_len=6, + max_total_draft_tokens=31, + use_dynamic_tree=True, + dynamic_tree_max_topK=10, + ) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + ) + try: + assert mgr.intermediate_ssm_states.shape[2] == 32 + assert mgr.intermediate_conv_states.shape[2] == 32 + assert mgr._kv_reserve_draft_tokens == 31 + assert mgr._num_reserved_dummy_slots == 1 + assert not mgr.use_replay_state_update + finally: + mgr.shutdown() + + +@skip_no_cuda +@pytest.mark.parametrize( + "builder", + [_build_hybrid_with_mamba_layer, _build_v2_hybrid_with_mamba_layer], + ids=["cpp", "v2"], +) +def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = builder( + max_batch_size=4, + spec_config=spec_config, + use_replay_state_update=True, + ) + try: + replay_metadata = mgr.get_replay_state_update_metadata() + assert mgr.use_replay_state_update is True + assert replay_metadata is not None + assert replay_metadata.replay_step_width == spec_config.tokens_per_gen_step + assert replay_metadata.replay_history_size == max( + MIN_REPLAY_HISTORY_SIZE, spec_config.tokens_per_gen_step + ) + layer_cache = mgr.mamba_layer_cache(0) + _assert_replay_layer_cache_uses_history_size( + layer_cache, replay_metadata.replay_history_size + ) + finally: + mgr.shutdown() + + +def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.local_num_mamba_layers = 1 + mgr._request_id_to_state_index = { + 100: 0, + 101: 1, + 102: 2, + 103: 3, + } + mgr._request_id_to_is_dummy = { + 100: False, + 101: False, + 102: True, + 103: False, + } + mgr._dummy_request_mask = torch.zeros(4, dtype=torch.bool) + mgr._dummy_request_mask_host = torch.zeros(4, dtype=torch.bool) + mgr._use_replay_state_update = True + mgr.replay_step_width = 5 + mgr.replay_history_size = 16 + mgr.prev_num_accepted_tokens = torch.full((4,), 13, dtype=torch.int32) + mgr.cache_buf_idx = torch.ones(4, dtype=torch.int32) + mgr.intermediate_state_indices = torch.arange(4, dtype=torch.int32) + mgr.all_ssm_states = [] + mgr.all_conv_states = [torch.empty(0)] + mgr.intermediate_conv_states = torch.empty(0) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + lambda *args, **kwargs: None, + ) + + request_ids = [100, 101, 102, 103] + state_indices = torch.tensor( + mgr.get_state_indices(request_ids, [False, False, False, True]), + dtype=torch.int32, + ) + assert mgr._dummy_request_mask.tolist() == [False, False, True, True] + + mgr.update_mamba_states( + SimpleNamespace(num_seqs=4, num_contexts=1), + torch.tensor([1, 3, 3, 3], dtype=torch.int32), + state_indices=state_indices, + ) + + assert mgr.prev_num_accepted_tokens.tolist() == [13, 3, 13, 13] + assert mgr.cache_buf_idx.tolist() == [1, 0, 1, 1] + + +def test_v2_hybrid_dynamic_tree_promotes_accepted_leaf_state(monkeypatch): + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.local_num_mamba_layers = 1 + mgr._use_replay_state_update = False + mgr.intermediate_state_indices = torch.arange(2, dtype=torch.int32) + mgr.all_ssm_states = [torch.empty(0)] + mgr.all_conv_states = [torch.empty(0)] + mgr.intermediate_ssm_states = torch.empty((1, 2, 8)) + mgr.intermediate_conv_states = torch.empty((1, 2, 8)) + + promoted_positions = [] + + def capture_promoted_position(_dst, _src, _src_indices, positions, _dst_indices): + promoted_positions.append(positions.clone()) + + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + capture_promoted_position, + ) + + mgr.update_mamba_states( + SimpleNamespace(num_seqs=3, num_contexts=1), + torch.tensor([1, 2, 3], dtype=torch.int32), + state_indices=torch.tensor([9, 10, 11], dtype=torch.int32), + accepted_leaf_positions=torch.tensor([4, 7], dtype=torch.int64), + ) + + assert len(promoted_positions) == 2 + assert all( + torch.equal(positions, torch.tensor([4, 7], dtype=torch.int32)) + for positions in promoted_positions ) @@ -639,7 +2493,7 @@ def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): """With block reuse enabled, the block-reuse branch must not drop the live-state + CUDA-graph-padding floor. - With max_batch_size=4, mamba_state_cache_interval=256, max_tokens=512: + With max_batch_size=4, periodic_snapshot_interval=256, max_tokens=512: naive: max_snapshots = 512 // 256 = 2 (drops live-state floor!) fixed: max_snapshots = max(2, 4 + 1) = 5 """ @@ -648,7 +2502,7 @@ def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): spec_config=None, max_batch_size=max_batch_size, enable_block_reuse=True, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, ) recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] assert recurrent_primary >= max_batch_size + 1, ( @@ -671,7 +2525,7 @@ def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): spec_config=None, max_batch_size=max_batch_size, enable_block_reuse=True, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, is_estimating_kv_cache=True, ) recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] @@ -693,9 +2547,8 @@ def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): # - call the real parent KVCacheManager with the union layer_mask and # num_layers=num_layers (not mamba_num_layers + num_layers), # - skip allocating any mamba-only state, and -# - leave self.requests = [] so the guards on prepare_resources / -# update_mamba_states / _setup_state_indices can no-op without touching -# uninitialized state. +# - leave self.requests = [] so Mamba-only hooks and metadata preparation +# can no-op without touching uninitialized state. # # We exercise the same Python branch with world_size=1 (so the real C++ # KVCacheManager init doesn't need MPI) and a layer mask that contains zero @@ -705,7 +2558,7 @@ def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): def _build_zero_mamba_hybrid( enable_block_reuse=False, - mamba_state_cache_interval=256, + periodic_snapshot_interval=256, ): """Construct a real CppMambaHybridCacheManager whose this-rank slice has no mamba layers. world_size=1 / pp_size=1 keeps the real parent @@ -722,7 +2575,7 @@ def _build_zero_mamba_hybrid( kv_cache_config = KvCacheConfig( max_tokens=128, enable_block_reuse=enable_block_reuse, - mamba_state_cache_interval=mamba_state_cache_interval, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=periodic_snapshot_interval), ) mgr = CppMambaHybridCacheManager( @@ -756,11 +2609,11 @@ def _build_zero_mamba_hybrid( @skip_no_cuda def test_cpp_hybrid_zero_local_mamba_layers(): """End-to-end: real parent KVCacheManager + real early-exit. Verifies - early-exit invariants on the manager state AND that the three guarded - methods no-op without raising on uninitialized mamba-only state.""" + early-exit invariants on the manager state and that Mamba hooks and + metadata preparation do not touch uninitialized Mamba-only state.""" mgr = _build_zero_mamba_hybrid( enable_block_reuse=True, - mamba_state_cache_interval=64, + periodic_snapshot_interval=64, ) # Early-exit indicators. @@ -819,3 +2672,20 @@ def test_cpp_hybrid_zero_local_mamba_layers(): mgr.prepare_resources(empty_batch) # super() runs, then guard returns mgr.update_mamba_states(attn_metadata=None, num_accepted_tokens=None, state_indices=None) mgr._setup_state_indices() + + metadata = Mamba2Metadata(max_batch_size=1, chunk_size=8) + seq_lens = torch.tensor([8], dtype=torch.int32) + metadata.prepare( + SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=1, + num_ctx_tokens=8, + kv_cache_manager=mgr, + request_ids=[123], + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int32) + ), + ) + ) + assert metadata.state_indices[0].item() == 0 diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index c2a9a56a87bb..d878a50c3c24 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -23,7 +23,7 @@ _MLA_KV_CACHE_REUSE_SUPPORTED_SM_VERSIONS, ) from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType -from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig +from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, ContextChunkingPolicy from tensorrt_llm.quantization import QuantAlgo @@ -162,7 +162,11 @@ def _make_llm_args(): enable_partial_reuse=False, tokens_per_block=32, max_attention_window=None, - mamba_state_cache_interval=1, + mamba_state_config=SimpleNamespace( + periodic_snapshot_interval=256, + additional_snapshot_offsets_from_start=[], + additional_snapshot_offsets_from_end=[], + ), ) scheduler_config = SimpleNamespace( context_chunking_policy=None, @@ -210,6 +214,8 @@ def _run_create_py_executor( enable_flash_mla=False, model_max_seq_len=128, enable_chunked_prefill=False, + is_hybrid_linear_model=False, + ctx_chunk_configs=None, ): """Execute create_py_executor with mocked dependencies and return MLA runtime flags. @@ -225,6 +231,8 @@ def _run_create_py_executor( enable_flash_mla: Whether to emulate the FlashMLA block-size override. model_max_seq_len: Effective sequence length reported by the model engine. enable_chunked_prefill: Whether to request MLA chunked prefill support. + is_hybrid_linear_model: Whether to emulate a hybrid linear model. + ctx_chunk_configs: Optional list that receives the executor chunk config. Returns: Tuple of (kv_cache_reuse_flag, runtime_cache_reuse_flag, @@ -263,7 +271,11 @@ def _run_create_py_executor( monkeypatch.setattr(py_executor_creator, "_adjust_torch_mem_fraction", lambda: None) monkeypatch.setattr(py_executor_creator, "log_memory_usage", lambda *args, **kwargs: None) monkeypatch.setattr(py_executor_creator, "is_mla", lambda _: True) - monkeypatch.setattr(py_executor_creator, "is_hybrid_linear", lambda _: False) + monkeypatch.setattr( + py_executor_creator, + "is_hybrid_linear", + lambda _: is_hybrid_linear_model, + ) monkeypatch.setattr(py_executor_creator, "get_sm_version", lambda: sm_version) monkeypatch.setattr(py_executor_creator, "KvCacheCreator", _DummyKvCacheCreator) @@ -288,6 +300,8 @@ def _create_model_engine(**kwargs): monkeypatch.setattr(py_executor_creator, "PyTorchModelEngine", _create_model_engine) def _create_py_executor_instance(**kwargs): + if ctx_chunk_configs is not None: + ctx_chunk_configs.append(kwargs["ctx_chunk_config"]) return _DummyPyExecutor( resources=kwargs["resources"], model_engine=kwargs["model_engine"], @@ -416,6 +430,19 @@ def test_explicit_transceiver_buffer_size_is_preserved(monkeypatch): assert config.max_tokens_in_buffer == 256 +def test_hybrid_force_chunk_uses_block_alignment_unit(monkeypatch): + ctx_chunk_configs = [] + _run_create_py_executor( + monkeypatch, + sm_version=90, + kv_cache_quant_algo=QuantAlgo.NO_QUANT, + is_hybrid_linear_model=True, + ctx_chunk_configs=ctx_chunk_configs, + ) + + assert ctx_chunk_configs == [(ContextChunkingPolicy.FORCE_CHUNK, 32)] + + @pytest.mark.parametrize("sm_version", _MLA_CHUNKED_PREFILL_SUPPORTED_SM_VERSIONS) def test_mla_supported_configuration_preserves_chunked_prefill(monkeypatch, sm_version): """Verify every supported MLA SM preserves chunked prefill when requested.""" diff --git a/tests/unittest/_torch/modeling/test_modeling_multimodal.py b/tests/unittest/_torch/modeling/test_modeling_multimodal.py index 112b665dfb8c..7a57ba500eca 100644 --- a/tests/unittest/_torch/modeling/test_modeling_multimodal.py +++ b/tests/unittest/_torch/modeling/test_modeling_multimodal.py @@ -628,6 +628,7 @@ def get_hybrid_kv_cache_manager( } mamba_params = extract_mamba_kv_cache_params(text_config) + mamba_layer_mask, full_attention_layer_mask = mamba_params.get_layer_masks() if mamba_params.dtype not in dtype_map: raise ValueError( f"Unsupported dtype for hybrid cache manager: " @@ -640,7 +641,7 @@ def get_hybrid_kv_cache_manager( head_dim = text_config.hidden_size // text_config.num_attention_heads # CppMambaHybridCacheManager reads Pydantic-only fields - # (mamba_state_cache_interval, enable_block_reuse) so we have to + # (mamba_state_config, enable_block_reuse) so we have to # construct the llmapi.llm_args.KvCacheConfig here, not the C++ # bindings KvCacheConfig that the standard KVCacheManager path uses. kv_cache_config = PyKvCacheConfig(max_tokens=num_blocks * tokens_per_block) @@ -654,15 +655,15 @@ def get_hybrid_kv_cache_manager( mamba_params.n_groups, mamba_params.head_dim, mamba_params.num_mamba_layers, - mamba_params.mamba_layer_mask, + mamba_layer_mask, mamba_params.dtype, mamba_params.mamba_ssm_cache_dtype, # kv cache parameters (positional) kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, # kw-only - num_layers=mamba_params.num_full_attention_layers, - layer_mask=mamba_params.full_attention_layer_mask, + num_layers=sum(full_attention_layer_mask), + layer_mask=full_attention_layer_mask, num_kv_heads=text_config.num_key_value_heads, head_dim=head_dim, tokens_per_block=tokens_per_block, diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py index 33cf3d635661..33a3667d0db4 100644 --- a/tests/unittest/_torch/speculative/test_eagle3.py +++ b/tests/unittest/_torch/speculative/test_eagle3.py @@ -22,6 +22,8 @@ from tensorrt_llm._torch.pyexecutor.py_executor_creator import \ _extend_full_attention_windows_for_spec_decode from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSpecMetadata +from tensorrt_llm._torch.speculative.mtp_dynamic_tree import \ + MTPEagleDynamicTreeWorker from tensorrt_llm.executor.request import LoRARequest from tensorrt_llm.llmapi import (CudaGraphConfig, Eagle3DecodingConfig, KvCacheConfig, MoeConfig, MTPDecodingConfig) @@ -30,6 +32,77 @@ sys.path.append(os.path.join(os.path.dirname(__file__), '..')) +def test_dynamic_tree_metadata_forces_target_mask_prepare_each_step() -> None: + metadata = TrtllmAttentionMetadata( + seq_lens=None, + seq_lens_kv=None, + num_contexts=0, + max_num_requests=1, + max_num_tokens=1, + max_seq_len=1, + ) + common_kwargs = dict( + batch_size=0, + is_spec_decoding_enabled=False, + is_spec_dec_tree=True, + max_draft_len=1, + max_total_draft_tokens=1, + ) + + metadata.update_spec_dec_param(is_spec_dec_dynamic_tree=True, + **common_kwargs) + assert metadata.force_prepare_spec_dec_tree_mask + + metadata.update_spec_dec_param(is_spec_dec_dynamic_tree=False, + **common_kwargs) + assert not metadata.force_prepare_spec_dec_tree_mask + + +def test_mtp_dynamic_tree_relocation_uses_full_attention_window( + monkeypatch: pytest.MonkeyPatch) -> None: + worker = object.__new__(MTPEagleDynamicTreeWorker) + worker._kv_head_dim_bytes = 256 + worker._accepted_draft_indices_tensor = torch.tensor([[0, 1], [2, -1]], + dtype=torch.int32) + worker._num_accepted_tokens_buf = torch.tensor([2, 1], dtype=torch.int32) + + attention_pool_pointers = object() + attention_block_offsets = object() + cache_manager = SimpleNamespace( + num_kv_heads_per_layer=[0, 8, 0, 8], + kv_cache_pool_mapping=[[0, 0], [2, 0], [1, 0], [2, 1]], + kv_cache_pool_pointers=[object(), + object(), attention_pool_pointers], + max_attention_window_vec=[None], + max_seq_len=8192, + max_total_draft_tokens=31, + max_blocks_per_seq=256, + tokens_per_block=32, + ) + attention_metadata = SimpleNamespace( + kv_cache_manager=cache_manager, + kv_lens_cuda=torch.tensor([128, 256], dtype=torch.int32), + kv_cache_block_offsets=[object(), + object(), attention_block_offsets], + ) + update_op = MagicMock() + monkeypatch.setattr( + torch.ops.tensorrt_llm, + "update_kv_cache_draft_token_location_2d", + update_op, + ) + + worker._relocate_kv_eagerly(attention_metadata, batch_size=2) + + update_op.assert_called_once() + args = update_op.call_args.args + assert args[4] == 2 + assert args[5] == 8 + assert args[8] == cache_manager.max_seq_len + assert args[9] is attention_pool_pointers + assert args[10] is attention_block_offsets + + def test_eagle3_draft_kv_cache_uses_full_window_when_draft_has_no_swa() -> None: kv_cache_config = KvCacheConfig(max_attention_window=[128, 131072]) draft_pretrained_config = SimpleNamespace(num_hidden_layers=3) diff --git a/tests/unittest/disaggregated/region/test_page.py b/tests/unittest/disaggregated/region/test_page.py index 2e5d76782ba4..2c93bb5f7fea 100644 --- a/tests/unittest/disaggregated/region/test_page.py +++ b/tests/unittest/disaggregated/region/test_page.py @@ -43,15 +43,25 @@ def test_physical_pool_construction(): assert pool.base_address == 0x10000 assert pool.slot_bytes == 256 assert pool.num_slots == 4 + assert pool.slot_stride_bytes == 256 + assert pool.layer_stride_bytes == 1024 def test_physical_pool_roundtrip(): - pool = PhysicalPool(base_address=0x10000, slot_bytes=256, num_slots=4) + pool = PhysicalPool( + base_address=0x10000, + slot_bytes=256, + num_slots=4, + slot_stride_bytes=2048, + layer_stride_bytes=512, + ) d = pool.to_dict() restored = PhysicalPool.from_dict(d) assert restored.base_address == pool.base_address assert restored.slot_bytes == pool.slot_bytes assert restored.num_slots == pool.num_slots + assert restored.slot_stride_bytes == pool.slot_stride_bytes + assert restored.layer_stride_bytes == pool.layer_stride_bytes def test_pool_view_roundtrip(): diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 9cd9c1698f1e..69fb84b26887 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -15,10 +15,12 @@ import numpy as np import pytest +import torch from tensorrt_llm._torch.disaggregation.base.region import MemRegionGroup, SpecRegion from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( KVRegionExtractorV1, + _build_v2_mamba_state_pool, build_page_table, build_page_table_from_manager, ) @@ -30,6 +32,7 @@ get_num_layer_groups, get_num_layers, get_physical_pool, + get_slot_address, get_unique_layers, ) from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import Role @@ -581,8 +584,20 @@ def test_v2_builder_validates_role_mapper_declaration(): def test_mamba_layer_group_serialization(): from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup, PhysicalPool - conv_pool = PhysicalPool(base_address=1000, slot_bytes=128, num_slots=10) - ssm_pool = PhysicalPool(base_address=8000, slot_bytes=256, num_slots=8) + conv_pool = PhysicalPool( + base_address=1000, + slot_bytes=128, + num_slots=10, + slot_stride_bytes=512, + layer_stride_bytes=256, + ) + ssm_pool = PhysicalPool( + base_address=8000, + slot_bytes=256, + num_slots=8, + slot_stride_bytes=1024, + layer_stride_bytes=512, + ) mlg = MambaLayerGroup( pool_group_idx=1, mamba_layer_offsets={10: 0, 11: 1, 12: 2}, @@ -602,14 +617,128 @@ def test_mamba_layer_group_serialization(): assert isinstance(restored, MambaLayerGroup) assert restored.mamba_layer_offsets == {10: 0, 11: 1, 12: 2} assert restored.conv_states.base_address == 1000 - assert restored.conv_states.slot_bytes == 128 - assert restored.conv_states.num_slots == 10 + assert restored.conv_states.slot_stride_bytes == 512 + assert get_slot_address(restored.conv_states, 3) == 1000 + 3 * 512 assert restored.ssm_states.base_address == 8000 - assert restored.ssm_states.slot_bytes == 256 - assert restored.ssm_states.num_slots == 8 + assert restored.ssm_states.slot_stride_bytes == 1024 assert restored.conv_section_bytes == [512, 256, 256] assert restored.ssm_bytes_per_head == 128 + legacy_pool = PhysicalPool.from_dict({"base_address": 1000, "slot_bytes": 128, "num_slots": 10}) + assert legacy_pool.slot_stride_bytes == legacy_pool.slot_bytes + assert legacy_pool.layer_stride_bytes == legacy_pool.num_slots * legacy_pool.slot_bytes + + +def test_v2_mamba_state_pool_uses_affine_layer_and_slot_strides(): + num_slots = 3 + state_bytes = 64 + storage = torch.empty((num_slots, 4, state_bytes), dtype=torch.uint8) + states = [storage[:, 0, :], storage[:, 2, :]] + + pool = _build_v2_mamba_state_pool(states) + + assert pool.base_address == states[0].data_ptr() + assert pool.slot_bytes == state_bytes + assert pool.num_slots == num_slots + assert pool.slot_stride_bytes == 4 * state_bytes + assert pool.layer_stride_bytes == 2 * state_bytes + + +def test_v2_mamba_single_layer_pool_preserves_shared_role_footprint(): + state_bytes = 64 + storage = torch.empty((3, 2, state_bytes), dtype=torch.uint8) + + pool = _build_v2_mamba_state_pool([storage[:, 1, :]]) + + assert pool.slot_bytes == state_bytes + assert pool.slot_stride_bytes == 2 * state_bytes + assert pool.layer_stride_bytes == 2 * state_bytes + + +def test_v2_mamba_state_pool_rejects_non_affine_layer_offsets(): + storage = torch.empty((3, 6, 64), dtype=torch.uint8) + states = [storage[:, 0, :], storage[:, 2, :], storage[:, 5, :]] + + with pytest.raises(ValueError, match="uniform layer stride"): + _build_v2_mamba_state_pool(states) + + +def test_v2_mamba_registration_uses_coalesced_physical_pool(): + from tensorrt_llm._torch.disaggregation.resource.page import ( + KVCachePageTable, + MambaLayerGroup, + PhysicalPool, + PhysicalPoolGroup, + ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + state_bytes = 64 + num_layers = 2 + num_slots = 8 + # Equal-sized SSM and convolution states share one interleaved V2 pool. + physical_slot_bytes = state_bytes * num_layers * 2 + physical_pool = PhysicalPool( + base_address=1000, + slot_bytes=physical_slot_bytes, + num_slots=num_slots, + ) + mamba_group = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=PhysicalPool( + base_address=1000 + state_bytes, + slot_bytes=state_bytes, + num_slots=num_slots, + slot_stride_bytes=physical_slot_bytes, + layer_stride_bytes=2 * state_bytes, + ), + ssm_states=PhysicalPool( + base_address=1000, + slot_bytes=state_bytes, + num_slots=num_slots, + slot_stride_bytes=physical_slot_bytes, + layer_stride_bytes=2 * state_bytes, + ), + ) + page_table = KVCachePageTable( + tokens_per_block=16, + layer_groups=[mamba_group], + pool_groups=[PhysicalPoolGroup(pools=[physical_pool])], + ) + + assert get_unique_pool_memory_descs(page_table, device_id=3) == [ + (1000, physical_slot_bytes * num_slots, 3, "kv_cache_memory_pool0") + ] + + +def test_legacy_mamba_registration_uses_layer_major_pools(): + from tensorrt_llm._torch.disaggregation.resource.page import ( + KVCachePageTable, + MambaLayerGroup, + PhysicalPool, + ) + from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs + + num_layers = 3 + conv_pool = PhysicalPool(base_address=1000, slot_bytes=128, num_slots=10) + ssm_pool = PhysicalPool(base_address=8000, slot_bytes=256, num_slots=8) + mamba_group = MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={10: 0, 11: 1, 12: 2}, + conv_states=conv_pool, + ssm_states=ssm_pool, + ) + page_table = KVCachePageTable( + tokens_per_block=16, + layer_groups=[mamba_group], + pool_groups=[], + ) + + assert get_unique_pool_memory_descs(page_table, device_id=3) == [ + (1000, num_layers * conv_pool.num_slots * conv_pool.slot_bytes, 3, "kv_cache_memory_pool0"), + (8000, num_layers * ssm_pool.num_slots * ssm_pool.slot_bytes, 3, "kv_cache_memory_pool1"), + ] + def test_mixed_page_table_serialization(): import numpy as np diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 8688d75f9a53..dd701a0bc5ba 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -17,6 +17,7 @@ import uuid from typing import Dict, List +import numpy as np import pytest import torch @@ -29,13 +30,18 @@ import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams +from tensorrt_llm._torch.disaggregation.native.mixers.ssm import peer +from tensorrt_llm._torch.disaggregation.resource import page from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, LlmRequestType, ) -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + MixedMambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp @@ -189,8 +195,14 @@ def _init(rank): return results -def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=False): - """Create MixedMambaHybridCacheManagers for all TP ranks (PP=1). +def _create_managers( + tp, + max_batch_size=MAX_BATCH_SIZE, + enable_attention_dp=False, + use_v2=False, + conv_state_layout="x_b_c", +): + """Create Mamba hybrid cache managers for all TP ranks (PP=1). Layer 0 is a dummy attention layer required by page table infrastructure. Layers 1..NUM_MAMBA_LAYERS are mamba layers under test. @@ -200,7 +212,16 @@ def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=Fals mapping = Mapping( world_size=tp, rank=rank, tp_size=tp, pp_size=1, enable_attention_dp=enable_attention_dp ) - mgr = MixedMambaHybridCacheManager( + manager_cls = MambaHybridCacheManagerV2 if use_v2 else MixedMambaHybridCacheManager + manager_kwargs = ( + { + "is_disagg": True, + "conv_state_layout": conv_state_layout, + } + if use_v2 + else {} + ) + mgr = manager_cls( mamba_d_state=MAMBA_D_STATE, mamba_d_conv=MAMBA_D_CONV, mamba_num_heads=MAMBA_NUM_HEADS, @@ -226,19 +247,83 @@ def _create_managers(tp, max_batch_size=MAX_BATCH_SIZE, enable_attention_dp=Fals max_batch_size=max_batch_size, mapping=mapping, dtype=DataType.FLOAT, + **manager_kwargs, ) managers.append(mgr) return managers +def _mamba_layer_ids(manager): + if isinstance(manager, MambaHybridCacheManagerV2): + return manager.mamba_layer_offsets + return manager._impl.mamba_layer_offsets + + +def _mamba_state_slot(manager, request_id): + if isinstance(manager, MambaHybridCacheManagerV2): + return manager._request_id_to_state_index[request_id] + return manager.mamba_cache_index[request_id] + + +def _zero_mamba_states(manager): + for layer_idx in _mamba_layer_ids(manager): + manager.get_conv_states(layer_idx).zero_() + manager.get_ssm_states(layer_idx).zero_() + + +def test_mamba_policy_layer_major_v1_ptrs(): + pool = page.PhysicalPool(base_address=100, slot_bytes=10, num_slots=8) + + ptrs = peer.MambaPolicy._build_layer_ptrs( + pool=pool, + layer_offsets={1: 0, 2: 1}, + overlapping_layers=[1, 2], + slot=3, + ) + + np.testing.assert_array_equal(ptrs, [130, 210]) + + +def test_mamba_policy_slot_major_interleaved_role_ptrs(): + """Layer/role offsets remain inside each coalesced V2 physical slot.""" + state_bytes = 64 + physical_slot_bytes = 4 * state_bytes + conv_pool = page.PhysicalPool( + base_address=1000 + state_bytes, + slot_bytes=state_bytes, + num_slots=8, + slot_stride_bytes=physical_slot_bytes, + layer_stride_bytes=2 * state_bytes, + ) + + ptrs = peer.MambaPolicy._build_layer_ptrs( + pool=conv_pool, + layer_offsets={1: 0, 2: 1}, + overlapping_layers=[1, 2], + slot=3, + ) + + np.testing.assert_array_equal( + ptrs, + [ + 1000 + state_bytes + 3 * physical_slot_bytes, + 1000 + 3 * state_bytes + 3 * physical_slot_bytes, + ], + ) + + # --------------------------------------------------------------------------- # Ground truth: generate, shard, write, compute expected, read actual # --------------------------------------------------------------------------- -def _full_conv_section_dims() -> List[int]: - """Full (unsharded) first-dim sizes: [x(d_inner) | B(ng*ds) | C(ng*ds)].""" +def _full_conv_section_dims(conv_state_layout="x_b_c") -> List[int]: + """Full first-dimension sizes in the model's convolution-state order.""" d_inner = MAMBA_HEAD_DIM * MAMBA_NUM_HEADS ng_ds = MAMBA_N_GROUPS * MAMBA_D_STATE - return [d_inner, ng_ds, ng_ds] + if conv_state_layout == "x_b_c": + return [d_inner, ng_ds, ng_ds] + if conv_state_layout == "q_k_v": + return [ng_ds, ng_ds, d_inner] + raise ValueError(f"Unsupported convolution state layout: {conv_state_layout!r}") def _generate_ground_truth(num_requests: int, seed: int = 12345): @@ -276,29 +361,48 @@ def _shard_ssm(full_ssm: torch.Tensor, tp: int, tp_rank: int) -> torch.Tensor: return full_ssm[tp_rank * n : (tp_rank + 1) * n].clone() -def _shard_conv(full_conv: torch.Tensor, tp: int, tp_rank: int) -> torch.Tensor: - """Shard conv per-section along dim 0: [x | B | C] each independently.""" +def _shard_conv( + full_conv: torch.Tensor, + tp: int, + tp_rank: int, + conv_state_layout="x_b_c", +) -> torch.Tensor: + """Shard each semantic convolution-state section independently.""" parts = [] offset = 0 - for sec_dim in _full_conv_section_dims(): + for sec_dim in _full_conv_section_dims(conv_state_layout): n = sec_dim // tp parts.append(full_conv[offset + tp_rank * n : offset + (tp_rank + 1) * n]) offset += sec_dim return torch.cat(parts, dim=0).clone() -def _write_ground_truth_to_ctx(managers, tp, ground_truth, request_ids): +def _write_ground_truth_to_ctx( + managers, + tp, + ground_truth, + request_ids, + conv_state_layout="x_b_c", +): """Write sharded ground truth into ctx managers' allocated mamba slots.""" for rank, mgr in enumerate(managers): for req_idx, rid in enumerate(request_ids): - slot = mgr.mamba_cache_index[rid] - for layer_idx in mgr._impl.mamba_layer_offsets: + slot = _mamba_state_slot(mgr, rid) + for layer_idx in _mamba_layer_ids(mgr): full = ground_truth[req_idx][layer_idx] mgr.get_ssm_states(layer_idx)[slot] = _shard_ssm(full["ssm"], tp, rank) - mgr.get_conv_states(layer_idx)[slot] = _shard_conv(full["conv"], tp, rank) - - -def _compute_expected(ground_truth, gen_managers, gen_tp, gen_request_ids) -> Dict: + mgr.get_conv_states(layer_idx)[slot] = _shard_conv( + full["conv"], tp, rank, conv_state_layout + ) + + +def _compute_expected( + ground_truth, + gen_managers, + gen_tp, + gen_request_ids, + conv_state_layout="x_b_c", +) -> Dict: """Compute expected mamba states BEFORE transfer. Returns: {(gen_rank, req_idx, layer_idx): {"conv": Tensor, "ssm": Tensor}} @@ -306,11 +410,11 @@ def _compute_expected(ground_truth, gen_managers, gen_tp, gen_request_ids) -> Di expected = {} for gen_rank, mgr in enumerate(gen_managers): for req_idx in range(len(gen_request_ids)): - for layer_idx in mgr._impl.mamba_layer_offsets: + for layer_idx in _mamba_layer_ids(mgr): full = ground_truth[req_idx][layer_idx] expected[(gen_rank, req_idx, layer_idx)] = { "ssm": _shard_ssm(full["ssm"], gen_tp, gen_rank), - "conv": _shard_conv(full["conv"], gen_tp, gen_rank), + "conv": _shard_conv(full["conv"], gen_tp, gen_rank, conv_state_layout), } return expected @@ -323,8 +427,8 @@ def _read_actual(gen_managers, gen_request_ids) -> Dict: actual = {} for gen_rank, mgr in enumerate(gen_managers): for req_idx, rid in enumerate(gen_request_ids): - slot = mgr.mamba_cache_index[rid] - for layer_idx in mgr._impl.mamba_layer_offsets: + slot = _mamba_state_slot(mgr, rid) + for layer_idx in _mamba_layer_ids(mgr): actual[(gen_rank, req_idx, layer_idx)] = { "conv": mgr.get_conv_states(layer_idx)[slot].cpu().clone(), "ssm": mgr.get_ssm_states(layer_idx)[slot].cpu().clone(), @@ -364,14 +468,26 @@ def test_mamba_disagg_attention_dp_dummy_with_batch_size_one(): mgr.shutdown() -def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): +def run_mamba_transfer_test( + ctx_tp: int, + gen_tp: int, + use_v2: bool = False, + conv_state_layout: str = "x_b_c", +): """Test mamba transfer: ctx_tp -> gen_tp (PP=1, no DP).""" # -- 1. Create managers, zero mamba caches -- - ctx_mgrs = _create_managers(ctx_tp) - gen_mgrs = _create_managers(gen_tp) + ctx_mgrs = _create_managers( + ctx_tp, + use_v2=use_v2, + conv_state_layout=conv_state_layout, + ) + gen_mgrs = _create_managers( + gen_tp, + use_v2=use_v2, + conv_state_layout=conv_state_layout, + ) for mgr in ctx_mgrs + gen_mgrs: - mgr._impl.mamba_cache.conv.zero_() - mgr._impl.mamba_cache.temporal.zero_() + _zero_mamba_states(mgr) # -- 2. Create transceivers -- config = CacheTransceiverConfig( @@ -425,14 +541,29 @@ def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): # -- 4. Allocate slots -- ctx_batch = ScheduledRequests() - ctx_batch.reset_context_requests(ctx_reqs) - for mgr in ctx_mgrs: - mgr.prepare_resources(ctx_batch) + if use_v2: + ctx_batch.context_requests_last_chunk = ctx_reqs + for mgr in ctx_mgrs: + for req in ctx_reqs: + assert mgr.prepare_context(req) + assert mgr.resize_context(req, req.context_chunk_size) + mgr.prepare_resources(ctx_batch) + else: + ctx_batch.reset_context_requests(ctx_reqs) + for mgr in ctx_mgrs: + mgr.prepare_resources(ctx_batch) gen_batch = ScheduledRequests() - gen_batch.reset_context_requests(gen_reqs) - for mgr in gen_mgrs: - mgr.prepare_resources(gen_batch) + if use_v2: + gen_batch.context_requests_last_chunk = gen_reqs + for mgr in gen_mgrs: + for req in gen_reqs: + assert mgr.prepare_disagg_gen_init(req) + mgr.prepare_resources(gen_batch) + else: + gen_batch.reset_context_requests(gen_reqs) + for mgr in gen_mgrs: + mgr.prepare_resources(gen_batch) for req in ctx_reqs + gen_reqs: req.context_current_position = req.prompt_len @@ -444,10 +575,22 @@ def run_mamba_transfer_test(ctx_tp: int, gen_tp: int): # -- 5. Ground truth -> shard -> write to ctx -- ground_truth = _generate_ground_truth(len(REQUEST_LENGTHS)) - _write_ground_truth_to_ctx(ctx_mgrs, ctx_tp, ground_truth, ctx_rids) + _write_ground_truth_to_ctx( + ctx_mgrs, + ctx_tp, + ground_truth, + ctx_rids, + conv_state_layout, + ) # -- 6. Compute expected BEFORE transfer -- - expected = _compute_expected(ground_truth, gen_mgrs, gen_tp, gen_rids) + expected = _compute_expected( + ground_truth, + gen_mgrs, + gen_tp, + gen_rids, + conv_state_layout, + ) # -- 7. Transfer -- for rank in range(gen_tp): @@ -505,3 +648,23 @@ def test_mamba_transfer(ctx_tp, gen_tp): print(f"\nMamba transfer test: ctx_tp={ctx_tp} -> gen_tp={gen_tp}") run_mamba_transfer_test(ctx_tp, gen_tp) print("PASSED") + + +@pytest.mark.timeout(180) +@pytest.mark.parametrize( + "ctx_tp,gen_tp,conv_state_layout", + [ + (2, 2, "x_b_c"), + (2, 4, "q_k_v"), + (4, 2, "x_b_c"), + ], + ids=["same_tp_xbc", "expand_tp_qkv", "contract_tp_xbc"], +) +def test_v2_mamba_transfer(ctx_tp, gen_tp, conv_state_layout): + """Transfer slot-major V2 Mamba states through Python/NIXL.""" + run_mamba_transfer_test( + ctx_tp, + gen_tp, + use_v2=True, + conv_state_layout=conv_state_layout, + ) diff --git a/tests/unittest/disaggregated/test_peer.py b/tests/unittest/disaggregated/test_peer.py index 6aa67c38ab9d..9b52f754dbdb 100644 --- a/tests/unittest/disaggregated/test_peer.py +++ b/tests/unittest/disaggregated/test_peer.py @@ -230,6 +230,37 @@ def test_no_overlap(): assert overlap.ranks == [] +@pytest.mark.parametrize( + ("self_heads", "peer_heads"), + [(0, 0), (0, 2), (2, 0)], +) +def test_attention_free_stage_has_no_head_duplication( + self_heads: int, + peer_heads: int, +) -> None: + self_ri = make_rankinfo( + "self", + tp_size=2, + tp_rank=0, + kv_heads_per_rank=self_heads, + layer_num_per_pp=[2], + ) + peer_ri = make_rankinfo( + "peer", + tp_size=2, + tp_rank=0, + kv_heads_per_rank=peer_heads, + layer_num_per_pp=[2], + ) + reg, peer_ri = _make_peer_registrar_and_peer_ri(self_ri, peer_ri) + + overlap = reg.get_peer_overlap(peer_ri, peer_dp_rank=0) + + assert overlap.duplicate_head_factor == 1 + assert overlap.peer_duplicate_head_factor == 1 + assert overlap.ranks == [0] + + def test_pp_ratio_peer_smaller(): self_ri = make_rankinfo( "self", diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index b135c189a587..e209e3d5cd50 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -21,8 +21,8 @@ from tensorrt_llm import bindings from tensorrt_llm._torch.disaggregation.native import rank_info as rank_info_module from tensorrt_llm._torch.disaggregation.native.auxiliary import AuxBufferMeta +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo -from tensorrt_llm.bindings import DataType def test_rank_info_construction(): @@ -126,9 +126,42 @@ def test_from_kv_cache_manager_uses_first_nonzero_kv_head_count(monkeypatch) -> assert info.attention.kv_heads_per_rank == 8 +def test_from_kv_cache_manager_preserves_attention_dp_on_attention_free_stage( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(rank_info_module, "build_page_table_from_manager", lambda _: None) + mapping = SimpleNamespace( + rank=1, + tp_size=2, + tp_rank=1, + pp_size=1, + pp_rank=0, + dp_size=1, + cp_size=1, + cp_rank=0, + enable_attention_dp=True, + ) + manager = SimpleNamespace( + mapping=mapping, + num_kv_heads_per_layer=[0, 0], + pp_layers=[0, 1], + tokens_per_block=32, + head_dim=128, + dtype=bindings.DataType.HALF, + kv_factor=2, + ) + + info = RankInfo.from_kv_cache_manager("ctx", manager, device_id=0) + + assert info.attention is not None + assert info.attention.kv_heads_per_rank == 0 + assert info.attention.enable_attention_dp + assert MambaPolicy._mamba_tp(info) == (1, 0) + + @pytest.mark.parametrize( ("dtype", "expected_element_bytes", "expected_type"), - [(DataType.NVFP4, 0.5, float), (DataType.HALF, 2, int)], + [(bindings.DataType.NVFP4, 0.5, float), (bindings.DataType.HALF, 2, int)], ) def test_rank_info_represents_cache_element_bytes( monkeypatch, dtype, expected_element_bytes, expected_type diff --git a/tests/unittest/executor/test_stats_serializer.py b/tests/unittest/executor/test_stats_serializer.py index bb061e4ec80d..b051947966ae 100644 --- a/tests/unittest/executor/test_stats_serializer.py +++ b/tests/unittest/executor/test_stats_serializer.py @@ -24,6 +24,8 @@ KVCacheV2IterationStatsReport, KVCacheV2LifeCycleIterationStats, KVCacheV2PoolGroupIterationStats, + KVCacheV2SsmLifeCycleIterationStats, + KVCacheV2SsmSnapshotIterationStats, ) from tensorrt_llm.executor.base_worker import BaseWorker @@ -387,7 +389,20 @@ def test_serializer_with_v2_pool_group_stats(self): window_size=16, kind="attention", stats=life_cycle_stats, - ) + ), + 4: KVCacheV2SsmLifeCycleIterationStats( + life_cycle_id=4, + pool_group_id=8, + snapshot_stats=KVCacheV2SsmSnapshotIterationStats( + iter_snapshot_lookups=4, + iter_snapshot_hits=3, + iter_snapshot_misses=1, + iter_reused_tokens=96, + iter_unreused_tokens=32, + iter_aligned_snapshot_hits=2, + iter_unaligned_snapshot_hits=1, + ), + ), }, ) @@ -425,6 +440,23 @@ def test_serializer_with_v2_pool_group_stats(self): assert life_cycle["iterReusedBlocks"] == 5 assert life_cycle["iterMissedBlocks"] == 3 assert "iterGenAllocBlocks" not in life_cycle + ssm_life_cycle = d["kvCacheIterationStatsByLifecycle"]["4"] + assert ssm_life_cycle == { + "lifeCycleId": 4, + "poolGroupId": 8, + "windowSize": None, + "kind": "ssm", + "snapshotStats": { + "iterSnapshotLookups": 4, + "iterSnapshotHits": 3, + "iterSnapshotMisses": 1, + "iterSnapshotHitRate": 0.75, + "iterReusedTokens": 96, + "iterUnreusedTokens": 32, + "iterAlignedSnapshotHits": 2, + "iterUnalignedSnapshotHits": 1, + }, + } def test_v2_peak_block_stats_reset_tracks_interval_peak(self): """Peak block stats should cover the interval since the previous reset.""" 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 d69ed8c4de69..469c66a61f3f 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 @@ -1989,6 +1989,89 @@ def test_no_reuse_with_ssm(self) -> None: kv_cache.resume(stream) kv_cache.close() + @parameterized.expand( + [ + ("miss", None, 48, False, (1, 0, 1, 0, 48, 0, 0)), + ("aligned_hit", 32, 48, False, (1, 1, 0, 32, 16, 1, 0)), + ("unaligned_hit", 48, 64, True, (1, 1, 0, 48, 16, 0, 1)), + ] + ) + def test_ssm_snapshot_iteration_stats( + self, + _name: str, + snapshot_length: int | None, + lookup_length: int, + enable_partial_reuse: bool, + expected: tuple[int, int, int, int, int, int, int], + ) -> None: + tokens_per_block = 32 + cfg = self._make_ssm_config( + tokens_per_block=tokens_per_block, + enable_partial_reuse=enable_partial_reuse, + ) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(lookup_length)] + + if snapshot_length is not None: + seed = self.manager.create_kv_cache() + seed.resume(stream) + seed.capacity = snapshot_length + seed.history_length = snapshot_length + seed.commit(prompt[:snapshot_length], is_end=True) + seed.close() + + reused = self.manager.create_kv_cache( + input_tokens=prompt, + id=101, + # This is only a sizing hint; lookup telemetry must use the + # actual input_tokens length. + expected_prompt_length=lookup_length + 17, + ) + self.assertEqual(reused.num_committed_tokens, expected[3]) + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), {101}) + reused.commit_pending_stats() + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), set()) + + assert self.manager._life_cycles.ssm_life_cycle_id is not None + ssm_life_cycle_id = self.manager._life_cycles.ssm_life_cycle_id + snapshot_stats = self.manager.get_and_reset_ssm_snapshot_iteration_stats() + self.assertEqual(set(snapshot_stats), {ssm_life_cycle_id}) + stats = snapshot_stats[ssm_life_cycle_id] + self.assertEqual( + ( + stats.iter_snapshot_lookups, + stats.iter_snapshot_hits, + stats.iter_snapshot_misses, + stats.iter_reused_tokens, + stats.iter_unreused_tokens, + stats.iter_aligned_snapshot_hits, + stats.iter_unaligned_snapshot_hits, + ), + expected, + ) + self.assertEqual(stats.iter_snapshot_hit_rate, expected[1] / expected[0]) + self.assertEqual(self.manager.get_and_reset_ssm_snapshot_iteration_stats(), {}) + + reused.resume(stream) + reused.close() + + def test_discard_ssm_snapshot_stats_clears_dirty_state(self) -> None: + cfg = self._make_ssm_config() + self.manager = KVCacheManager(cfg) + tokens = [self.next_token() for _ in range(16)] + + kv_cache = self.manager.create_kv_cache(input_tokens=tokens, id=101) + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), {101}) + kv_cache.discard_pending_stats() + + self.assertEqual(self.manager.get_dirty_stats_kv_cache_ids(), set()) + self.assertEqual(self.manager.get_and_reset_ssm_snapshot_iteration_stats(), {}) + stream_holder = CachedCudaStream() + kv_cache.resume(cast(CudaStream, stream_holder.handle)) + kv_cache.close() + def test_ssm(self) -> None: """Inference with SSM layer: prefill 63 tokens, decode 52 tokens.""" cfg = self._make_ssm_config() @@ -2140,6 +2223,184 @@ def test_ssm_reuse_keeps_snapshots_from_multiple_commits(self) -> None: kv4.resume(stream) kv4.close() + def test_ssm_planned_drop_targets_latest_snapshot_with_shared_plans(self) -> None: + """Shared plans drop only their conversation endpoint snapshot.""" + cfg = self._make_ssm_config(tokens_per_block=32) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(64)] + + kv_cache = self.manager.create_kv_cache() + kv_cache.resume(stream) + kv_cache.capacity = 32 + kv_cache.commit(prompt[:32]) + kv_cache.capacity = 64 + kv_cache.commit(prompt[32:]) + kv_cache.stop_committing() + first_handle = kv_cache.plan_committed_block_drop() + second_handle = kv_cache.plan_committed_block_drop() + self.assertIsNotNone(first_handle) + self.assertIsNotNone(second_handle) + kv_cache.close() + + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 64) + assert first_handle is not None + first_handle.drop() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 64) + assert second_handle is not None + second_handle.drop() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 32) + + empty_cache = self.manager.create_kv_cache() + empty_cache.resume(stream) + empty_cache.stop_committing() + self.assertIsNone(empty_cache.plan_committed_block_drop()) + empty_cache.close() + + def test_ssm_planned_drop_includes_partial_swa_window(self) -> None: + """Hybrid plans include SSM and every partial SWA-window page.""" + cfg = self._make_ssm_config( + tokens_per_block=32, + num_attn_layers=1, + num_ssm_layers=1, + window_size=32, + ) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(48)] + + kv_cache = self.manager.create_kv_cache() + kv_cache.resume(stream) + kv_cache.capacity = len(prompt) + kv_cache.commit(prompt) + kv_cache.stop_committing() + drop_handle = kv_cache.plan_committed_block_drop() + self.assertIsNotNone(drop_handle) + + match = self.manager._radix_tree.match( + ReuseScope(), prompt, self.manager.enable_partial_match + ) + self.assertEqual(match.num_tokens, len(prompt)) + attn_lc_id = next(iter(self.manager._life_cycles.attention_life_cycles()))[0] + assert self.manager._life_cycles.ssm_life_cycle_id is not None + ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id + planned_pages = [] + for block in match.blocks: + page_ref = block.storage[attn_lc_id] + assert page_ref is not None + planned_pages.append(unwrap_rawref(page_ref)) + ssm_page_ref = match.blocks[-1].storage[ssm_lc_id] + assert ssm_page_ref is not None + planned_pages.append(unwrap_rawref(ssm_page_ref)) + self.assertTrue(all(page.planned_drop_count == 1 for page in planned_pages)) + planned_pages.clear() + del page_ref, ssm_page_ref + + kv_cache.close() + assert drop_handle is not None + drop_handle.drop() + self.assertEqual(self.manager.probe_reuse(input_tokens=prompt), 0) + + def test_ssm_same_block_snapshots_support_monotonic_multi_turn_reuse(self) -> None: + cfg = self._make_ssm_config(tokens_per_block=32, enable_partial_reuse=True) + self.manager = KVCacheManager(cfg) + engine = FakeEngine(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + + prompt = [self.next_token() for _ in range(64)] + + for snapshot_length, expected_reuse in ((10, 0), (20, 10), (25, 20)): + kv_cache = self.manager.create_kv_cache(input_tokens=prompt[:snapshot_length]) + self.assertEqual(kv_cache.num_committed_tokens, expected_reuse) + kv_cache.resume(stream) + kv_cache.capacity = snapshot_length + engine.execute( + [ + Step( + kv_cache, + prompt[expected_reuse:snapshot_length], + prompt[:expected_reuse], + ) + ], + stream, + ) + kv_cache.history_length = snapshot_length + kv_cache.commit(prompt[expected_reuse:snapshot_length]) + kv_cache.close() + + exact = self.manager.create_kv_cache(input_tokens=prompt[:25]) + self.assertEqual(exact.num_committed_tokens, 25) + exact.resume(stream) + exact.capacity = 25 + exact.history_length = 25 + engine.execute([Step(exact, [], prompt[:25])], stream) + exact.close() + + def test_ssm_same_block_forks_only_reuse_safe_snapshots(self) -> None: + cfg = self._make_ssm_config(tokens_per_block=32, enable_partial_reuse=True) + self.manager = KVCacheManager(cfg) + engine = FakeEngine(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(64)] + + source = self.manager.create_kv_cache() + source.resume(stream) + commit_start = 0 + for commit_end in (10, 20): + source.capacity = commit_end + chunk = prompt[commit_start:commit_end] + engine.execute([Step(source, chunk, prompt[:commit_start])], stream) + source.history_length = commit_end + source.commit(chunk) + commit_start = commit_end + source.close() + + # The retained 20-token state is in the future of a fork at token 15. + # Falling back to zero reuse is safe; reusing that state would corrupt + # the fork's SSM history. + early_fork = prompt[:15] + [self.next_token() for _ in range(25)] + early = self.manager.create_kv_cache(input_tokens=early_fork) + self.assertEqual(early.num_committed_tokens, 0) + early.resume(stream) + early.capacity = len(early_fork) + engine.execute([Step(early, early_fork, [])], stream) + early.history_length = len(early_fork) + engine.execute([Step(early, [], early_fork)], stream) + early.close() + + later_fork = prompt[:25] + [self.next_token() for _ in range(15)] + later = self.manager.create_kv_cache(input_tokens=later_fork) + self.assertEqual(later.num_committed_tokens, 20) + later.resume(stream) + later.capacity = len(later_fork) + engine.execute([Step(later, later_fork[20:], later_fork[:20])], stream) + later.history_length = len(later_fork) + engine.execute([Step(later, [], later_fork)], stream) + later.close() + + aligned = self.manager.create_kv_cache(input_tokens=prompt[:32]) + self.assertEqual(aligned.num_committed_tokens, 20) + aligned.resume(stream) + aligned.capacity = 32 + engine.execute([Step(aligned, prompt[20:32], prompt[:20])], stream) + aligned.history_length = 32 + aligned.commit(prompt[20:32]) + aligned.close() + + aligned_fork = prompt[:40] + [self.next_token() for _ in range(8)] + reused = self.manager.create_kv_cache(input_tokens=aligned_fork) + self.assertEqual(reused.num_committed_tokens, 32) + reused.resume(stream) + reused.capacity = len(aligned_fork) + engine.execute([Step(reused, aligned_fork[32:], aligned_fork[:32])], stream) + reused.history_length = len(aligned_fork) + engine.execute([Step(reused, [], aligned_fork)], stream) + reused.close() + def test_ssm_partial_snapshot_respects_partial_reuse_setting(self) -> None: """Partial SSM snapshots are created, but partial prompt reuse remains optional.""" tokens_per_block = 32 @@ -2377,6 +2638,9 @@ def tearDown(self) -> None: # Non-power-of-2 sizes so granularity rounding is non-trivial. PG0_SLOT_SIZE = 786432 # 768KB (windowed) PG1_SLOT_SIZE = 1310720 # 1280KB (non-windowed) + SSM_STATE_SLOT_SIZE = 23592960 + SSM_CONV_SLOT_SIZE = 829440 + ATTN_SLOT_SIZE = 245760 def _make_config( self, @@ -2431,6 +2695,38 @@ def _make_config( swa_scratch_reuse=(SwaScratchReuseConfig() if enable_swa_scratch_reuse else None), ) + def _make_hybrid_config(self, gpu_quota: int = 128 << 20) -> KVCacheManagerConfig: + return KVCacheManagerConfig( + tokens_per_block=self.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)], + layers=[ + SsmLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig( + role=DataRole("ssm_state"), + size=self.SSM_STATE_SLOT_SIZE, + ), + BufferConfig( + role=DataRole("conv_state"), + size=self.SSM_CONV_SLOT_SIZE, + ), + ], + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[ + BufferConfig( + role=DataRole("key"), + size=self.ATTN_SLOT_SIZE, + ), + ], + ), + ], + enable_partial_reuse=False, + commit_min_snapshot=True, + ) + def test_default_init_ratio(self): """Without typical_step or constraints, uses hardcoded fallback.""" cfg = self._make_config() @@ -2468,6 +2764,25 @@ def test_typical_step_long_sequences(self): self.assertLess(ratio[0], 0.15) manager.shutdown() + def test_zero_capacity_request_reserves_only_an_ssm_slot(self): + """Every request reserves one SSM slot, including a zero-token dummy.""" + manager = KVCacheManager(self._make_hybrid_config()) + ssm_lc = manager._life_cycles.ssm_life_cycle_id + assert ssm_lc is not None + ssm_pg = manager._storage.get_pool_group_index(ssm_lc) + attn_pg = 1 - ssm_pg + + batch = BatchDesc( + kv_caches=[ + KVCacheDesc(capacity=64, history_length=63), + KVCacheDesc(capacity=0, history_length=0), + ] + ) + slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) + self.assertEqual(slots[ssm_pg], 2) + self.assertEqual(slots[attn_pg], 2) + manager.shutdown() + def test_constraints_floor_typical_step(self): """Constraints clamp the typical_step ratio from below.""" typical = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4000)] * 32) 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 index fa43bf2d0992..ea25d28e541c 100644 --- 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 @@ -26,6 +26,7 @@ KVCacheManager, KVCacheManagerConfig, KVCacheStatsDelta, + SsmSnapshotIterationStatsDelta, ) pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") @@ -63,6 +64,20 @@ def test_stats_delta_arithmetic() -> None: assert iteration.empty assert iteration.iter_cache_hit_rate == 0.0 + snapshot = SsmSnapshotIterationStatsDelta( + iter_snapshot_lookups=4, + iter_snapshot_hits=3, + iter_snapshot_misses=1, + iter_reused_tokens=96, + iter_unreused_tokens=32, + iter_aligned_snapshot_hits=2, + iter_unaligned_snapshot_hits=1, + ) + assert snapshot.iter_snapshot_hit_rate == 0.75 + snapshot.clear() + assert snapshot.empty + assert snapshot.iter_snapshot_hit_rate == 0.0 + @pytest.mark.parametrize("enable_stats", [False, True]) def test_manager_stats_config_and_api(enable_stats: bool) -> None: @@ -72,6 +87,7 @@ def test_manager_stats_config_and_api(enable_stats: bool) -> None: assert manager.init_config.enable_stats is enable_stats assert manager.get_committed_stats() == KVCacheStatsDelta() assert manager.get_and_reset_iteration_stats() == {} + assert manager.get_and_reset_ssm_snapshot_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 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index a43818a18090..9c130f12977e 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py @@ -60,6 +60,7 @@ class _StatsRequest: state: LlmRequestState = LlmRequestState.GENERATION_IN_PROGRESS context_current_position: int = 0 context_chunk_size: int = 0 + expect_snapshot_points: list[int] = field(default_factory=list) prepopulated_prompt: tuple[int, int] | None = None kv_cache_perf_metric_calls: list[dict[str, int]] = field(default_factory=list) multimodal_hashes: None = None diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 558a57c70208..183481805aed 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -18,6 +18,7 @@ from utils.util import force_ampere import tensorrt_llm.bindings.executor as tle +import tensorrt_llm.llmapi as public_llmapi import tensorrt_llm.llmapi.llm_args as llm_args_mod from tensorrt_llm import LLM as TorchLLM from tensorrt_llm._torch.auto_deploy.llm_args import \ @@ -42,7 +43,8 @@ ExecutorMemoryType, ExtendedRuntimePerfKnobConfig, KvCacheConfig, - LookaheadDecodingConfig, MoeConfig, + LookaheadDecodingConfig, + MambaStateConfig, MoeConfig, MTPDecodingConfig, MultimodalConfig, MultimodalEncoderCudaGraphConfig, PeftCacheConfig, PybindMirror, @@ -52,7 +54,8 @@ StrictBaseModel, TorchCompileConfig, TorchLlmArgs, UserProvidedDecodingConfig, - update_llm_args_with_extra_dict) + update_llm_args_with_extra_dict, + update_llm_args_with_extra_options) # fmt: on from tensorrt_llm.llmapi.llm_utils import (_resolve_kv_cache_manager_v2_auto, _resolve_transceiver_runtime_auto, @@ -180,6 +183,37 @@ def test_llm_args_with_kvcache_config(self): 1024, 1024, 1024 ] + def test_from_yaml_migrates_legacy_mamba_interval(self, tmp_path): + yaml_path = tmp_path / "legacy_mamba_interval.yaml" + yaml_path.write_text( + yaml.safe_dump({ + "model": str(llama_model_path), + "kv_cache_config": { + "mamba_state_cache_interval": 64, + }, + }), + encoding="utf-8", + ) + + llm_args = TorchLlmArgs.from_yaml(yaml_path) + + assert llm_args.kv_cache_config.mamba_state_config.periodic_snapshot_interval == 64 + + def test_from_yaml_empty_file_reports_missing_model(self, tmp_path): + yaml_path = tmp_path / "empty.yaml" + yaml_path.write_text("", encoding="utf-8") + + with pytest.raises(ValidationError, match="model"): + TorchLlmArgs.from_yaml(yaml_path) + + def test_from_yaml_rejects_non_mapping_root(self, tmp_path): + yaml_path = tmp_path / "list.yaml" + yaml_path.write_text("[]", encoding="utf-8") + + with pytest.raises(ValueError, + match="Configuration file root must be a mapping"): + TorchLlmArgs.from_yaml(yaml_path) + def test_llm_args_with_pydantic_options(self): yaml_content = """ max_batch_size: 16 @@ -460,7 +494,9 @@ def test_kv_cache_manager_v2_auto_uses_model_default(self, explicit_auto): model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} apply_model_defaults_to_llm_args(llm_args, model_defaults) - _resolve_kv_cache_manager_v2_auto(llm_args, model_defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, + model_defaults, + original_setting="auto") assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True @@ -471,6 +507,45 @@ def test_kv_cache_manager_v2_auto_falls_back_to_false(self): assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + @pytest.mark.parametrize( + ("backend", "runtime"), + [ + ("NIXL", "CPP"), + ("UCX", None), + ("MPI", None), + ], + ) + def test_kv_cache_manager_v2_auto_falls_back_for_incompatible_disagg( + self, backend, runtime): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend=backend, transceiver_runtime=runtime), + ) + model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + + apply_model_defaults_to_llm_args(llm_args, model_defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, + model_defaults, + original_setting="auto") + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is False + + def test_kv_cache_manager_v2_auto_keeps_python_nixl_model_default(self): + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON"), + ) + model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + + apply_model_defaults_to_llm_args(llm_args, model_defaults) + _resolve_kv_cache_manager_v2_auto(llm_args, + model_defaults, + original_setting="auto") + + assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is True + @pytest.mark.parametrize("user_setting", [False, True]) def test_kv_cache_manager_v2_explicit_value_overrides_model_default( self, user_setting): @@ -586,6 +661,8 @@ def get_model_defaults(cls, llm_args): def test_KvCacheConfig_declaration(): + assert KvCacheConfig().mamba_state_cache_interval is None + assert KvCacheConfig().mamba_state_config.periodic_snapshot_interval == 0 assert KvCacheConfig().kv_cache_event_hash_algo == "auto" assert KvCacheConfig().block_reuse_policy == "all_reusable" assert KvCacheConfig().enable_swa_scratch_reuse is False @@ -613,6 +690,11 @@ def test_KvCacheConfig_declaration(): enable_swa_scratch_reuse=True, enable_partial_reuse=True, copy_on_partial_reuse=True, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[128], + additional_snapshot_offsets_from_end=[0, 32], + ), pool_ratio=[0.25, 0.75], avg_seq_len=2048, block_reuse_policy="per_request", @@ -635,6 +717,13 @@ def test_KvCacheConfig_declaration(): assert config.pool_ratio == [0.25, 0.75] assert config.avg_seq_len == 2048 assert config.block_reuse_policy == "per_request" + assert config.mamba_state_config.periodic_snapshot_interval == 0 + assert config.mamba_state_config.additional_snapshot_offsets_from_start == [ + 128 + ] + assert config.mamba_state_config.additional_snapshot_offsets_from_end == [ + 0, 32 + ] assert not hasattr(pybind_config, "pool_ratio") assert not hasattr(pybind_config, "avg_seq_len") assert not hasattr(pybind_config, "block_reuse_policy") @@ -656,6 +745,149 @@ def test_KvCacheConfig_declaration(): KvCacheConfig(block_reuse_policy="invalid") +def test_MambaStateConfig_defaults_use_independent_lists(): + first = MambaStateConfig() + second = MambaStateConfig() + + assert first.periodic_snapshot_interval == 0 + first.additional_snapshot_offsets_from_start.append(128) + first.additional_snapshot_offsets_from_end.append(0) + + assert second.additional_snapshot_offsets_from_start == [] + assert second.additional_snapshot_offsets_from_end == [] + assert public_llmapi.MambaStateConfig is MambaStateConfig + + +def test_MambaStateConfig_rejects_unknown_fields(): + with pytest.raises(ValidationError, match="extra_forbidden"): + MambaStateConfig(unknown_snapshot_policy=1) + + +@pytest.mark.parametrize( + ("field", "value"), + [ + ("periodic_snapshot_interval", -1), + ("additional_snapshot_offsets_from_start", [0]), + ("additional_snapshot_offsets_from_start", [True]), + ("additional_snapshot_offsets_from_end", [-1]), + ("additional_snapshot_offsets_from_end", [1.5]), + ], +) +def test_MambaStateConfig_rejects_invalid_snapshot_offsets(field, value): + with pytest.raises(ValidationError, match=field): + MambaStateConfig(**{field: value}) + + +@pytest.mark.parametrize( + ("field", "offsets"), + [ + ("additional_snapshot_offsets_from_start", [128]), + ("additional_snapshot_offsets_from_end", [0]), + ], +) +def test_KvCacheConfig_requires_v2_for_additional_snapshot_offsets( + field, offsets): + state_config = MambaStateConfig(**{field: offsets}) + + with pytest.raises(ValidationError, match="use_kv_cache_manager_v2=True"): + KvCacheConfig( + mamba_state_config=state_config, + use_kv_cache_manager_v2=False, + ) + + config = KvCacheConfig( + mamba_state_config=state_config, + use_kv_cache_manager_v2=True, + ) + assert getattr(config.mamba_state_config, field) == offsets + + +def test_KvCacheConfig_migrates_deprecated_mamba_interval(monkeypatch): + warnings_seen = [] + monkeypatch.setattr(llm_args_mod.logger, "warning", + lambda message: warnings_seen.append(message)) + + config = KvCacheConfig(mamba_state_cache_interval=64) + + assert config.mamba_state_cache_interval == 64 + assert config.mamba_state_config.periodic_snapshot_interval == 64 + assert any("mamba_state_cache_interval' is deprecated" in message + for message in warnings_seen) + assert "mamba_state_cache_interval" not in config.model_dump() + + +def test_KvCacheConfig_warns_when_disabling_periodic_conversation_snapshots( + monkeypatch): + warnings_seen = [] + monkeypatch.setattr(llm_args_mod.logger, "warning", + lambda message: warnings_seen.append(message)) + + config = KvCacheConfig( + block_reuse_policy="per_conversation", + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64, + additional_snapshot_offsets_from_end=[0], + ), + ) + + assert config.mamba_state_config.periodic_snapshot_interval == 0 + assert config.mamba_state_config.additional_snapshot_offsets_from_end == [0] + assert len(warnings_seen) == 1 + assert "periodic_snapshot_interval=64" in warnings_seen[0] + assert "block_reuse_policy=per_conversation" in warnings_seen[0] + assert "setting it to 0" in warnings_seen[0] + + warnings_seen.clear() + KvCacheConfig(block_reuse_policy="per_conversation") + assert warnings_seen == [] + + +def test_update_llm_args_with_empty_options_file(tmp_path): + yaml_path = tmp_path / "empty.yaml" + yaml_path.write_text("", encoding="utf-8") + llm_args = {"model": "dummy"} + + assert update_llm_args_with_extra_options(llm_args, + str(yaml_path)) == llm_args + + +def test_config_file_merge_migrates_legacy_mamba_interval_without_mutating_input( +): + yaml_dict = { + "kv_cache_config": { + "mamba_state_cache_interval": 64, + "mamba_state_config": { + "additional_snapshot_offsets_from_end": [0], + }, + }, + } + + merged = update_llm_args_with_extra_dict({"model": "dummy"}, yaml_dict) + + assert merged[ + "kv_cache_config"].mamba_state_config.periodic_snapshot_interval == 64 + assert merged[ + "kv_cache_config"].mamba_state_config.additional_snapshot_offsets_from_end == [ + 0 + ] + assert yaml_dict["kv_cache_config"]["mamba_state_cache_interval"] == 64 + + +def test_config_file_merge_rejects_legacy_and_new_mamba_intervals(): + with pytest.raises(ValueError, match="Cannot set both"): + update_llm_args_with_extra_dict( + {"model": "dummy"}, + { + "kv_cache_config": { + "mamba_state_cache_interval": 64, + "mamba_state_config": { + "periodic_snapshot_interval": 64, + }, + }, + }, + ) + + def test_KvCacheConfig_disk_cache_validation(tmp_path): config = KvCacheConfig(disk_cache_size=2048, disk_cache_path=str(tmp_path)) @@ -3123,6 +3355,126 @@ def get_preferred_transceiver_runtime(cls, pretrained_config=None): return None +class _NoModelDefaults: + + @classmethod + def get_model_defaults(cls, llm_args): + return {} + + +class TestMambaSnapshotConfigResolution: + + @staticmethod + def _load_config(monkeypatch, args, architecture): + from unittest.mock import MagicMock + + from tensorrt_llm._torch.pyexecutor import \ + model_loader as model_loader_mod + + monkeypatch.setattr( + model_loader_mod.AutoModelForCausalLM, + "_resolve_class", + staticmethod(lambda config: _NoModelDefaults), + ) + fake_loader = MagicMock() + fake_config = MagicMock() + fake_config.pretrained_config.architectures = [architecture] + fake_config.pretrained_config.hybrid_override_pattern = None + fake_loader.load_config.return_value = fake_config + return model_loader_mod.ModelLoader.load_config_and_apply_defaults( + "/tmp/dummy_model", args, fake_loader) + + @staticmethod + def _capture_warnings(monkeypatch): + from tensorrt_llm._torch.pyexecutor import \ + model_loader as model_loader_mod + + messages = [] + monkeypatch.setattr( + model_loader_mod.logger, "warning", + lambda message, *args: messages.append(message % args + if args else message)) + return messages + + @pytest.mark.parametrize( + ("kv_cache_config", "expected_reuse", "expected_warning"), + [ + (KvCacheConfig(), False, True), + ( + KvCacheConfig( + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64), + use_kv_cache_manager_v2=False, + ), + True, + False, + ), + ( + KvCacheConfig( + mamba_state_config=MambaStateConfig( + additional_snapshot_offsets_from_end=[0]), + use_kv_cache_manager_v2=True, + ), + True, + False, + ), + ( + KvCacheConfig( + block_reuse_policy="per_conversation", + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=64), + use_kv_cache_manager_v2=True, + ), + False, + True, + ), + ], + ids=["none", "periodic", "fixed", "per-conversation-no-fixed"], + ) + def test_hybrid_snapshot_policy_controls_block_reuse( + self, + monkeypatch, + kv_cache_config, + expected_reuse, + expected_warning, + ): + args = TorchLlmArgs(model="/tmp/dummy_model", + kv_cache_config=kv_cache_config) + warnings = self._capture_warnings(monkeypatch) + + self._load_config(monkeypatch, args, "Qwen3NextForCausalLM") + + assert args.kv_cache_config.enable_block_reuse is expected_reuse + assert bool(warnings) is expected_warning + if expected_warning: + assert "no Mamba state snapshot policy" in warnings[0] + + def test_hybrid_fixed_snapshot_rejects_auto_resolved_v1(self, monkeypatch): + args = TorchLlmArgs( + model="/tmp/dummy_model", + kv_cache_config=KvCacheConfig( + mamba_state_config=MambaStateConfig( + additional_snapshot_offsets_from_start=[128]), + use_kv_cache_manager_v2="auto", + ), + ) + + with pytest.raises( + ValueError, + match="use_kv_cache_manager_v2=True after resolving"): + self._load_config(monkeypatch, args, "Qwen3NextForCausalLM") + + def test_non_hybrid_without_snapshot_policy_preserves_block_reuse( + self, monkeypatch): + args = TorchLlmArgs(model="/tmp/dummy_model") + warnings = self._capture_warnings(monkeypatch) + + self._load_config(monkeypatch, args, "LlamaForCausalLM") + + assert args.kv_cache_config.enable_block_reuse is True + assert warnings == [] + + class TestTransceiverRuntimeAutoResolution: """Tests for the transceiver_runtime 'auto' selection mechanism.""" @@ -3181,13 +3533,32 @@ def test_default_backend_resolves_to_nixl_and_adopts_preference( # creation time. assert args.cache_transceiver_config.backend == "DEFAULT" - def test_default_backend_env_override_falls_back_to_cpp(self, monkeypatch): - """DEFAULT + TRTLLM_USE_UCX_KVCACHE=1 means effective UCX -> C++.""" - monkeypatch.delenv("TRTLLM_USE_NIXL_KVCACHE", raising=False) - monkeypatch.setenv("TRTLLM_USE_UCX_KVCACHE", "1") + @pytest.mark.parametrize( + "backend_env", + ["TRTLLM_USE_UCX_KVCACHE", "TRTLLM_USE_MPI_KVCACHE"], + ) + def test_default_backend_env_override_falls_back_to_v1_cpp( + self, monkeypatch, backend_env): + """An incompatible DEFAULT route falls back to the V1 C++ path.""" + for env_var in ( + "TRTLLM_USE_NIXL_KVCACHE", + "TRTLLM_USE_UCX_KVCACHE", + "TRTLLM_USE_MOONCAKE_KVCACHE", + "TRTLLM_USE_MPI_KVCACHE", + ): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setenv(backend_env, "1") args = self._disagg_args(backend="DEFAULT") + model_defaults = {"kv_cache_config": {"use_kv_cache_manager_v2": True}} + apply_model_defaults_to_llm_args(args, model_defaults) + _resolve_transceiver_runtime_auto(args, _PreferPythonTransceiverModel) + _resolve_kv_cache_manager_v2_auto(args, + model_defaults, + original_setting="auto") + assert args.cache_transceiver_config.transceiver_runtime is None + assert args.kv_cache_config.use_kv_cache_manager_v2 is False def test_disagg_disabled_is_noop(self): """Resolver never creates a config when cache_transceiver_config is None.""" diff --git a/tests/unittest/metrics/test_collector.py b/tests/unittest/metrics/test_collector.py index e66975ca6f8e..2727a36fb73e 100644 --- a/tests/unittest/metrics/test_collector.py +++ b/tests/unittest/metrics/test_collector.py @@ -814,7 +814,20 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): "iterFullReusedBlocks": 4, "iterPartialReusedBlocks": 1, "iterMissedBlocks": 3, - } + }, + "1": { + "kind": "ssm", + "snapshotStats": { + "iterSnapshotLookups": 2, + "iterSnapshotHits": 1, + "iterSnapshotMisses": 1, + "iterSnapshotHitRate": 0.5, + "iterReusedTokens": 128, + "iterUnreusedTokens": 64, + "iterAlignedSnapshotHits": 1, + "iterUnalignedSnapshotHits": 0, + }, + }, }, "kvCacheIterationStatsByPoolGroup": { "0": { @@ -829,6 +842,7 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): } before_reused = _get_counter_value(collector, "kv_cache_iter_reused_blocks") + before_missed = _get_counter_value(collector, "kv_cache_iter_missed_blocks") before_gen_alloc = _get_counter_value(collector, "kv_cache_gen_alloc_blocks_total") before_onboard = _get_counter_value(collector, "kv_cache_onboard_bytes_total") @@ -839,6 +853,9 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): assert _get_counter_value( collector, "kv_cache_iter_reused_blocks" ) - before_reused == pytest.approx(5) + assert _get_counter_value( + collector, "kv_cache_iter_missed_blocks" + ) - before_missed == pytest.approx(3) assert _get_counter_value( collector, "kv_cache_gen_alloc_blocks_total" ) - before_gen_alloc == pytest.approx(2) @@ -846,6 +863,47 @@ def test_v2_lifecycle_and_pool_group_stats_are_aggregated(self): collector, "kv_cache_onboard_bytes_total" ) - before_onboard == pytest.approx(4096) + def test_v2_ssm_only_lifecycle_falls_back_to_attention_window_stats(self): + """An SSM lifecycle entry must not hide attention's window aggregate.""" + collector = _make_kv_iter_collector() + stats = { + "kvCacheIterationStats": { + "16": { + "iterReusedBlocks": 2, + "iterFullReusedBlocks": 2, + "iterPartialReusedBlocks": 0, + "iterMissedBlocks": 2, + } + }, + "kvCacheIterationStatsByLifecycle": { + "1": { + "kind": "ssm", + "snapshotStats": { + "iterSnapshotLookups": 1, + "iterSnapshotHits": 1, + "iterSnapshotMisses": 0, + "iterSnapshotHitRate": 1.0, + "iterReusedTokens": 32, + "iterUnreusedTokens": 0, + "iterAlignedSnapshotHits": 1, + "iterUnalignedSnapshotHits": 0, + }, + } + }, + } + before_reused = _get_counter_value(collector, "kv_cache_iter_reused_blocks") + before_missed = _get_counter_value(collector, "kv_cache_iter_missed_blocks") + + collector.log_iteration_stats(stats) + + assert _get_gauge_value(collector, "kv_cache_iter_reuse_rate") == pytest.approx(0.5) + assert _get_counter_value( + collector, "kv_cache_iter_reused_blocks" + ) - before_reused == pytest.approx(2) + assert _get_counter_value( + collector, "kv_cache_iter_missed_blocks" + ) - before_missed == pytest.approx(2) + def test_multiple_windows_aggregated(self): """Stats from multiple window sizes should be summed.""" collector = _make_kv_iter_collector()