diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index e8c8816ba4c8..91af6eed3631 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -337,14 +337,24 @@ def _build_layer_ptrs( layer_offsets: Dict[int, int], overlapping_layers: List[int], slot: int, + layer_slot0_addresses: Optional[Dict[int, int]] = None, + physical_slot_stride_bytes: Optional[int] = None, ) -> 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 - ) + if layer_slot0_addresses is not None: + ptrs.append( + layer_slot0_addresses[glid] + + slot * (physical_slot_stride_bytes or pool.slot_bytes) + ) + else: + lid = layer_offsets[glid] + ptrs.append( + pool.base_address + + lid * pool.num_slots * pool.slot_bytes + + slot * pool.slot_bytes + ) return np.array(ptrs, dtype=np.int64) @staticmethod @@ -430,11 +440,41 @@ def build_mamba_frags( (self_mlg.conv_states, peer_mlg.conv_states, True), (self_mlg.ssm_states, peer_mlg.ssm_states, False), ]: + self_layer_slot0_addresses = ( + self_mlg.conv_layer_slot0_addresses + if is_conv + else self_mlg.ssm_layer_slot0_addresses + ) + peer_layer_slot0_addresses = ( + peer_mlg.conv_layer_slot0_addresses + if is_conv + else peer_mlg.ssm_layer_slot0_addresses + ) + self_physical_slot_stride_bytes = ( + self_mlg.conv_physical_slot_stride_bytes + if is_conv + else self_mlg.ssm_physical_slot_stride_bytes + ) + peer_physical_slot_stride_bytes = ( + peer_mlg.conv_physical_slot_stride_bytes + if is_conv + else peer_mlg.ssm_physical_slot_stride_bytes + ) 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, + self_layer_slot0_addresses, + self_physical_slot_stride_bytes, ) 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, + peer_layer_slot0_addresses, + peer_physical_slot_stride_bytes, ) src_region = SpecRegion( diff --git a/tensorrt_llm/_torch/disaggregation/native/rank_info.py b/tensorrt_llm/_torch/disaggregation/native/rank_info.py index 12a614ca2dad..4bea2ee1bb5b 100644 --- a/tensorrt_llm/_torch/disaggregation/native/rank_info.py +++ b/tensorrt_llm/_torch/disaggregation/native/rank_info.py @@ -59,6 +59,7 @@ def from_kv_cache_manager( m = kv_cache_manager.mapping kvm = kv_cache_manager enable_attention_dp = m.enable_attention_dp + kv_heads_per_rank = next((h for h in kvm.num_kv_heads_per_layer if h > 0), 0) return cls( instance_name=instance_name, instance_rank=m.rank, @@ -77,7 +78,7 @@ def from_kv_cache_manager( self_endpoint="", transfer_engine_info=bytes(), attention=AttentionInfo( - kv_heads_per_rank=kvm.num_kv_heads_per_layer[0], + kv_heads_per_rank=kv_heads_per_rank, tokens_per_block=kvm.tokens_per_block, dims_per_head=kvm.head_dim, element_bytes=get_size_in_bytes(1, kvm.dtype), diff --git a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py index 92f0a9c2f4ad..a766be60b6bb 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py +++ b/tensorrt_llm/_torch/disaggregation/resource/kv_extractor.py @@ -21,7 +21,10 @@ PoolView, ) from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool -from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MambaHybridCacheManager +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MambaHybridCacheManager, + V2MambaHybridCacheManager, +) 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 @@ -132,6 +135,65 @@ def _build_layer_group_for_mamba( ) +def _physical_slot_stride_bytes(tensor) -> int: + return int(tensor.stride(0) * tensor.element_size()) + + +def _build_layer_group_for_v2_mamba( + manager: V2MambaHybridCacheManager, 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() + } + + first_conv_state = manager.all_conv_states[0] + first_ssm_state = manager.all_ssm_states[0] + conv_physical_slot_stride_bytes = _physical_slot_stride_bytes(first_conv_state) + ssm_physical_slot_stride_bytes = _physical_slot_stride_bytes(first_ssm_state) + conv_slot_bytes = int(first_conv_state[0].numel() * first_conv_state.element_size()) + ssm_slot_bytes = int(first_ssm_state[0].numel() * first_ssm_state.element_size()) + num_slots = int(first_ssm_state.shape[0]) + + conv_layer_slot0_addresses = { + int(global_layer_id): int(manager.all_conv_states[offset].data_ptr()) + for global_layer_id, offset in mamba_layer_offsets.items() + } + ssm_layer_slot0_addresses = { + int(global_layer_id): int(manager.all_ssm_states[offset].data_ptr()) + for global_layer_id, offset in mamba_layer_offsets.items() + } + + d_conv_m1 = manager.conv_state_shape[1] + conv_elem_size = first_conv_state.element_size() + nheads, 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=PhysicalPool( + base_address=int(first_conv_state.data_ptr()), + slot_bytes=conv_slot_bytes, + num_slots=num_slots, + ), + ssm_states=PhysicalPool( + base_address=int(first_ssm_state.data_ptr()), + slot_bytes=ssm_slot_bytes, + num_slots=num_slots, + ), + conv_section_bytes=conv_section_bytes, + ssm_bytes_per_head=ssm_bytes_per_head, + conv_layer_slot0_addresses=conv_layer_slot0_addresses, + ssm_layer_slot0_addresses=ssm_layer_slot0_addresses, + conv_physical_slot_stride_bytes=conv_physical_slot_stride_bytes, + ssm_physical_slot_stride_bytes=ssm_physical_slot_stride_bytes, + ) + + def build_page_table(kv_cache_manager: KVCacheManager) -> KVCachePageTable: """Build a KVCachePageTable from a KVCacheManager (V1).""" if kv_cache_manager.dtype == DataType.NVFP4: @@ -339,6 +401,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, V2MambaHybridCacheManager) 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 = [ @@ -392,7 +462,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, V2MambaHybridCacheManager + ): 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 81514c06d6a0..e15b54ed4803 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/page.py +++ b/tensorrt_llm/_torch/disaggregation/resource/page.py @@ -204,6 +204,10 @@ class MambaLayerGroup(LayerGroup): ssm_states: Optional[PhysicalPool] = None conv_section_bytes: Optional[List[int]] = None ssm_bytes_per_head: Optional[int] = None + conv_layer_slot0_addresses: Optional[Dict[int, int]] = None + ssm_layer_slot0_addresses: Optional[Dict[int, int]] = None + conv_physical_slot_stride_bytes: Optional[int] = None + ssm_physical_slot_stride_bytes: Optional[int] = None def to_dict(self) -> dict: return { @@ -213,12 +217,26 @@ def to_dict(self) -> dict: "ssm_states": self.ssm_states.to_dict(), "conv_section_bytes": self.conv_section_bytes, "ssm_bytes_per_head": self.ssm_bytes_per_head, + "conv_layer_slot0_addresses": { + int(k): int(v) for k, v in (self.conv_layer_slot0_addresses or {}).items() + } + if self.conv_layer_slot0_addresses is not None + else None, + "ssm_layer_slot0_addresses": { + int(k): int(v) for k, v in (self.ssm_layer_slot0_addresses or {}).items() + } + if self.ssm_layer_slot0_addresses is not None + else None, + "conv_physical_slot_stride_bytes": self.conv_physical_slot_stride_bytes, + "ssm_physical_slot_stride_bytes": self.ssm_physical_slot_stride_bytes, } @classmethod def from_dict(cls, data: dict) -> "MambaLayerGroup": conv_section_bytes = data.get("conv_section_bytes") ssm_bytes_per_head = data.get("ssm_bytes_per_head") + conv_layer_slot0_addresses = data.get("conv_layer_slot0_addresses") + ssm_layer_slot0_addresses = data.get("ssm_layer_slot0_addresses") return cls( pool_group_idx=int(data["pool_group_idx"]), mamba_layer_offsets={int(k): int(v) for k, v in data["mamba_layer_offsets"].items()}, @@ -228,6 +246,24 @@ def from_dict(cls, data: dict) -> "MambaLayerGroup": if conv_section_bytes is not None else None, ssm_bytes_per_head=int(ssm_bytes_per_head) if ssm_bytes_per_head is not None else None, + conv_layer_slot0_addresses={ + int(k): int(v) for k, v in conv_layer_slot0_addresses.items() + } + if conv_layer_slot0_addresses is not None + else None, + ssm_layer_slot0_addresses={int(k): int(v) for k, v in ssm_layer_slot0_addresses.items()} + if ssm_layer_slot0_addresses is not None + else None, + conv_physical_slot_stride_bytes=( + int(data["conv_physical_slot_stride_bytes"]) + if data.get("conv_physical_slot_stride_bytes") is not None + else None + ), + ssm_physical_slot_stride_bytes=( + int(data["ssm_physical_slot_stride_bytes"]) + if data.get("ssm_physical_slot_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 21c4d98bd2aa..e6e04a97c7a4 100644 --- a/tensorrt_llm/_torch/disaggregation/resource/utils.py +++ b/tensorrt_llm/_torch/disaggregation/resource/utils.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 __future__ import annotations from typing import Dict, List, Set @@ -117,9 +131,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 + is_v2_layout = ( + lg.conv_layer_slot0_addresses is not None + or lg.ssm_layer_slot0_addresses is not None + ) + if is_v2_layout: + 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 4fe7fa28e50f..aea6c7db1c1e 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -30,7 +30,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, + V2MambaHybridCacheManager, +) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import nvtx_range from tensorrt_llm.bindings import LlmRequestState @@ -137,7 +140,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, V2MambaHybridCacheManager + ): 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( @@ -230,7 +235,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, V2MambaHybridCacheManager) + and self._kv_cache_manager.local_num_mamba_layers > 0 + ): + mamba_state_index = self._kv_cache_manager.get_state_indices([req.py_request_id])[0] + 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 5c2e60c4706a..c033d538489e 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 @@ -991,6 +991,13 @@ def get_model_defaults(cls, llm_args: "TorchLlmArgs") -> dict: # is supported for Mamba/SSM-based models return {"kv_cache_config": {"enable_block_reuse": False}} + @classmethod + def get_preferred_transceiver_runtime(cls, + pretrained_config: object + | None = None) -> Literal["PYTHON"]: + """Use the Python transceiver for V2 hybrid-state transfers.""" + return "PYTHON" + @staticmethod def lora_config(model_dir: str): """Nemotron-H-specific LoRA configuration. diff --git a/tensorrt_llm/_torch/models/modeling_qwen3_5.py b/tensorrt_llm/_torch/models/modeling_qwen3_5.py index 6fc43129a353..54f5b051046f 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 @@ -674,6 +674,13 @@ def get_model_defaults(cls, llm_args): # would silently fall back to the global default (block reuse on). 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 V2 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 0d91b4ebad74..3562a08fab8b 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3_next.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3_next.py @@ -18,7 +18,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 @@ -992,6 +992,13 @@ def get_model_defaults(cls, llm_args: 'TorchLlmArgs') -> dict: # is supported for Mamba/SSM-based models return {"kv_cache_config": {"enable_block_reuse": False}} + @classmethod + def get_preferred_transceiver_runtime(cls, + pretrained_config: object + | None = None) -> Literal["PYTHON"]: + """Use the Python transceiver for V2 hybrid-state transfers.""" + return "PYTHON" + def load_weights(self, weights: dict, weight_mapper: BaseWeightMapper, diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9049b1b90a73..ecee5ee24e70 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -58,6 +58,7 @@ from .mamba_cache_manager import (BaseMambaCacheManager, CppMambaHybridCacheManager, MixedMambaHybridCacheManager, + V2MambaHybridCacheManager, use_py_mamba_cache_manager) from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor @@ -86,6 +87,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, @@ -93,14 +110,17 @@ 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 the default + ``V2MambaHybridCacheManager`` 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. + V2 disaggregated serving requires the Python transceiver with the NIXL + backend. Other disaggregated configurations keep their legacy 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 @@ -121,43 +141,64 @@ def get_kv_cache_manager_cls( f"Sparse attention algorithm {sparse_attn_algorithm!r} is not " "supported with hybrid Mamba / linear-attention models.") + if is_disagg: + backend, runtime = _resolve_disagg_transceiver_route( + cache_transceiver_config) + if runtime != "PYTHON": + return CppMambaHybridCacheManager + if backend != "NIXL": + return CppMambaHybridCacheManager + # 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 use_py_mamba_cache_manager() and not is_disagg: 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 V2MambaHybridCacheManager.") 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 + default_cls = V2MambaHybridCacheManager 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 kv_cache_config.enable_block_reuse: + raise ValueError( + "TLLM_MAMBA_MANAGER_PREFERENCE=MIXED forces " + "MixedMambaHybridCacheManager, which does not support " + "block reuse. Disable block reuse or use " + "TLLM_MAMBA_MANAGER_PREFERENCE=V2/CPP.") 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': 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: + if env_override == 'V2': logger.warning( - f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " - f"Expected 'CPP' or 'MIXED'. Using default {default_cls.__name__}." - ) + "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=V2 " + "overrides the default Mamba cache manager to " + "V2MambaHybridCacheManager.") + return V2MambaHybridCacheManager + logger.warning( + f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " + f"Expected 'CPP', 'MIXED', or 'V2'. Using default {default_cls.__name__}." + ) + if kv_cache_config.enable_block_reuse: + return V2MambaHybridCacheManager + # transceiver_runtime == "PYTHON" needs no manager override: the + # Python transceiver (KvCacheTransceiverV2) drives the default + # V2MambaHybridCacheManager natively. Mixed remains available via + # TLLM_MAMBA_MANAGER_PREFERENCE=MIXED. return default_cls elif sparse_attn_config is not None: return get_sparse_attn_kv_cache_manager(sparse_attn_config) @@ -357,22 +398,21 @@ 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 + # The V1-route hybrid mamba managers (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. 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' \ + uses_v1_mamba_route = \ + os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' \ or self._speculative_config is not None - if uses_v1_mamba_route: + if uses_v1_mamba_route and not issubclass( + cls, V2MambaHybridCacheManager): 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( @@ -414,6 +454,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): + logger.warning( + "KVCacheManagerV2-backed MambaHybridCacheManager is " + "not supported with %s. Falling back to " + "CppMambaHybridCacheManager.", incompat_str) + return CppMambaHybridCacheManager # 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. @@ -1838,6 +1884,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, V2MambaHybridCacheManager): + manager_extra_kwargs["is_disagg"] = is_disagg if is_mla(config): kv_cache_manager = kv_cache_manager_cls( diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 46d3c55a0771..637664356300 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -118,6 +118,8 @@ class Role: VALUE = DataRole("value") KEY_BLOCK_SCALE = DataRole("key_block_scale") VALUE_BLOCK_SCALE = DataRole("value_block_scale") + SSM_STATE = DataRole("ssm_state") + CONV_STATE = DataRole("conv_state") # Sparse-attention per-layer index-K cache (MiniMax-M3 and similar # sparse-block-selection backends). Registered as a native V2 # BufferConfig on sparse layers via the extra_buffers_per_layer hook on @@ -698,6 +700,7 @@ def __init__( self._kv_reserve_draft_tokens = max(self.max_total_draft_tokens, draft_loop_tokens) self.event_buffer_max_size = kv_cache_config.event_buffer_max_size + self.kv_cache_config = kv_cache_config self.enable_stats = enable_stats kv_cache_event_hash_algo = get_effective_kv_cache_event_hash_algo( kv_cache_config.kv_cache_event_hash_algo, @@ -909,7 +912,9 @@ def append_to_kv_heads_per_layer( "tier (cuMemHostRegister may have failed). " "Retrying without host cache tier." ) - cache_tiers_gpu_only = [t for t in cache_tiers if isinstance(t, GpuCacheTierConfig)] + cache_tiers_gpu_only = [ + tier for tier in cache_tiers if isinstance(tier, GpuCacheTierConfig) + ] config = self._build_cache_config( kv_cache_config, tokens_per_block=tokens_per_block, @@ -1039,63 +1044,66 @@ 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] + page_index_role = self._get_pool_page_index_role(pool_id) kv_cache_pool_pointers_list.append( [ self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED + layer_id, page_index_role, PageIndexMode.SHARED ), 0, ] ) if self.dtype == DataType.NVFP4: + block_scale_role = self._get_block_scale_role_for_pool(pool_id) 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) - if self.dtype != DataType.NVFP4: - key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED - ) - - key_base_addr + page_index_role = self._get_pool_page_index_role(layer_group_id) + index_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] + addr_offset = ( + self.impl.get_mem_pool_base_address( + layer_id, page_index_role, PageIndexMode.SHARED ) + - index_base_addr + ) + offset_divisor = self.impl.get_page_stride(layer_id, page_index_role) + if self._get_pool_paired_role(layer_group_id) is not None: + offset_divisor *= self.kv_factor + offset = exact_div(addr_offset, offset_divisor) + + if self.dtype != DataType.NVFP4: + block_scale_offset = None else: - key_base_addr = kv_cache_pool_pointers_list[layer_group_id][0] - block_scale_base_addr = block_scale_pool_pointers_list[layer_group_id][0] - addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY, PageIndexMode.SHARED + block_scale_role = self._get_block_scale_role_for_pool(layer_group_id) + 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 ) - - key_base_addr - ) - block_scale_addr_offset = ( - self.impl.get_mem_pool_base_address( - layer_id, Role.KEY_BLOCK_SCALE, PageIndexMode.SHARED + block_scale_divisor = self.impl.get_page_stride(layer_id, block_scale_role) + if self._get_pool_paired_role(layer_group_id) is not None: + block_scale_divisor *= self.kv_factor + block_scale_offset = exact_div( + block_scale_addr_offset, + block_scale_divisor, ) - - block_scale_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 = exact_div( - addr_offset, - self.get_layer_bytes_per_token(layer_id, Role.KEY) - * self.kv_factor - * self.tokens_per_block, - ) - if self.dtype == DataType.NVFP4: + if block_scale_offset is not None: assert block_scale_offset == offset, ( "Block scale offset and offset should be the same" ) @@ -1130,12 +1138,16 @@ 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: + page_index_role = self._get_pool_page_index_role(pool_id) + self.index_scales[pool_id] = self.impl.get_page_index_scale(layer_id, page_index_role) + paired_role = self._get_pool_paired_role(pool_id) + if paired_role 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, paired_role, PageIndexMode.SHARED) + - self.impl.get_mem_pool_base_address( + layer_id, page_index_role, PageIndexMode.SHARED + ), + self.impl.get_page_stride(layer_id, page_index_role), ) else: self.kv_offset[pool_id] = 0 @@ -1233,6 +1245,19 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: context_extra_quota = context_tokens * context_swa_size_per_token return int(generation_quota + context_extra_quota) + def _get_pool_page_index_role(self, pool_id: int) -> DataRole: + """Return the role whose page indices and strides define this pool.""" + return Role.KEY + + def _get_pool_paired_role(self, pool_id: int) -> Optional[DataRole]: + """Return the optional role co-indexed with the pool's index role.""" + return Role.VALUE if self.kv_cache_type != CacheTypeCpp.SELFKONLY else None + + def _get_block_scale_role_for_pool(self, pool_id: int) -> Optional[DataRole]: + if self.dtype != DataType.NVFP4: + return None + return Role.KEY_BLOCK_SCALE if self._get_pool_page_index_role(pool_id) == Role.KEY else None + def _get_event_num_blocks_per_cache_level( self, cache_tiers: List[CacheTierConfig], @@ -1939,6 +1964,18 @@ def _resume_and_restore(self, req_id: int, kv_cache) -> bool: self._restore_page_index_bufs(req_id, kv_cache) return True + def _mark_context_position_as_history(self, request: LlmRequest, kv_cache) -> None: + """Advance history without making later tokens 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 KV cache for request " + f"{request.py_request_id} to {history_length} tokens" + ) + def prepare_context(self, req: LlmRequest) -> bool: """Create _KVCache, handle block reuse, and resume. Does NOT resize. @@ -2602,7 +2639,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, @@ -2784,27 +2824,33 @@ def release_resources( return requests - def try_commit_blocks(self, request: LlmRequest) -> None: + 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 - kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + kv_cache = self.kv_cache_map.get(request.py_request_id) if kv_cache is None: return - if request.context_current_position > kv_cache.num_committed_tokens: + commit_limit = request.block_reuse_commit_limit() + commit_end = min(request.context_current_position, commit_limit) + if commit_end > kv_cache.num_committed_tokens: + commit_start = 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=request.context_current_position, + start=commit_start, + end=commit_end, ) # TODO: On a disaggregated prefill server, pass is_end=True for # the last context chunk to improve performance. 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() @@ -2831,6 +2877,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self.impl.clear_stats_excluded(request.py_request_id) return kv_cache.discard_pending_stats() + self.try_commit_blocks(request, kv_cache) kv_cache.close() self.impl.clear_stats_excluded(request.py_request_id) if request.py_request_id in self._early_freed_index_requests: @@ -3132,8 +3179,16 @@ def update_context_resources(self, scheduled_batch: ScheduledRequests): self.enable_block_reuse and not self.is_draft and not req.is_dummy_request ) is_all_reusable = self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE - should_resize = not should_block_reuse or not is_all_reusable - should_commit = is_all_reusable or req.context_remaining_length == 0 + is_snapshot_boundary = req.should_save_ssm_snapshot(req.context_current_position) + has_pending_snapshot = any( + point > req.context_current_position for point in req.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 req.context_remaining_length == 0 + ) if should_resize: success = kv_cache.resize(None, req.context_current_position) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py index c64f780a0229..8928a403a422 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -15,7 +15,9 @@ from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, - CppMambaHybridCacheManager) + CppMambaHybridCacheManager, + MixedMambaHybridCacheManager, + V2MambaHybridCacheManager) from .resource_manager import KVCacheManager CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver @@ -129,6 +131,12 @@ def create_kv_cache_transceiver( if cache_transceiver_config.transceiver_runtime == "auto": cache_transceiver_config.transceiver_runtime = None + if (cache_transceiver_config.transceiver_runtime != "PYTHON" + and isinstance(mamba_cache_manager, MixedMambaHybridCacheManager)): + raise ValueError( + "MixedMambaHybridCacheManager requires the Python transceiver " + "runtime with the NIXL backend in disaggregated serving.") + _validate_disagg_inflight_cancel_config(cache_transceiver_config) if cache_transceiver_config.backend == "DEFAULT": @@ -150,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. + # + # V2MambaHybridCacheManager is backed by the Python KVCacheManagerV2 core, + # not the C++ BaseKVCacheManager binding required by CacheTransceiverCpp. + is_v2_mamba_hybrid = isinstance(mamba_cache_manager, + V2MambaHybridCacheManager) + use_python_transceiver = ( + cache_transceiver_config.transceiver_runtime == "PYTHON") + + if is_v2_mamba_hybrid and not use_python_transceiver: + raise ValueError( + "V2MambaHybridCacheManager 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 " + "(TLLM_MAMBA_MANAGER_PREFERENCE=V2).") + # 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/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 63bfb318e284..c860a08acc3a 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -675,6 +675,7 @@ def __init__( self.py_logits_post_processors = kwargs.pop("py_logits_post_processors", None) self.py_lora_path: str | None = kwargs.pop("py_lora_path", None) + self.expect_snapshot_points: list[int] = [] # Multimodal data self.py_multimodal_data = kwargs.pop("py_multimodal_data", None) encoder_input_tokens = kwargs.get("encoder_input_tokens") @@ -832,6 +833,14 @@ def set_exclude_last_generation_logits( self.py_result.set_exclude_last_generation_logits( exclude_last_generation_logits) + def block_reuse_commit_limit(self) -> int: + if not self.expect_snapshot_points: + return self.prompt_len + return min(max(self.expect_snapshot_points), self.prompt_len) + + def should_save_ssm_snapshot(self, commit_end: int) -> bool: + return commit_end in self.expect_snapshot_points + @property def cached_tokens(self) -> int: return self._cached_tokens diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index df137d96de6d..990df4b7db5d 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -27,19 +27,30 @@ 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 (AttentionLayerConfig, + BatchDesc, BufferConfig, + CacheTierConfig, + KVCacheDesc) +from tensorrt_llm.runtime.kv_cache_manager_v2 import \ + KVCacheManagerConfig as KVCacheManagerConfigPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, SsmLayerConfig, + SwaScratchReuseConfig) GB = 1 << 30 @@ -1136,24 +1147,22 @@ def update_resources(self, kv_cache_dtype_byte_size) -def calc_context_stop_positions(prompt_len: int, - tokens_per_block: int, - mamba_state_cache_interval: int, - save_last_snapshot: bool = False) -> list[int]: +def calc_context_stop_positions( + prompt_len: int, tokens_per_block: int, + mamba_state_cache_interval: Optional[int]) -> list[int]: """Compute token positions at which mamba state snapshots should be saved. - Returns positions spaced by ``mamba_state_cache_interval`` plus the final - prompt length (and optionally the last block-aligned position). + Returns regular interval boundaries. ``tokens_per_block`` is kept in the + signature because the C++/V1 path used it to derive block-aligned points; + V2 can snapshot exact partial boundaries. """ - stop_positions = list( - range(mamba_state_cache_interval, prompt_len, - mamba_state_cache_interval)) - last_ckpt = prompt_len // tokens_per_block * tokens_per_block - if save_last_snapshot and (last_ckpt not in stop_positions): - stop_positions.append(last_ckpt) - if prompt_len not in stop_positions: - stop_positions.append(prompt_len) - return stop_positions + del tokens_per_block + stop_positions = [] + if mamba_state_cache_interval is not None and mamba_state_cache_interval > 0: + stop_positions.extend( + range(mamba_state_cache_interval, prompt_len + 1, + mamba_state_cache_interval)) + return sorted({pos for pos in stop_positions if 0 < pos <= prompt_len}) @triton.jit @@ -1583,15 +1592,12 @@ def get_cache_size_per_token( 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 + # live mamba state). With pipeline parallelism, multiple microbatches + # can be in-flight concurrently on the same rank. + 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) + static_slots_per_request) # Regular-snapshot bytes per token. None / non-positive intervals # mean "no regular snapshots", so the mamba contribution is zero. @@ -2040,10 +2046,9 @@ def calc_next_context_chunk_size(self, request: LlmRequest) -> int: f"disabled, but got {current}") return prompt_len - current step = self.linear_attention_metadata.states_snapshot_interval - stop_positions = calc_context_stop_positions(prompt_len, + stop_positions = calc_context_stop_positions(request.prompt_len, self.tokens_per_block, step) - stop_positions = sorted(set(stop_positions)) for pos in stop_positions: if pos > current: return pos - current @@ -2229,3 +2234,964 @@ 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) + + +class V2MambaHybridCacheManager(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. + """ + + 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, + model_type: str = "nemotron_hybrid", + **kwargs, + ) -> None: + 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._full_attention_layer_mask = full_attention_layer_mask + self._combined_layer_mask = combined_layer_mask + self.requests = [] + 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] = self.replay_step_width + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + self._seed_request_counter = 0 + 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._mamba_state_cache_interval = ( + kv_cache_config.mamba_state_cache_interval) + self._mamba_block_reuse_enabled = kv_cache_config.enable_block_reuse + + 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 + conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state + 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 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 + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + ng_ds_local = self._n_groups_per_rank * mamba_d_state + d_inner_local = mamba_head_dim * nheads + if model_type == "qwen3_next": + self.conv_section_dims = [ + ng_ds_local, ng_ds_local, d_inner_local + ] + elif model_type == "nemotron_hybrid": + self.conv_section_dims = [ + d_inner_local, ng_ds_local, ng_ds_local + ] + else: + raise ValueError(f"Unsupported model type: {model_type}") + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + 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 + else: + logger.info( + "No local mamba layers for this rank, skipping mamba state views" + ) + self._n_groups_per_rank = 0 + self.conv_section_dims = [] + self.conv_state_shape = [] + self.ssm_state_shape = [] + self.ssm_count = 0 + self.conv_count = 0 + self.ssm_bytes = 0 + self.conv_bytes = 0 + + 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 + + self._setup_mtp_intermediate_states(spec_config, max_batch_size) + + kv_cache_config = kv_cache_config.model_copy(deep=True) + if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: + # SSM reuse is valid only at explicit snapshot boundaries. + kv_cache_config.block_reuse_policy = BlockReusePolicy.PER_REQUEST.value + + 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, + **kwargs, + ) + + self.mamba_layer_offsets = { + layer_id: offset + for offset, layer_id in enumerate(self.mamba_pp_layers) + } + self.mamba_local_layer_ids = [ + self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers + ] + self._request_id_to_state_index = {} + self.kv_cache_config = kv_cache_config + self.is_estimating_kv_cache = is_estimating_kv_cache + + self.cuda_state_indices = torch.zeros([self.max_batch_size], + dtype=torch.int32, + device="cuda") + self._host_state_indices = torch.zeros([self.max_batch_size], + dtype=torch.int32, + pin_memory=prefer_pinned()) + + if self.local_num_mamba_layers > 0: + first_mamba_local_layer = self.mamba_local_layer_ids[0] + self.ssm_life_cycle_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), Role.SSM_STATE) + self._setup_states() + self._setup_replay_buffers(spec_config) + else: + self.ssm_life_cycle_id = None + self._ssm_page_index_scale = 1 + self.all_ssm_states = [] + self.all_conv_states = [] + self._setup_replay_buffers(spec_config) + + @staticmethod + def get_cache_size_per_token(model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + **kwargs): + return CppMambaHybridCacheManager.get_cache_size_per_token( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + num_layers=num_layers, + **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_page_index_role(self, pool_id: int): + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_local_mamba_layer(layer_id): + return Role.SSM_STATE + return Role.KEY + + def _get_pool_paired_role(self, pool_id: int): + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_local_mamba_layer(layer_id): + return None + return super()._get_pool_paired_role(pool_id) + + 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 + + interval = self._mamba_state_cache_interval + num_snapshots = 0 + if interval is not None and interval > 0: + num_snapshots += capacity // interval + return num_snapshots + + def _ssm_slots_per_request_for_typical_batch( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> List[int]: + total_live_state_slots = self.max_batch_size + total_snapshot_slots = self._num_ssm_snapshots_for_capacity( + capacity, kv_cache_config) + total_slots = total_live_state_slots + total_snapshot_slots + slots_per_request, extra_slots = divmod(total_slots, + self.max_batch_size) + return [ + slots_per_request + (1 if i < extra_slots else 0) + for i in range(self.max_batch_size) + ] + + def _build_cache_config( + self, + kv_cache_config: KvCacheConfig, + *, + tokens_per_block: int, + vocab_size: Optional[int], + cache_tiers: List[CacheTierConfig], + ): + # Kept in the virtual method contract for cache-manager subclasses. + # The generic V2 config no longer stores the vocabulary size. + del vocab_size + buffer_type = [Role.KEY] + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + buffer_type.append(Role.VALUE) + if kv_cache_config.dtype == "nvfp4": + for layer_idx, hd in enumerate(self.head_dim_per_layer): + if self._is_local_mamba_layer(layer_idx): + continue + assert hd % 2 == 0, ( + f"head_dim must be divisible by 2 for nvfp4 kv cache, " + f"but layer {layer_idx} has head_dim={hd}") + buffer_type.append(Role.KEY_BLOCK_SCALE) + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + buffer_type.append(Role.VALUE_BLOCK_SCALE) + + scratch_reuse_config = None + if self.enable_swa_scratch_reuse: + scratch_reuse_config = SwaScratchReuseConfig( + max_rewind_len=self.num_extra_kv_tokens) + + layers = [] + for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): + layer_id = LayerId(local_layer_idx) + if self._mamba_layer_mask[global_layer_idx]: + layers.append( + SsmLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig(role=Role.SSM_STATE, + size=self.ssm_bytes), + BufferConfig(role=Role.CONV_STATE, + size=self.conv_bytes), + ], + )) + else: + layers.append( + AttentionLayerConfig( + layer_id=layer_id, + buffers=[ + BufferConfig( + role=role, + size=self.get_layer_bytes_per_token( + local_layer_idx=layer_id, data_role=role) * + tokens_per_block, + ) for role in buffer_type + ], + sliding_window_size=self.max_attention_window_vec[ + global_layer_idx % + len(self.max_attention_window_vec)], + num_sink_tokens=None, + )) + + typical_ssm_slots = self._ssm_slots_per_request_for_typical_batch( + self.max_seq_len, kv_cache_config) + + return KVCacheManagerConfigPy( + tokens_per_block=tokens_per_block, + cache_tiers=cache_tiers, + max_util_for_resume=kv_cache_config.max_util_for_resume, + initial_pool_ratio=kv_cache_config.pool_ratio, + layers=layers, + enable_partial_reuse=kv_cache_config.enable_partial_reuse, + constraints=[ + BatchDesc([ + KVCacheDesc(capacity=2049, + history_length=2048, + ssm_snapshots=1) + for _ in range(self.max_batch_size) + ]) + ], + typical_step=BatchDesc([ + KVCacheDesc(capacity=self.max_seq_len, + history_length=max(0, self.max_seq_len - 1), + ssm_snapshots=typical_ssm_slots[i]) + for i in range(self.max_batch_size) + ]), + swa_scratch_reuse=scratch_reuse_config, + commit_min_snapshot=True, + ) + + def _build_pool_mapping_tensors(self): + kv_cache_pool_pointers = torch.tensor( + [[ + self.impl.get_mem_pool_base_address( + self.impl.layer_grouping[pool_id][0], + self._get_pool_page_index_role(pool_id)), + 0, + ] for pool_id in range(self.num_pools)], + dtype=torch.int64, + device="cpu", + pin_memory=prefer_pinned(), + ) + + if self.dtype == DataType.NVFP4: + kv_cache_block_scale_pool_pointers = [] + for pool_id in range(self.num_pools): + layer_id = self.impl.layer_grouping[pool_id][0] + if self._is_local_mamba_layer(int(layer_id)): + kv_cache_block_scale_pool_pointers.append([0, 0]) + else: + kv_cache_block_scale_pool_pointers.append([ + self.impl.get_mem_pool_base_address( + layer_id, Role.KEY_BLOCK_SCALE), + 0, + ]) + kv_cache_pool_pointers = torch.stack( + [ + kv_cache_pool_pointers, + torch.tensor( + kv_cache_block_scale_pool_pointers, + dtype=torch.int64, + device="cpu", + pin_memory=prefer_pinned(), + ), + ], + dim=-1, + ) + + kv_cache_pool_mapping_list = [] + for local_layer_idx in range(self.num_local_layers): + layer_group_id = self.impl.get_layer_group_id( + LayerId(local_layer_idx)) + if self._is_local_mamba_layer(local_layer_idx): + offset = 0 + elif self.dtype != DataType.NVFP4: + addr_offset = (self.impl.get_mem_pool_base_address( + LayerId(local_layer_idx), Role.KEY) - + int(kv_cache_pool_pointers[layer_group_id][0])) + offset = addr_offset // (self.get_layer_bytes_per_token( + local_layer_idx=LayerId(local_layer_idx), + data_role=Role.KEY) * self.kv_factor * + self.tokens_per_block) + else: + addr_offset = ( + self.impl.get_mem_pool_base_address( + LayerId(local_layer_idx), Role.KEY) - + int(kv_cache_pool_pointers[layer_group_id][0][0])) + offset = addr_offset // (self.get_layer_bytes_per_token( + local_layer_idx=LayerId(local_layer_idx), + data_role=Role.KEY) * self.kv_factor * + self.tokens_per_block) + kv_cache_pool_mapping_list.append([layer_group_id, offset]) + + return ( + kv_cache_pool_pointers, + torch.tensor(kv_cache_pool_mapping_list, + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned()), + ) + + 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) + 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: + self.all_ssm_states = [ + self._get_state_buffer(local_layer_idx, Role.SSM_STATE, + self.ssm_state_dtype, self.ssm_state_shape) + for local_layer_idx in self.mamba_local_layer_ids + ] + self.all_conv_states = [ + self._get_state_buffer(local_layer_idx, Role.CONV_STATE, + self.conv_state_dtype, self.conv_state_shape) + for local_layer_idx in self.mamba_local_layer_ids + ] + + 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 and self.local_num_mamba_layers > 0: + speculative_num_draft_tokens = self.spec_config.tokens_per_gen_step - 1 + if not self._use_replay_state_update: + self.intermediate_ssm_states = torch.zeros( + size=[ + self.local_num_mamba_layers, max_batch_size, + speculative_num_draft_tokens + 1 + ] + 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, + speculative_num_draft_tokens + 1 + ] + 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 _setup_replay_buffers(self, spec_config) -> None: + 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 + + if (self.local_num_mamba_layers == 0 + or (not self._use_replay_state_update + and not self._mamba_ssm_stochastic_rounding)): + return + + 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 + 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 + + T = spec_config.tokens_per_gen_step + nheads, head_dim, d_state = self.ssm_state_shape + 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(self.local_num_mamba_layers, + cache_size, + 2, + T, + nheads, + head_dim, + dtype=self.conv_state_dtype, + device=device) + self.old_B = torch.zeros(self.local_num_mamba_layers, + cache_size, + 2, + T, + self._n_groups_per_rank, + d_state, + dtype=self.conv_state_dtype, + device=device) + self.old_dt = torch.zeros(self.local_num_mamba_layers, + cache_size, + 2, + nheads, + T, + dtype=torch.float32, + device=device) + self.old_dA_cumsum = torch.zeros(self.local_num_mamba_layers, + cache_size, + 2, + nheads, + T, + dtype=torch.float32, + device=device) + + def get_cache_bytes_per_token(self) -> int: + data_roles = [Role.KEY] + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + data_roles.append(Role.VALUE) + if self.dtype == DataType.NVFP4: + data_roles.append(Role.KEY_BLOCK_SCALE) + if self.kv_cache_type != CacheTypeCpp.SELFKONLY: + data_roles.append(Role.VALUE_BLOCK_SCALE) + + attention_bytes = 0 + for local_layer_idx in range(self.num_local_layers): + if self._is_local_mamba_layer(local_layer_idx): + continue + for data_role in data_roles: + attention_bytes += self.get_layer_bytes_per_token( + local_layer_idx=local_layer_idx, data_role=data_role) + + if self._mamba_block_reuse_enabled and self._mamba_state_cache_interval: + attention_bytes += (self.local_num_mamba_layers * + (self.ssm_bytes + self.conv_bytes) // + self._mamba_state_cache_interval) + if attention_bytes == 0 and self.local_num_mamba_layers > 0: + return max( + 1, + self.local_num_mamba_layers * + (self.ssm_bytes + self.conv_bytes)) + return max(1, attention_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, Role.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 + + @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 super().get_buffers(layer_idx, kv_layout) + + 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()) + pool_handled = set() + + for layer_id, layer_offset in self.layer_offsets.items(): + if self._is_local_mamba_layer(layer_offset): + continue + pool_id = self.layer_to_pool_mapping_dict[layer_offset] + if pool_id in pool_handled: + continue + buffer = super().get_buffers(layer_id) + for i in range(0, buffer.shape[0], 256): + buffer_slice = buffer[i:i + 256] + try: + has_invalid_values.logical_or_( + torch.isnan(buffer_slice).any()) + has_invalid_values.logical_or_( + torch.isinf(buffer_slice).any()) + except NotImplementedError: + some_checks_unavailable = True + if fill_with_zero: + buffer.zero_() + pool_handled.add(pool_id) + torch.cuda.synchronize() + + if some_checks_unavailable: + logger.warning( + "`torch.isnan` or `torch.isinf` is not implemented for current kv cache dtype, " + "related checks are skipped") + return bool(has_invalid_values) + + 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: + self.requests.extend(requests) + if prepare_resource: + # self.requests may still contain transfer-pending requests + # from an older batch. Only the new dummy requests participate + # in this forward pass. + self._setup_state_indices(requests) + return requests + + 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) + 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 + self.requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) + self._setup_state_indices() + num_contexts = len(scheduled_batch.context_requests) + if num_contexts == 0: + return + ctx_slots = self.cuda_state_indices[:num_contexts].long() + if self._use_replay_state_update and self.prev_num_accepted_tokens 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 + 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 + + def flush_state_transfers(self) -> None: + return + + def update_resources(self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None): + super().update_resources(scheduled_batch, attn_metadata, + kv_cache_dtype_byte_size) + + def _setup_state_indices(self, + requests: Optional[List[LlmRequest]] = None + ) -> None: + if self.local_num_mamba_layers == 0: + return + requests = self.requests if requests is None else requests + 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_life_cycle_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) + for i, req in enumerate(requests): + self._request_id_to_state_index[ + req.py_request_id] = self._host_state_indices[i].item() + + def get_state_indices(self, + request_ids: Optional[List[int]] = None, + is_padding: Optional[List[bool]] = None): + if request_ids is not None: + return [self._request_id_to_state_index[rid] for rid in request_ids] + return self.cuda_state_indices + + def get_max_resource_count(self) -> int: + return self.max_batch_size + + def is_speculative(self) -> bool: + return self.spec_config is not None + + def update_mamba_states(self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: 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) + 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: + # Mirror Mamba2Metadata's replay checkpoint predicate: only a + # checkpoint write resets PNAT and flips the active buffer. + 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) + self.prev_num_accepted_tokens[slots] = torch.where( + wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) + cache_buf_idx = self.cache_buf_idx[slots] + self.cache_buf_idx[slots] = torch.where(wrote_checkpoint, + 1 - cache_buf_idx, + cache_buf_idx) + 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, + num_accepted_draft_tokens, + 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, + num_accepted_draft_tokens, + state_indices_d, + ) + + 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 = {} + 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) + + def prepare_expect_snapshot_points(self, + requests: List[LlmRequest]) -> None: + """Set absolute context positions where FORCE_CHUNK should save snapshots.""" + if not self.kv_cache_config.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + interval = self._mamba_state_cache_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 = calc_context_stop_positions( + request.prompt_len, self.tokens_per_block, interval) + + def calc_next_context_chunk_size(self, request: LlmRequest) -> int: + prompt_len = request.prompt_len + current = request.context_current_position + if current >= prompt_len: + return 0 + if not self.kv_cache_config.enable_block_reuse: + assert current == 0, ( + "Expected context_current_position to be 0 when block reuse " + f"is disabled, but got {current}") + return prompt_len - current + stop_positions = request.expect_snapshot_points + for pos in stop_positions: + if pos > current: + return pos - current + return prompt_len - current + + @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 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.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/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a792e8f9a634..eed1add6ff93 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5216,6 +5216,10 @@ def _waiting_requests(self, context_requests: list[LlmRequest], @nvtx_range("_schedule") def _schedule(self): + if hasattr(self.kv_cache_manager, "prepare_expect_snapshot_points"): + self.kv_cache_manager.prepare_expect_snapshot_points( + self.active_requests) + scheduler_output = self.scheduler.schedule_request( self.active_requests, self.inflight_req_ids) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index a4c2e48ddc53..4e2e79b16661 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -668,23 +668,9 @@ def drafting_loop_wrapper(model): max_num_tokens = model_engine.max_num_tokens sparse_attention_config = model_engine.sparse_attention_config - # Set default value for cache_transceiver_config.max_tokens_in_buffer - if cache_transceiver_config and cache_transceiver_config.max_tokens_in_buffer is None: - cache_transceiver_config.max_tokens_in_buffer = net_max_seq_len - 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 @@ -734,6 +720,17 @@ def drafting_loop_wrapper(model): if draft_model_engine is not None: draft_model_engine.attn_runtime_features.chunked_prefill = False + # Set default value for cache_transceiver_config.max_tokens_in_buffer. + # Placed after the FlashMLA tokens_per_block override and rounded up to a + # tokens_per_block multiple: CacheTransBufferManager requires + # max_tokens_in_buffer % tokens_per_block == 0 (cacheTransBuffer.cpp), + # and net_max_seq_len is in general not aligned (e.g. max_seq_len plus a + # non-power-of-two seq_len offset). + if cache_transceiver_config and cache_transceiver_config.max_tokens_in_buffer is None: + cache_transceiver_config.max_tokens_in_buffer = ( + (net_max_seq_len + tokens_per_block - 1) // tokens_per_block * + tokens_per_block) + if enable_chunked_context: chunk_unit_size = tokens_per_block max_attention_window = kv_cache_config.max_attention_window @@ -890,17 +887,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.") + "use TLLM_MAMBA_MANAGER_PREFERENCE=MIXED 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 5bac35ab06d0..804a02fc4952 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -2541,6 +2541,8 @@ def update_resources( kv_cache_dtype_byte_size: Optional[float] = None, ): for _, resource_manager in self.resource_managers.items(): + if hasattr(resource_manager, "update_context_resources"): + resource_manager.update_context_resources(scheduled_batch) if hasattr(resource_manager, "update_resources"): if isinstance(resource_manager, KVCacheManager): resource_manager.update_resources(scheduled_batch, diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py index 22d7ac2a8821..8f43bfecee1d 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler.py @@ -116,6 +116,28 @@ def _get_lora_task_id(req: LlmRequest): return (1, lora_id) +def _next_expected_snapshot_point(req: LlmRequest) -> int | None: + return min( + (point for point in req.expect_snapshot_points if point > req.context_current_position), + default=None, + ) + + +def _get_forced_context_chunk_size(req: LlmRequest, default_chunk_size: int) -> int: + if not req.expect_snapshot_points: + return min(req.context_remaining_length, default_chunk_size) + next_point = _next_expected_snapshot_point(req) + if next_point is None: + return req.context_remaining_length + next_position = min(next_point, req.prompt_len) + return max(0, next_position - req.context_current_position) + + +def _is_forced_context_chunk_boundary(req: LlmRequest, chunk_size: int) -> bool: + next_position = req.context_current_position + chunk_size + return next_position >= req.prompt_len or next_position in req.expect_snapshot_points + + class ScheduledRequests: """Scheduled requests separated into disjoint sets. @@ -883,25 +905,38 @@ def _chunk_fcfs( if capacity is not None: current_compute_capacity -= actual_model_cost - def _chunk_forced(self, requests: RequestList, capacity: Optional[int], unit_size: int): + @staticmethod + def _force_chunk_size(req: LlmRequest, tokens_per_block: int) -> int: + assert isinstance(req.expect_snapshot_points, list) + return _get_forced_context_chunk_size(req, tokens_per_block) + + def _chunk_forced(self, requests: RequestList, capacity: Optional[int], tokens_per_block: int): """Mirrors the kFORCE_CHUNK specialization of setCtxRequestsChunkSize (microBatchScheduler.cpp). - Every request is assigned exactly min(context_remaining_length, unit_size) tokens. - Requests that would exceed the capacity budget are zeroed out. + Requests use expect_snapshot_points when present; otherwise each request + gets min(context_remaining_length, tokens_per_block). Requests that + would exceed the capacity budget are rounded down to the nearest lower + tokens_per_block multiple. This policy is designed for linear attention / Mamba2 state caching, which doesn't support estimating reusable tokens, so we don't subtract them from the budget. """ - if self.max_context_length is not None and self.max_context_length < unit_size: + if self.max_context_length is not None and self.max_context_length < tokens_per_block: raise ValueError( - f"The forced chunk size ({unit_size}) exceeds the " + f"The forced chunk size ({tokens_per_block}) exceeds the " f"max context length ({self.max_context_length})" ) total_tokens = 0 for req in requests: - req.context_chunk_size = min(req.context_remaining_length, unit_size) - if capacity is not None and total_tokens + req.context_chunk_size > capacity: - req.context_chunk_size = 0 + chunk_size = self._force_chunk_size(req, tokens_per_block) + if self.max_context_length is not None and chunk_size > self.max_context_length: + chunk_size = (self.max_context_length // tokens_per_block) * tokens_per_block + if capacity is not None and total_tokens + chunk_size > capacity: + remaining_capacity = max(0, capacity - total_tokens) + chunk_size = ( + min(chunk_size, remaining_capacity) // tokens_per_block + ) * tokens_per_block + req.context_chunk_size = int(chunk_size) total_tokens += req.context_chunk_size assert capacity is None or total_tokens <= capacity diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 2accbcc152b3..779406248b23 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -17,7 +17,7 @@ import os from typing import Optional -from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy +from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy from tensorrt_llm.logger import logger from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length @@ -25,7 +25,9 @@ RequestList, RequestScheduler, SchedulerOutput, + _get_forced_context_chunk_size, _get_lora_task_id, + _is_forced_context_chunk_boundary, drop_decoder_context_requests_waiting_for_encoder_output, ) @@ -175,8 +177,9 @@ def __init__( self.policy = scheduler_policy self.peft_cache_manager = peft_cache_manager - # Chunking config — only FCFS supported + # Chunking config. self.chunking_enabled = False + self.chunking_policy: Optional[ContextChunkingPolicy] = None self.chunk_unit_size = 0 self.max_context_length = max_num_tokens self.tokens_per_block = kv_cache_manager.tokens_per_block @@ -194,6 +197,7 @@ def __init__( ) if ctx_chunk_config is not None: self.chunking_enabled = True + self.chunking_policy = ctx_chunk_config[0] self.chunk_unit_size = ctx_chunk_config[1] # State value caches for fast comparison. @@ -574,10 +578,13 @@ def _try_schedule_context_chunked( """ remaining_budget = budget.remaining_tokens pre_prepare_context_remaining = req.context_remaining_length + force_chunk = self.chunking_policy == ContextChunkingPolicy.FORCE_CHUNK - # Min budget check — need at least one chunk unit - if remaining_budget is not None and remaining_budget < self.chunk_unit_size: - return ScheduleAction.SKIP, 0, False + if remaining_budget is not None: + no_budget = remaining_budget <= 0 + fcfs_under_min = not force_chunk and remaining_budget < self.chunk_unit_size + if no_budget or fcfs_under_min: + return ScheduleAction.SKIP, 0, False # Prepare context (create _KVCache, block reuse, resume — no resize) if not self.kv_cache_manager.prepare_context(req): @@ -587,22 +594,35 @@ def _try_schedule_context_chunked( # Calculate chunk size from remaining budget # (context_remaining_length is now correct after block reuse) context_remaining = req.context_remaining_length - budget_context_remaining = ( - context_remaining - if self.enable_prefix_aware_scheduling - else pre_prepare_context_remaining - ) - chunk_size = ( - min(remaining_budget, budget_context_remaining) - if remaining_budget is not None - else budget_context_remaining - ) + if force_chunk: + # Snapshot boundaries can be shorter than chunk_unit_size. + assert isinstance(req.expect_snapshot_points, list) + chunk_size = _get_forced_context_chunk_size(req, context_remaining) + budget_context_remaining = context_remaining + else: + budget_context_remaining = ( + context_remaining + if self.enable_prefix_aware_scheduling + else pre_prepare_context_remaining + ) + chunk_size = ( + min(remaining_budget, budget_context_remaining) + if remaining_budget is not None + else budget_context_remaining + ) if self.max_context_length is not None: chunk_size = min(chunk_size, self.max_context_length) - # Round down to chunk_unit_size boundary (unless last chunk). - if chunk_size < budget_context_remaining: + chunk_size = min( + chunk_size, remaining_budget if remaining_budget is not None else chunk_size + ) + + # Round down to chunk_unit_size boundary only when not hitting the end + # or a checkpoint. + if chunk_size < budget_context_remaining and not ( + force_chunk and _is_forced_context_chunk_boundary(req, chunk_size) + ): chunk_size = (chunk_size // self.chunk_unit_size) * self.chunk_unit_size if chunk_size <= 0: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index 61ddc4bdbb4f..02c8120e4a85 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -153,6 +153,7 @@ LayerConfig = AttentionLayerConfig | SsmLayerConfig class KVCacheDesc: capacity: int history_length: int + ssm_snapshots: int | None = None @dataclass(slots=True) class BatchDesc: diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py index 8f374e7a37f7..52dd9ecf6e1f 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_config.py @@ -145,9 +145,11 @@ def __post_init__(self) -> None: class KVCacheDesc: capacity: int history_length: int + ssm_snapshots: int | None = None def __post_init__(self) -> None: assert 0 <= self.history_length <= self.capacity + assert self.ssm_snapshots is None or self.ssm_snapshots >= 0 # A batch of requests, working as a use case the KVCacheManager must always support. diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py index 765d3bdca068..c9c844d7dc42 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py @@ -953,8 +953,13 @@ def _compute_slots_for_batch( for lc_idx, lc in typed_enumerate(self._life_cycles.get()): pg_idx = self.get_pool_group_index(lc_idx) if lc_idx == ssm_lc_idx: - # SSM: always 1 dedicated block per request, never shared. - num_slots[pg_idx] += len(batch.kv_caches) + # SSM state pages do not scale with token count. They scale + # with the number of live states and reusable snapshots the + # batch shape asks the storage planner to retain. + num_slots[pg_idx] += sum( + kv.ssm_snapshots if kv.ssm_snapshots is not None else 1 + for kv in batch.kv_caches + ) continue # Shared sys blocks (counted once): union of non-stale sys blocks across all requests. # A sys block needs memory if it's non-stale for ANY request. diff --git a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py index 9af481b2f686..910b49242408 100644 --- a/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py +++ b/tests/unittest/_torch/executor/test_disagg_index_mapper_early_release.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); 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..f5bdc6de75e4 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, + V2MambaHybridCacheManager, +) 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__(V2MambaHybridCacheManager) + + 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__(V2MambaHybridCacheManager) + + 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 be1f2936abf8..c80059b11bd7 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 from types import SimpleNamespace +from unittest.mock import MagicMock import pytest @@ -15,12 +16,19 @@ class _FakeKVCache: def __init__(self, num_committed_tokens: int) -> None: self.num_committed_tokens = num_committed_tokens self.committed_tokens: list[int] | None = None + self.history_length = 0 + self.capacity = num_committed_tokens self.stopped_committing = False def commit(self, tokens: list[int]) -> None: self.committed_tokens = tokens self.num_committed_tokens += len(tokens) + def resize(self, capacity: int, history_length: int | None = None) -> bool: + self.capacity = capacity + self.history_length = history_length + return True + def stop_committing(self) -> None: self.stopped_committing = True @@ -87,9 +95,11 @@ def test_propagates_partial_reuse_config(enable_partial_reuse: bool) -> None: def test_try_commit_blocks_commits_partial_block_at_context_end() -> None: request = SimpleNamespace( py_request_id=1, + is_dummy=False, is_dummy_request=False, context_current_position=10, context_remaining_length=0, + block_reuse_commit_limit=lambda: 10, get_tokens=lambda beam_id: list(range(10)), ) kv_cache = _FakeKVCache(num_committed_tokens=4) @@ -103,4 +113,63 @@ def test_try_commit_blocks_commits_partial_block_at_context_end() -> None: assert kv_cache.committed_tokens == [4, 5, 6, 7, 8, 9] assert kv_cache.num_committed_tokens == 10 + assert kv_cache.history_length == 10 + assert kv_cache.stopped_committing + + +def test_try_commit_blocks_stops_at_reusable_prompt_boundary() -> None: + request = SimpleNamespace( + py_request_id=1, + is_dummy=False, + is_dummy_request=False, + context_current_position=10, + context_remaining_length=0, + block_reuse_commit_limit=lambda: 8, + get_tokens=lambda beam_id: list(range(10)), + ) + kv_cache = _FakeKVCache(num_committed_tokens=4) + manager = object.__new__(KVCacheManagerV2) + manager.enable_block_reuse = True + manager.is_draft = False + manager.kv_cache_map = {request.py_request_id: kv_cache} + manager._augment_tokens_for_block_reuse = lambda tokens, request, start, end: tokens[start:end] + + manager.try_commit_blocks(request) + + assert kv_cache.committed_tokens == [4, 5, 6, 7] + assert kv_cache.num_committed_tokens == 8 + assert kv_cache.history_length == 10 assert kv_cache.stopped_committing + + +@pytest.mark.parametrize(("position", "should_commit"), [(32, False), (64, True)]) +def test_update_context_resources_commits_only_at_snapshot_boundary( + position: int, should_commit: bool +) -> None: + request = SimpleNamespace( + py_request_id=1, + is_dummy_request=False, + context_current_position=position, + context_remaining_length=128 - position, + expect_snapshot_points=[64, 128], + should_save_ssm_snapshot=lambda commit_end: commit_end == 64, + ) + kv_cache = SimpleNamespace( + is_active=True, + resize=MagicMock(return_value=True), + enable_swa_scratch_reuse=True, + ) + manager = object.__new__(KVCacheManagerV2) + manager.enable_block_reuse = True + manager.is_draft = False + manager.block_reuse_policy = BlockReusePolicy.PER_REQUEST + manager.kv_cache_map = {request.py_request_id: kv_cache} + manager.try_commit_blocks = MagicMock() + + manager.update_context_resources(SimpleNamespace(context_requests=[request])) + + kv_cache.resize.assert_not_called() + if should_commit: + manager.try_commit_blocks.assert_called_once_with(request) + else: + manager.try_commit_blocks.assert_not_called() 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 135740961927..432a90e99bb3 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -23,7 +23,7 @@ import pytest from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState -from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy +from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy # --------------------------------------------------------------------------- # State value constants @@ -75,6 +75,7 @@ def make_ctx_request( req.context_remaining_length = context_remaining_length req.prompt_len = prompt_len or context_remaining_length req.context_current_position = 0 + req.expect_snapshot_points = [] req.num_draft_tokens = num_draft_tokens req.has_draft_tokens = num_draft_tokens > 0 req.py_draft_tokens = [0] * num_draft_tokens if num_draft_tokens > 0 else [] @@ -1456,6 +1457,72 @@ def test_has_chunking_flag_set(self): # We verify indirectly: the request should be scheduled assert len(out.context_requests) == 1 + def test_force_chunk_uses_expected_snapshot_points(self): + mgr = make_kv_cache_manager(tokens_per_block=32) + sched = make_scheduler( + mgr, + max_num_tokens=8192, + ctx_chunk_config=(ContextChunkingPolicy.FORCE_CHUNK, 256), + ) + req = make_ctx_request(0, context_remaining_length=1176) + req.expect_snapshot_points = [256, 512, 768, 1024, 1176] + + out = sched.schedule_request([req], set()) + + assert ids(out.context_requests) == [0] + assert req.context_chunk_size == 256 + + def test_force_chunk_defaults_to_remaining_context_without_snapshot_points(self): + mgr = make_kv_cache_manager(tokens_per_block=32) + sched = make_scheduler( + mgr, + max_num_tokens=8192, + ctx_chunk_config=(ContextChunkingPolicy.FORCE_CHUNK, 256), + ) + req = make_ctx_request(0, context_remaining_length=1176) + + out = sched.schedule_request([req], set()) + + assert ids(out.context_requests) == [0] + assert req.context_chunk_size == 1176 + + def test_force_chunk_uses_next_expected_snapshot_point(self): + mgr = make_kv_cache_manager(tokens_per_block=32) + sched = make_scheduler( + mgr, + max_num_tokens=8192, + ctx_chunk_config=(ContextChunkingPolicy.FORCE_CHUNK, 256), + ) + req = make_ctx_request( + 0, + context_remaining_length=920, + prompt_len=1176, + is_first_context_chunk=False, + is_last_context_chunk=False, + ) + req.context_current_position = 256 + req.expect_snapshot_points = [256, 512, 768, 1024, 1176] + + out = sched.schedule_request([req], set()) + + assert ids(out.context_requests) == [0] + assert req.context_chunk_size == 256 + + def test_force_chunk_allows_prompt_end_shorter_than_unit_size(self): + mgr = make_kv_cache_manager(tokens_per_block=32) + sched = make_scheduler( + mgr, + max_num_tokens=8192, + ctx_chunk_config=(ContextChunkingPolicy.FORCE_CHUNK, 256), + ) + req = make_ctx_request(0, context_remaining_length=150) + req.expect_snapshot_points = [150] + + out = sched.schedule_request([req], set()) + + assert ids(out.context_requests) == [0] + assert req.context_chunk_size == 150 + 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 0e9e56feaf52..6d8cf911afc3 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -5,28 +5,215 @@ 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.pyexecutor._util import KvCacheCreator, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID -from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID +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.mamba_cache_manager import ( MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, + MixedMambaHybridCacheManager, PythonMambaCacheManager, + V2MambaHybridCacheManager, _get_mamba_hybrid_pool_size, + calc_context_stop_positions, ) from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType 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, + MTPDecodingConfig, + TorchLlmArgs, +) +from tensorrt_llm.llmapi.llm_utils import _resolve_transceiver_runtime_auto from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import GpuCacheTierConfig, LayerId +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_hybrid_cache_manager_factory_defaults_to_v2(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_CPP_MAMBA", raising=False) + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) + 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, + get_num_mamba_layers=lambda: 1, + ) + + assert ( + get_kv_cache_manager_cls(model_config, KvCacheConfig(enable_block_reuse=False)) + is V2MambaHybridCacheManager + ) + assert ( + get_kv_cache_manager_cls(model_config, KvCacheConfig(enable_block_reuse=True)) + is V2MambaHybridCacheManager + ) + + +def test_hybrid_cache_manager_factory_honors_cpp_preference_with_block_reuse(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_CPP_MAMBA", raising=False) + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "CPP") + 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, + get_num_mamba_layers=lambda: 1, + ) + + assert ( + get_kv_cache_manager_cls(model_config, KvCacheConfig(enable_block_reuse=True)) + is CppMambaHybridCacheManager + ) + + +@pytest.mark.parametrize( + ("backend", "runtime", "backend_env", "expected_cls"), + [ + ("NIXL", "PYTHON", None, V2MambaHybridCacheManager), + ("DEFAULT", "PYTHON", None, V2MambaHybridCacheManager), + ("DEFAULT", "PYTHON", "TRTLLM_USE_UCX_KVCACHE", CppMambaHybridCacheManager), + ("UCX", "PYTHON", None, CppMambaHybridCacheManager), + ("NIXL", "auto", None, CppMambaHybridCacheManager), + ("NIXL", None, None, CppMambaHybridCacheManager), + ("NIXL", "CPP", None, CppMambaHybridCacheManager), + ("UCX", None, None, CppMambaHybridCacheManager), + ], +) +def test_hybrid_cache_manager_factory_gates_v2_disagg_backend( + monkeypatch, backend, runtime, backend_env, expected_cls +): + 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") + 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, + get_num_mamba_layers=lambda: 1, + ) + transceiver_config = CacheTransceiverConfig(backend=backend, transceiver_runtime=runtime) + + assert ( + get_kv_cache_manager_cls( + model_config, + KvCacheConfig(enable_block_reuse=False), + is_disagg=True, + cache_transceiver_config=transceiver_config, + ) + is expected_cls + ) + + +def test_hybrid_disagg_python_runtime_keeps_mixed_override(monkeypatch): + monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) + monkeypatch.setenv("TLLM_MAMBA_MANAGER_PREFERENCE", "MIXED") + 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, + get_num_mamba_layers=lambda: 1, + ) + + assert ( + get_kv_cache_manager_cls( + model_config, + KvCacheConfig(enable_block_reuse=False), + is_disagg=True, + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON" + ), + ) + is MixedMambaHybridCacheManager + ) + + +def test_hybrid_models_resolve_auto_to_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"), + ) + _resolve_transceiver_runtime_auto(llm_args, model_cls) + assert llm_args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +def test_v2_hybrid_incompatibility_falls_back_to_cpp_manager(): + 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 = None + creator._max_beam_width = 2 + + assert ( + creator._fallback_if_unsupported_kv_cache_manager_v2( + V2MambaHybridCacheManager, model_config, KvCacheConfig() + ) + is CppMambaHybridCacheManager + ) + + def _make_mgr( max_batch_size=4, max_draft_len=2, enable_attention_dp=False, use_replay_state_update=False ): @@ -342,6 +529,142 @@ def test_non_mtp_pytorch_prepare_and_get_state_indices_flow(): ] +def test_calc_context_stop_positions_returns_snapshot_points(): + assert calc_context_stop_positions(128, 32, 64) == [64, 128] + assert calc_context_stop_positions(70, 32, 256) == [] + assert calc_context_stop_positions(70, 32, 0) == [] + assert calc_context_stop_positions(70, 32, None) == [] + + +def test_v2_hybrid_prepare_expect_snapshot_points(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + mamba_state_cache_interval=64, + ) + mgr.tokens_per_block = 32 + mgr._mamba_state_cache_interval = 64 + request = SimpleNamespace(prompt_len=150) + + mgr.prepare_expect_snapshot_points([request]) + + assert request.expect_snapshot_points == [64, 128] + + +def test_request_block_reuse_commit_limit_uses_snapshot_points(): + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[137]) + + assert LlmRequest.block_reuse_commit_limit(request) == 137 + + +def test_request_block_reuse_commit_limit_defaults_to_prompt_len(): + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[]) + + assert LlmRequest.block_reuse_commit_limit(request) == 150 + + +def test_v2_block_reuse_commit_saves_ssm_snapshot_at_snapshot_point(): + mgr = object.__new__(KVCacheManagerV2) + 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, + block_reuse_commit_limit=lambda: 137, + ) + 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) + + +def test_v2_hybrid_add_dummy_requests_forwards_encoder_output_lens(mocker): + mgr = object.__new__(V2MambaHybridCacheManager) + 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] + + +def test_v2_hybrid_prepare_expect_snapshot_points_clears_when_reuse_disabled(): + mgr = object.__new__(V2MambaHybridCacheManager) + mgr.kv_cache_config = KvCacheConfig(enable_block_reuse=False) + mgr.tokens_per_block = 32 + mgr._mamba_state_cache_interval = 64 + request = SimpleNamespace(prompt_len=150) + + mgr.prepare_expect_snapshot_points([request]) + + assert request.expect_snapshot_points == [] + + +@skip_no_cuda +def test_v2_hybrid_pool_ratio_controls_allocated_memory(): + def allocated_memory(pool_ratio): + mgr = object.__new__(V2MambaHybridCacheManager) + 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.max_seq_len = 128 + mgr._mamba_state_cache_interval = 64 + mgr.enable_swa_scratch_reuse = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + + config = mgr._build_cache_config( + KvCacheConfig( + pool_ratio=pool_ratio, + enable_partial_reuse=False, + ), + tokens_per_block=32, + vocab_size=1024, + cache_tiers=[GpuCacheTierConfig(quota=64 << 20)], + ) + 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 # @@ -365,6 +688,7 @@ def _build_hybrid_with_mamba_layer( 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 @@ -402,7 +726,335 @@ 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, + spec_config=None, + use_replay_state_update=False, + model_type="nemotron_hybrid", + enable_block_reuse=False, + enable_partial_reuse=True, + enable_attention_dp=False, +): + """Construct a real V2MambaHybridCacheManager.""" + num_attention_layers = 1 + mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers + attn_mask = [False] * num_mamba_layers + [True] * num_attention_layers + 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, + ) + return V2MambaHybridCacheManager( + 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=4, + 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, + model_type=model_type, + ) + + +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.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_dummy_indices_keep_cuda_buffer_address(): + max_batch_size = 4 + mgr = _build_v2_hybrid_with_mamba_layer(max_batch_size=max_batch_size, enable_attention_dp=True) + try: + stale_request = mgr.add_dummy_requests([123], token_nums=[8], is_gen=False)[0] + state_indices_ptr = mgr.cuda_state_indices.data_ptr() + + # Model a transfer-pending prior batch. The new ADP dummy is the only + # request participating in the next forward pass. + mgr.requests = [stale_request] * max_batch_size + new_requests = mgr.add_dummy_requests([456], token_nums=[8], is_gen=False) + + assert len(new_requests) == 1 + assert len(mgr.requests) == max_batch_size + 1 + assert mgr.cuda_state_indices.data_ptr() == state_indices_ptr + assert mgr.cuda_state_indices[0].item() == mgr.get_state_indices([456], [False])[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 + + mgr.requests.clear() + mgr.free_resources(request) + + assert request_id not in mgr._request_id_to_state_index + 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_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 + + for local_layer_idx, ssm_state, conv_state in zip( + mgr.mamba_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, Role.SSM_STATE) + conv_scale = mgr.impl.get_page_index_scale(layer_id, Role.CONV_STATE) + assert ssm_state.stride(0) == mgr.ssm_count * ssm_scale + assert conv_state.stride(0) == mgr.conv_count * conv_scale + assert ( + ssm_state.shape[0] + == (mgr.impl.get_page_index_upper_bound(layer_id, Role.SSM_STATE) + ssm_scale - 1) + // ssm_scale + ) + assert ( + conv_state.shape[0] + == (mgr.impl.get_page_index_upper_bound(layer_id, Role.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_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, model_type="qwen3_next") + 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_cpp_hybrid_replay_buffers_size_by_tokens_per_gen_step(): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = _build_hybrid_with_mamba_layer( + spec_config=spec_config, + max_batch_size=4, + 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() + + +@skip_no_cuda +def test_v2_hybrid_replay_buffers_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, + 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 == 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() + + +@skip_no_cuda +def test_v2_hybrid_replay_bookkeeping_matches_checkpoint_predicate(monkeypatch): + 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, + use_replay_state_update=True, + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.mamba_cache_manager._promote_mamba_state_triton", + lambda *args, **kwargs: None, ) + try: + slot = torch.tensor([0], dtype=torch.int32, device="cuda") + attn_metadata = SimpleNamespace(num_seqs=1, num_contexts=0) + + mgr.update_mamba_states( + attn_metadata, + torch.tensor([3], dtype=torch.int32, device="cuda"), + state_indices=slot, + ) + assert mgr.prev_num_accepted_tokens[0].item() == 3 + assert mgr.cache_buf_idx[0].item() == 0 + + mgr.update_mamba_states( + attn_metadata, + torch.tensor([2], dtype=torch.int32, device="cuda"), + state_indices=slot, + ) + assert mgr.prev_num_accepted_tokens[0].item() == 2 + assert mgr.cache_buf_idx[0].item() == 1 + finally: + mgr.shutdown() @skip_no_cuda diff --git a/tests/unittest/_torch/executor/test_py_scheduler.py b/tests/unittest/_torch/executor/test_py_scheduler.py index 8776295bbda9..2b935527f448 100644 --- a/tests/unittest/_torch/executor/test_py_scheduler.py +++ b/tests/unittest/_torch/executor/test_py_scheduler.py @@ -1478,9 +1478,9 @@ def test_draft_tokens_no_chunking_discard_none(self): class TestForceChunkPolicy: """ Tests for FORCE_CHUNK chunking policy in PyMicroBatchScheduler. - FORCE_CHUNK always chunks every context request to at most chunk_unit_size - tokens per scheduling step, regardless of whether the full context would fit - in the budget. + FORCE_CHUNK chunks every context request at expected chunking points when + present; otherwise it uses at most chunk_unit_size tokens per scheduling + step, regardless of whether the full context would fit in the budget. Aligned with C++ ForceChunkTest in microBatchSchedulerTest.cpp. """ @@ -1575,7 +1575,8 @@ def test_multiple_requests(self): def test_capacity_limits(self): """ - When capacity is limited, later requests get chunk_size=0. + When remaining capacity is smaller than chunk_unit_size, later requests + are rounded down to chunk_size=0. C++ ref: ForceChunkTest.CapacityLimits """ config = ContextChunkingConfig(ChunkingPolicy.FORCE_CHUNK, chunk_unit_size=10) @@ -1608,6 +1609,39 @@ def test_capacity_exact_fit(self): assert reqs[0].context_chunk_size == 10 assert reqs[1].context_chunk_size == 10 + def test_expected_snapshot_points(self): + """ + Expected snapshot points are absolute context positions. + C++ ref: ForceChunkTest.ExpectedChunkingPoints + """ + reqs = [make_context_request(0, prompt_len=30)] + reqs[0].expect_snapshot_points = [12, 25] + + self._chunk_iteration(reqs, 10) + self._expect_positions(reqs, [12], "iter 1") + + self._chunk_iteration(reqs, 10) + self._expect_positions(reqs, [25], "iter 2") + + self._chunk_iteration(reqs, 10) + self._expect_positions(reqs, [30], "iter 3") + + def test_capacity_rounds_expected_snapshot_down_to_unit(self): + """ + Capacity truncation rounds expected chunks down to chunk_unit_size. + C++ ref: ForceChunkTest.CapacityRoundsExpectedChunkDownToUnit + """ + config = ContextChunkingConfig(ChunkingPolicy.FORCE_CHUNK, chunk_unit_size=10) + scheduler = PyMicroBatchScheduler( + max_batch_size=64, max_num_tokens=1000, ctx_chunk_config=config + ) + reqs = [make_context_request(0, prompt_len=50)] + reqs[0].expect_snapshot_points = [30] + + scheduler._set_ctx_requests_chunk_size(reqs, capacity=25) + + assert reqs[0].context_chunk_size == 20 + def test_multi_iteration(self): """ A request with prompt_len=25 and chunk_unit_size=10 processes in 3 diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index b4888797eb42..9c50f273558d 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -16,7 +16,8 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ( - KVCacheManager, PeftCacheManager, _merge_kv_cache_pool_pointers, + KVCacheManager, PeftCacheManager, ResourceManager, ResourceManagerType, + _merge_kv_cache_pool_pointers, _warn_if_unsupported_v1_kv_cache_event_hash_algo) from tensorrt_llm.bindings import LayerType from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp @@ -164,6 +165,29 @@ def test_default_nvfp4_pool_configuration_merges_scale_pointers(self): manager.shutdown() +def test_resource_manager_updates_context_resources_before_generation_resources( +): + scheduled_batch = object() + calls = [] + + class MockResourceManager: + + def update_context_resources(self, batch): + assert batch is scheduled_batch + calls.append("context") + + def update_resources(self, batch): + assert batch is scheduled_batch + calls.append("generation") + + manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: MockResourceManager()}) + + manager.update_resources(scheduled_batch) + + assert calls == ["context", "generation"] + + class TestResourceManager(unittest.TestCase): CPP_RESOURCES_DIR = os.path.join(str(root_dir), "cpp", "tests", "resources") CPP_DATA_DIR = os.path.join(CPP_RESOURCES_DIR, "data") diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 6fc2f68fef68..a969068735bb 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.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. + import numpy as np import pytest @@ -244,11 +258,19 @@ def test_mamba_layer_group_serialization(): ssm_states=ssm_pool, conv_section_bytes=[512, 256, 256], ssm_bytes_per_head=128, + conv_layer_slot0_addresses={10: 1000, 11: 2000, 12: 3000}, + ssm_layer_slot0_addresses={10: 8000, 11: 9000, 12: 10000}, + conv_physical_slot_stride_bytes=512, + ssm_physical_slot_stride_bytes=1024, ) d = mlg.to_dict() assert d["mamba_layer_offsets"] == {10: 0, 11: 1, 12: 2} assert d["conv_section_bytes"] == [512, 256, 256] + assert d["conv_layer_slot0_addresses"] == {10: 1000, 11: 2000, 12: 3000} + assert d["ssm_layer_slot0_addresses"] == {10: 8000, 11: 9000, 12: 10000} + assert d["conv_physical_slot_stride_bytes"] == 512 + assert d["ssm_physical_slot_stride_bytes"] == 1024 from tensorrt_llm._torch.disaggregation.resource.page import LayerGroup @@ -263,6 +285,93 @@ def test_mamba_layer_group_serialization(): assert restored.ssm_states.num_slots == 8 assert restored.conv_section_bytes == [512, 256, 256] assert restored.ssm_bytes_per_head == 128 + assert restored.conv_layer_slot0_addresses == {10: 1000, 11: 2000, 12: 3000} + assert restored.ssm_layer_slot0_addresses == {10: 8000, 11: 9000, 12: 10000} + assert restored.conv_physical_slot_stride_bytes == 512 + assert restored.ssm_physical_slot_stride_bytes == 1024 + + +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, + ), + ssm_states=PhysicalPool( + base_address=1000, + slot_bytes=state_bytes, + num_slots=num_slots, + ), + conv_layer_slot0_addresses={ + 1: 1000 + state_bytes, + 2: 1000 + state_bytes * 3, + }, + ssm_layer_slot0_addresses={ + 1: 1000, + 2: 1000 + state_bytes * 2, + }, + conv_physical_slot_stride_bytes=physical_slot_bytes, + ssm_physical_slot_stride_bytes=physical_slot_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(): diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index 0cf0b3209899..928947fcdce9 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -23,13 +23,19 @@ 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 import rank_info +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 ( + MixedMambaHybridCacheManager, + V2MambaHybridCacheManager, +) 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 @@ -183,8 +189,8 @@ 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): + """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. @@ -194,7 +200,9 @@ 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 = V2MambaHybridCacheManager if use_v2 else MixedMambaHybridCacheManager + manager_kwargs = {"is_disagg": True} 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, @@ -220,11 +228,97 @@ 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, V2MambaHybridCacheManager): + return manager.mamba_layer_offsets + return manager._impl.mamba_layer_offsets + + +def _mamba_state_slot(manager, request_id): + if isinstance(manager, V2MambaHybridCacheManager): + return manager.get_state_indices([request_id])[0] + 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_slot_major_layer_ptrs(): + """V2 Mamba state tensors are slot-major, not layer-major.""" + self_mlg = page.MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=page.PhysicalPool(base_address=100, slot_bytes=10, num_slots=8), + ssm_states=page.PhysicalPool(base_address=200, slot_bytes=20, num_slots=8), + conv_section_bytes=[10], + ssm_bytes_per_head=10, + conv_layer_slot0_addresses={1: 1000, 2: 2000}, + ssm_layer_slot0_addresses={1: 3000, 2: 4000}, + conv_physical_slot_stride_bytes=64, + ssm_physical_slot_stride_bytes=128, + ) + peer_mlg = page.MambaLayerGroup( + pool_group_idx=0, + mamba_layer_offsets={1: 0, 2: 1}, + conv_states=page.PhysicalPool(base_address=500, slot_bytes=10, num_slots=8), + ssm_states=page.PhysicalPool(base_address=600, slot_bytes=20, num_slots=8), + conv_section_bytes=[10], + ssm_bytes_per_head=10, + conv_layer_slot0_addresses={1: 5000, 2: 6000}, + ssm_layer_slot0_addresses={1: 7000, 2: 8000}, + conv_physical_slot_stride_bytes=64, + ssm_physical_slot_stride_bytes=128, + ) + self_ri = rank_info.RankInfo( + instance_name="self", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + server_endpoint="", + self_endpoint="", + transfer_engine_info=bytes(), + ) + peer_ri = rank_info.RankInfo( + instance_name="peer", + instance_rank=0, + tp_size=1, + tp_rank=0, + pp_size=1, + pp_rank=0, + layer_num_per_pp=[2], + sender_endpoints=[], + server_endpoint="", + self_endpoint="", + transfer_engine_info=bytes(), + ) + + src_frags, dst_frags, kv_sizes = peer.MambaPolicy.build_mamba_frags( + self_mlg=self_mlg, + peer_mlg=peer_mlg, + src_slot=3, + dst_slot=5, + self_ri=self_ri, + peer_ri=peer_ri, + ) + + assert src_frags == [1192, 2192, 3384, 4384] + assert dst_frags == [5320, 6320, 7640, 8640] + assert kv_sizes == [10, 10, 20, 20] + + # --------------------------------------------------------------------------- # Ground truth: generate, shard, write, compute expected, read actual # --------------------------------------------------------------------------- @@ -285,8 +379,8 @@ def _write_ground_truth_to_ctx(managers, tp, ground_truth, request_ids): """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) @@ -300,7 +394,7 @@ 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), @@ -317,8 +411,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(), @@ -358,14 +452,13 @@ 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): """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) + gen_mgrs = _create_managers(gen_tp, use_v2=use_v2) 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( @@ -419,14 +512,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 @@ -499,3 +607,14 @@ 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", + [(2, 2), (2, 4)], + ids=["same_tp", "tp_mismatch"], +) +def test_v2_mamba_transfer(ctx_tp, gen_tp): + """Transfer slot-major V2 Mamba states through Python/NIXL.""" + run_mamba_transfer_test(ctx_tp, gen_tp, use_v2=True) 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 9b8bda967ca3..d3c4f808a8fe 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 @@ -2335,6 +2335,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, @@ -2389,6 +2392,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() @@ -2426,6 +2461,35 @@ def test_typical_step_long_sequences(self): self.assertLess(ratio[0], 0.15) manager.shutdown() + def test_ssm_snapshots_control_ssm_slot_count(self): + """SSM sizing uses explicit state-slot count, not token capacity.""" + 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=1024, + history_length=1023, + ssm_snapshots=3, + ) + ] + * 2 + ) + slots = manager._storage._compute_slots_for_batch(batch, self.TOKENS_PER_BLOCK, None) + self.assertEqual(slots[ssm_pg], 6) + self.assertEqual(slots[attn_pg], 2 * div_up(1024, self.TOKENS_PER_BLOCK)) + + default_batch = BatchDesc(kv_caches=[KVCacheDesc(capacity=4096, history_length=4095)] * 2) + default_slots = manager._storage._compute_slots_for_batch( + default_batch, self.TOKENS_PER_BLOCK, None + ) + self.assertEqual(default_slots[ssm_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_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index a43818a18090..1281245f8e0e 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: set[int] = field(default_factory=set) prepopulated_prompt: tuple[int, int] | None = None kv_cache_perf_metric_calls: list[dict[str, int]] = field(default_factory=list) multimodal_hashes: None = None @@ -82,6 +83,14 @@ def get_tokens(self, beam_id: int = DEFAULT_BEAM_INDEX) -> list[int]: assert beam_id == DEFAULT_BEAM_INDEX return self.tokens + def block_reuse_commit_limit(self) -> int: + if not self.expect_snapshot_points: + return self.prompt_len + return min(max(self.expect_snapshot_points), self.prompt_len) + + def should_save_ssm_snapshot(self, commit_end: int) -> bool: + return commit_end in self.expect_snapshot_points + def set_prepopulated_prompt_len(self, length: int, tokens_per_block: int) -> None: self.prepopulated_prompt = (length, tokens_per_block) diff --git a/tests/unittest/others/test_kv_cache_transceiver.py b/tests/unittest/others/test_kv_cache_transceiver.py index cd84403c6104..0d6b31e367d1 100644 --- a/tests/unittest/others/test_kv_cache_transceiver.py +++ b/tests/unittest/others/test_kv_cache_transceiver.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import gc import multiprocessing import sys @@ -33,6 +36,23 @@ KV_TRANSFER_COMPLETION_MARGIN_S = 10.0 +@pytest.mark.parametrize("transceiver_runtime", ["CPP", "auto"]) +def test_cpp_transceiver_rejects_mixed_mamba_manager(transceiver_runtime): + config = CacheTransceiverConfig(backend="NIXL", + transceiver_runtime=transceiver_runtime) + mixed_manager = object.__new__(MixedMambaHybridCacheManager) + + with pytest.raises( + ValueError, + match="MixedMambaHybridCacheManager requires the Python"): + create_kv_cache_transceiver(mapping=None, + dist=None, + kv_cache_manager=None, + attention_type=AttentionTypeCpp.DEFAULT, + cache_transceiver_config=config, + mamba_cache_manager=mixed_manager) + + def create_kv_cache_manager(mapping, dtype, max_tokens=256, @@ -716,7 +736,7 @@ def create_hybrid_cache_manager(mapping, dtype, mamba_conv_dtype=torch.float16, mamba_ssm_dtype=torch.float16): - """Create a MixedMambaHybridCacheManager for testing hybrid models. + """Create a mixed hybrid manager for Python transceiver tests. This manager handles both KV cache (attention layers) and Mamba cache (RNN layers). @@ -790,7 +810,7 @@ def hybrid_dtypes(request): Returns (kv_dtype, mamba_conv_dtype, mamba_ssm_dtype) based on the parametrized string. KV dtype: fp8, bf16 - Conv dtype: fp8, bf16, fp32 + Conv dtype: bf16, fp32 SSM dtype: bf16, fp32 """ kv_dtype_str, conv_dtype_str, ssm_dtype_str = request.param @@ -813,7 +833,6 @@ def hybrid_dtypes(request): @pytest.mark.timeout(120) -@pytest.mark.parametrize("backend", ["NIXL", "UCX"], ids=["NIXL", "UCX"]) @pytest.mark.parametrize( "hybrid_dtypes", [ @@ -839,18 +858,20 @@ def hybrid_dtypes(request): ], indirect=["hybrid_dtypes"], ) -def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes): +def test_hybrid_cache_transceiver_single_process(hybrid_dtypes, request): mapping = Mapping(world_size=1, rank=0) kv_dtype, mamba_conv_dtype, mamba_ssm_dtype = hybrid_dtypes # Create hybrid cache managers (combines KV + Mamba) for context and generation hybrid_cache_manager_ctx = create_hybrid_cache_manager( mapping, kv_dtype, mamba_conv_dtype, mamba_ssm_dtype) + request.addfinalizer(hybrid_cache_manager_ctx.shutdown) hybrid_cache_manager_gen = create_hybrid_cache_manager( mapping, kv_dtype, mamba_conv_dtype, mamba_ssm_dtype) + request.addfinalizer(hybrid_cache_manager_gen.shutdown) - cache_transceiver_config = CacheTransceiverConfig(backend=backend, - max_tokens_in_buffer=512) + cache_transceiver_config = CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON", max_tokens_in_buffer=512) dist = Distributed.get(mapping) # Create transceivers - the hybrid manager serves as both kv_cache_manager and mamba_cache_manager @@ -861,6 +882,7 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes): AttentionTypeCpp.DEFAULT, cache_transceiver_config, mamba_cache_manager=hybrid_cache_manager_ctx) + request.addfinalizer(cache_transceiver_ctx.shutdown) cache_transceiver_gen = create_kv_cache_transceiver( mapping, @@ -869,6 +891,7 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes): AttentionTypeCpp.DEFAULT, cache_transceiver_config, mamba_cache_manager=hybrid_cache_manager_gen) + request.addfinalizer(cache_transceiver_gen.shutdown) # Fill both KV and Mamba cache buffers with random data fill_hybrid_cache_buffers(hybrid_cache_manager_ctx) @@ -883,6 +906,9 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes): sampling_params._get_sampling_config()), is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + ctx_request.py_disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="context_only", + disagg_request_id=uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF) # Prepare resources for hybrid manager (handles both KV and Mamba) scheduled_ctx = ScheduledRequests() @@ -902,6 +928,14 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes): is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, context_phase_params=ctx_request.context_phase_params) + gen_request.py_disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="generation_only", + disagg_request_id=ctx_request.py_disaggregated_params.disagg_request_id, + ctx_request_id=ctx_request.request_id, + ctx_dp_rank=ctx_request.context_phase_params.ctx_dp_rank, + ctx_info_endpoint=ctx_request.context_phase_params.disagg_info_endpoint, + first_gen_tokens=ctx_request.context_phase_params.first_gen_tokens, + draft_tokens=ctx_request.context_phase_params.draft_tokens) # Prepare resources for hybrid manager on gen side scheduled_gen = ScheduledRequests() @@ -911,6 +945,7 @@ def test_hybrid_cache_transceiver_single_process(backend, hybrid_dtypes): cache_transceiver_gen.request_and_receive_async(gen_request) completed_ctx_ids = set() + expected_ctx_id = get_context_completed_request_id(ctx_request, "PYTHON") def poll_transfers(): completed, failed = cache_transceiver_ctx.check_context_transfer_status( @@ -920,8 +955,7 @@ def poll_transfers(): cache_transceiver_gen.check_gen_transfer_status(1) def transfers_done(): - return (ctx_request.py_request_id in completed_ctx_ids - and gen_request.state + return (expected_ctx_id in completed_ctx_ids and gen_request.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE) wait_for_transfer_completion(poll_transfers, transfers_done) @@ -934,10 +968,10 @@ def transfers_done(): # independently-allocated slots on each side, so we check the # request's own slot instead of the full state buffer (which has # extra padding-dummy slots that only the ctx side touched). - slot_ctx = hybrid_cache_manager_ctx._impl.mamba_impl.get_cache_index( - ctx_request.py_request_id) - slot_gen = hybrid_cache_manager_gen._impl.mamba_impl.get_cache_index( - gen_request.py_request_id) + slot_ctx = hybrid_cache_manager_ctx.mamba_cache_index[ + ctx_request.py_request_id] + slot_gen = hybrid_cache_manager_gen.mamba_cache_index[ + gen_request.py_request_id] assert torch.equal( hybrid_cache_manager_gen.get_conv_states(1)[slot_gen], hybrid_cache_manager_ctx.get_conv_states(1)[slot_ctx]), ( @@ -947,20 +981,22 @@ def transfers_done(): hybrid_cache_manager_gen.get_ssm_states(1)[slot_gen], hybrid_cache_manager_ctx.get_ssm_states(1)[slot_ctx]), ( "different mamba ssm states") + shutdown_transceivers(cache_transceiver_gen, cache_transceiver_ctx) @pytest.mark.timeout(120) -@pytest.mark.parametrize("backend", ["NIXL", "UCX"], ids=["NIXL", "UCX"]) -def test_hybrid_cache_transceiver_cancel_request(backend): +def test_hybrid_cache_transceiver_cancel_request(request): mapping = Mapping(world_size=1, rank=0) dtype = DataType.HALF hybrid_cache_manager_ctx = create_hybrid_cache_manager(mapping, dtype) + request.addfinalizer(hybrid_cache_manager_ctx.shutdown) hybrid_cache_manager_gen = create_hybrid_cache_manager(mapping, dtype) + request.addfinalizer(hybrid_cache_manager_gen.shutdown) - cache_transceiver_config = CacheTransceiverConfig(backend=backend, - max_tokens_in_buffer=512) + cache_transceiver_config = CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON", max_tokens_in_buffer=512) dist = Distributed.get(mapping) cache_transceiver_ctx = create_kv_cache_transceiver( @@ -970,6 +1006,7 @@ def test_hybrid_cache_transceiver_cancel_request(backend): AttentionTypeCpp.DEFAULT, cache_transceiver_config, mamba_cache_manager=hybrid_cache_manager_ctx) + request.addfinalizer(cache_transceiver_ctx.shutdown) cache_transceiver_gen = create_kv_cache_transceiver( mapping, @@ -978,6 +1015,7 @@ def test_hybrid_cache_transceiver_cancel_request(backend): AttentionTypeCpp.DEFAULT, cache_transceiver_config, mamba_cache_manager=hybrid_cache_manager_gen) + request.addfinalizer(cache_transceiver_gen.shutdown) fill_hybrid_cache_buffers(hybrid_cache_manager_ctx) @@ -991,6 +1029,9 @@ def test_hybrid_cache_transceiver_cancel_request(backend): sampling_params._get_sampling_config()), is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY) + ctx_request.py_disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="context_only", + disagg_request_id=uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF) scheduled_ctx = ScheduledRequests() scheduled_ctx.context_requests_last_chunk = [ctx_request] @@ -999,13 +1040,6 @@ def test_hybrid_cache_transceiver_cancel_request(backend): # Send ctx request cache_transceiver_ctx.respond_and_send_async(ctx_request) - # Wait for ctx request to be sent - time.sleep(2) - - # Cancel ctx request - is_cancelled = cache_transceiver_ctx.cancel_request(ctx_request) - assert is_cancelled - # Init gen request gen_request = LlmRequest( request_id=0, @@ -1016,6 +1050,14 @@ def test_hybrid_cache_transceiver_cancel_request(backend): is_streaming=False, llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY, context_phase_params=ctx_request.context_phase_params) + gen_request.py_disaggregated_params = tensorrt_llm.DisaggregatedParams( + request_type="generation_only", + disagg_request_id=ctx_request.py_disaggregated_params.disagg_request_id, + ctx_request_id=ctx_request.request_id, + ctx_dp_rank=ctx_request.context_phase_params.ctx_dp_rank, + ctx_info_endpoint=ctx_request.context_phase_params.disagg_info_endpoint, + first_gen_tokens=ctx_request.context_phase_params.first_gen_tokens, + draft_tokens=ctx_request.context_phase_params.draft_tokens) scheduled_gen = ScheduledRequests() scheduled_gen.context_requests_last_chunk = [gen_request] @@ -1024,7 +1066,21 @@ def test_hybrid_cache_transceiver_cancel_request(backend): # Try to receive gen request cache_transceiver_gen.request_and_receive_async(gen_request) - # Block the main thread due to the async operation - time.sleep(2) - cache_transceiver_gen.check_gen_transfer_status(0) - assert gen_request.state == LlmRequestState.DISAGG_TRANS_ERROR + generation_cancelled = [False] + + def cancel_generation_transfer(): + generation_cancelled[0] = cache_transceiver_gen.cancel_request( + gen_request) + + wait_for_transfer_completion(cancel_generation_transfer, + lambda: generation_cancelled[0]) + assert cache_transceiver_gen.check_gen_transfer_complete() + + context_cancelled = [False] + + def cancel_context_transfer(): + context_cancelled[0] = cache_transceiver_ctx.cancel_request(ctx_request) + + wait_for_transfer_completion(cancel_context_transfer, + lambda: context_cancelled[0]) + shutdown_transceivers(cache_transceiver_gen, cache_transceiver_ctx)