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/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9049b1b90a73..d8716fe0086c 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 @@ -93,14 +94,13 @@ 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. + Aggregated serving defaults to V2. Disaggregated serving keeps the legacy + manager routing until the V2 transceiver/page-table adapter is enabled. """ config = model_config.pretrained_config sparse_attn_config = model_config.sparse_attention_config @@ -123,41 +123,64 @@ def get_kv_cache_manager_cls( # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. - if use_py_mamba_cache_manager(): + if is_disagg: + 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 model") + return MixedMambaHybridCacheManager + return CppMambaHybridCacheManager + + 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 +380,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 (disagg, 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' \ 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 +436,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. 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..8f58c828a96c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py @@ -15,7 +15,8 @@ from .llm_request import LlmRequest from .mamba_cache_manager import (BaseMambaCacheManager, - CppMambaHybridCacheManager) + CppMambaHybridCacheManager, + MixedMambaHybridCacheManager) from .resource_manager import KVCacheManager CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver @@ -129,6 +130,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": 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..ce71a2a825ca 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -668,10 +668,6 @@ 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)) @@ -734,6 +730,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 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_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..24b628064c14 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -5,17 +5,27 @@ import os from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch +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 @@ -23,10 +33,113 @@ from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MTPDecodingConfig 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 + ) + + +def test_hybrid_cache_manager_factory_keeps_legacy_disagg_route(monkeypatch): + 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_disagg=True, + ) + is CppMambaHybridCacheManager + ) + assert ( + get_kv_cache_manager_cls( + model_config, + KvCacheConfig(enable_block_reuse=False), + is_disagg=True, + cache_transceiver_config=SimpleNamespace(transceiver_runtime="PYTHON"), + ) + is MixedMambaHybridCacheManager + ) + + +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 +455,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 +614,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 +652,297 @@ 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_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/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)