From 90136ba43af9e749d9ea75158aec91e5f7e5f897 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 16 Sep 2026 22:41:09 +0800 Subject: [PATCH 01/25] [None][refactor] Move Mamba cache behavior into model managers (By Agent) Keep the V2 pool, slot, and snapshot lifecycle shared while moving model-specific state, replay, QSA, and PLE behavior into model modules. Register persistent recurrent buffer roles and typed views in the base manager so warmup cleanup preserves scratch, indices, seeds, and replay bookkeeping. Add lifecycle and cleanup regression coverage. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../attention/backends/sparse/qsa/__init__.py | 35 +- .../backends/sparse/qsa/cache_manager.py | 283 +- .../_torch/models/modeling_qwen4_exp.py | 5 +- .../_torch/modules/fla/cache_manager.py | 357 ++ .../_torch/modules/kimi_kda/cache_manager.py | 680 +++ .../_torch/modules/kimi_kda/kimi_kda_mixer.py | 8 +- .../_torch/modules/mamba/cache_manager.py | 504 ++ .../_torch/modules/mamba/gdn_mixer.py | 6 +- .../_torch/modules/mamba/mamba2_metadata.py | 17 +- .../_torch/modules/mamba/mamba2_mixer.py | 10 +- .../_torch/modules/qwen4_exp/cache_manager.py | 556 ++ tensorrt_llm/_torch/pyexecutor/_util.py | 223 +- .../_torch/pyexecutor/config_utils.py | 138 +- .../kv_cache/kv_cache_manager_v2.py | 69 + .../kv_cache/mamba_cache_manager.py | 4649 ----------------- .../kv_cache/mamba_cache_manager/__init__.py | 98 + .../kv_cache/mamba_cache_manager/common.py | 796 +++ .../kv_cache/mamba_cache_manager/legacy.py | 2320 ++++++++ .../mamba_cache_manager_v2.py | 1279 +++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 14 +- .../kv_cache/test_mamba_cache_manager.py | 468 +- .../kv_cache/test_mamba_model_managers.py | 1061 ++++ .../modeling/test_qsa_runtime_wiring.py | 13 +- .../_torch/modeling/test_qwen4_exp_support.py | 42 +- .../modules/mamba/test_mamba2_metadata.py | 52 +- 25 files changed, 8233 insertions(+), 5450 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/fla/cache_manager.py create mode 100644 tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py create mode 100644 tensorrt_llm/_torch/modules/mamba/cache_manager.py create mode 100644 tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py delete mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py create mode 100644 tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py index 89a6410e1cc7..443edf7415df 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py @@ -1,15 +1,26 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -from .backend import QSATrtllmAttention -from .cache_manager import QSAMambaHybridCacheManagerV2 -from .metadata import QSAAttentionMetadata -from .params import QSASparseMetadataParams, QSASparseParams - -__all__ = [ - "QSAAttentionMetadata", - "QSAMambaHybridCacheManagerV2", - "QSASparseMetadataParams", - "QSASparseParams", - "QSATrtllmAttention", -] +"""Lazy QSA exports; importing geometry does not construct the cache stack.""" + +from importlib import import_module + +_EXPORTS = { + "QSATrtllmAttention": "backend", + "QSAMambaHybridCacheManagerV2": "cache_manager", + "QSAAttentionMetadata": "metadata", + "QSASparseMetadataParams": "params", + "QSASparseParams": "params", +} + + +def __getattr__(name: str): + module = _EXPORTS.get(name) + if module is None: + raise AttributeError(name) + value = getattr(import_module(f"{__name__}.{module}"), name) + globals()[name] = value + return value + + +__all__ = list(_EXPORTS) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py index 7029f111cfdf..5e8bf2e1e8e7 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py @@ -1,284 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""V2 hybrid cache manager for QSA sparse attention.""" -from typing import TYPE_CHECKING, Optional +"""Compatibility imports for the Qwen4-owned QSA cache manager.""" -import torch - -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 -from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor -from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode -from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole - -from .constants import ( - QSA_INDEX_K_CACHE_DTYPE, - QSA_MAIN_KV_ROLES, - QSA_POSITION_CACHE_DTYPE, - QSA_POSITION_COORDINATE_AXES, - QSA_SPARSE_KV_CACHE_DTYPES, +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import QSA_INDEX_POSITION +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + Qwen4ExpHybridCacheManagerV2 as QSAMambaHybridCacheManagerV2, ) -from .params import QSASparseParams - -if TYPE_CHECKING: - from tensorrt_llm.llmapi.llm_args import SparseAttentionConfig - -# Storage-layout constants; model geometry comes from QSASparseParams. -_IDENTITY_PAGE_INDEX_EXPANSION = 1 -_SHARED_POSITION_ROLE_INDEX = 0 -_INDEX_K_ELEMENT_BYTES = torch.empty((), dtype=QSA_INDEX_K_CACHE_DTYPE).element_size() -_POSITION_ELEMENT_BYTES = torch.empty((), dtype=QSA_POSITION_CACHE_DTYPE).element_size() - -# Per-token RoPE/mRoPE coordinates used when raw index keys are compressed. -# They are request state shared by all local QSA layers. -QSA_INDEX_POSITION = DataRole("qsa_index_position") - - -class QSAMambaHybridCacheManagerV2(MambaHybridCacheManagerV2): - """Hybrid GDN/KV manager with lifecycle-coupled sparse side buffers. - - QSA adds a per-layer index-K cache and request-wide position coordinates. - These buffers share the main K/V allocation, eviction, and prefix-reuse - lifecycle. The parent hybrid manager owns GDN/Mamba recurrent state. - - Index geometry is resolved from the same serving/checkpoint configuration - as the indexer. A local fallback could otherwise give the indexer and its - paged cache different shapes or page strides. - """ - - def __init__( - self, - *args, - sparse_attention_config: Optional["SparseAttentionConfig"] = None, - pretrained_config: Optional[object] = None, - layer_mask: Optional[list[bool]] = None, - **kwargs, - ) -> None: - if sparse_attention_config is None: - raise ValueError("sparse_attention_config is required for the QSA cache manager") - if layer_mask is None: - raise ValueError("QSA cache allocation requires the full-attention layer mask") - sparse_params = sparse_attention_config.to_sparse_params( - pretrained_config=pretrained_config - ) - if not isinstance(sparse_params, QSASparseParams): - raise ValueError("QSA cache manager requires QSA sparse parameters") - - # Match the cache's width and head count exactly to the indexer. V2's - # page-index converter already handles roles that share a physical pool. - self.qsa_index_dim = sparse_params.index_head_dim - self.qsa_index_kv_heads = sparse_params.index_kv_heads - # layer_mask uses global IDs; layer_offsets later selects this PP rank. - self.qsa_sparse_layer_ids = [ - layer_idx for layer_idx, enabled in enumerate(layer_mask) if enabled - ] - self.qsa_position_layer_id: Optional[int] = None - super().__init__(*args, layer_mask=layer_mask, **kwargs) - - def _extra_buffers_per_layer( - self, - *, - tokens_per_block: int, - ) -> dict[int, list[BufferConfig]]: - """Register per-layer index K and lifecycle-aligned position pages.""" - index_size = ( - self.qsa_index_kv_heads * self.qsa_index_dim * _INDEX_K_ELEMENT_BYTES * tokens_per_block - ) - local_sparse_layers = [ - layer_id for layer_id in self.qsa_sparse_layer_ids if layer_id in self.layer_offsets - ] - # Coordinates are request-wide, so all local indexers use one view. - # Register the position role once; duplicating it on every sparse layer - # wastes one three-axis int32 page per layer without adding state. - self.qsa_position_layer_id = next(iter(local_sparse_layers), None) - result = { - self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=index_size)] - for layer_id in local_sparse_layers - } - if self.qsa_position_layer_id is not None: - local_idx = self.layer_offsets[self.qsa_position_layer_id] - result[local_idx].append( - BufferConfig( - role=QSA_INDEX_POSITION, - size=( - QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * tokens_per_block - ), - ) - ) - return result - - def get_index_k_buffer( - self, - layer_idx: int, - kv_layout: str = "NHD", - ) -> Optional[torch.Tensor]: - """Return the index-K view using the indexer's resolved geometry.""" - return super().get_index_k_buffer( - layer_idx, - num_heads=self.qsa_index_kv_heads, - head_dim=self.qsa_index_dim, - dtype=QSA_INDEX_K_CACHE_DTYPE, - kv_layout=kv_layout, - ) - - def get_buffers( - self, - layer_idx: int, - kv_layout: str = "NHD", - ) -> Optional[torch.Tensor]: - """Return this layer's adjacent K/V roles with the physical slot stride. - - QSA tables store lifecycle slot IDs. Preserving V2's coalesced stride is - therefore required even when several layers share one physical pool. - """ - if self.dtype not in QSA_SPARSE_KV_CACHE_DTYPES: - # The parent owns packed data and scale-page layouts such as - # NVFP4. QSA returns to the regular backend for these formats. - return super().get_buffers(layer_idx, kv_layout) - if kv_layout not in ("NHD", "HND"): - raise ValueError(f"Unsupported kv_layout: {kv_layout}") - if layer_idx not in self.layer_offsets: - return None - if self.kv_cache_type == CacheTypeCpp.SELFKONLY: - raise NotImplementedError("QSA sparse attention requires both K and V cache buffers") - - layer_offset = self.layer_offsets[layer_idx] - addr_key = self.impl.get_mem_pool_base_address(layer_offset, Role.KEY, PageIndexMode.SHARED) - addr_value = self.impl.get_mem_pool_base_address( - layer_offset, Role.VALUE, PageIndexMode.SHARED - ) - key_stride = self.impl.get_page_stride(layer_offset, Role.KEY) - value_stride = self.impl.get_page_stride(layer_offset, Role.VALUE) - if addr_key + key_stride != addr_value or key_stride != value_stride: - raise RuntimeError( - "QSA K/V buffers are not adjacent equal-sized V2 roles: " - f"layer={layer_idx}, key_stride={key_stride}, " - f"value_stride={value_stride}" - ) - - converter = self.impl.get_page_index_converter(layer_offset, Role.KEY) - scale = int(converter.scale) - if scale < QSA_MAIN_KV_ROLES: - raise RuntimeError( - f"QSA K/V roles do not share one lifecycle slot: layer={layer_idx}, {scale=}" - ) - if int(converter.expansion) != _IDENTITY_PAGE_INDEX_EXPANSION: - raise RuntimeError("QSA does not support expanded V2 page indices") - layer_offset_pages = int(converter.layer_offset) - page_upper = self.impl.get_page_index_upper_bound(layer_offset, Role.KEY) - num_slots_total = page_upper + layer_offset_pages - if num_slots_total % scale != 0: - raise RuntimeError( - "QSA K/V page mapping is inconsistent: " - f"{num_slots_total=} is not divisible by {scale=}" - ) - num_slots = num_slots_total // scale - - torch_dtype = binding_to_torch_dtype(self.dtype) - head_dim = self.head_dim_per_layer[layer_offset] - num_heads = self.num_kv_heads_per_layer[layer_offset] - if kv_layout == "NHD": - page_shape = [self.tokens_per_block, num_heads, head_dim] - else: - page_shape = [num_heads, self.tokens_per_block, head_dim] - full_view = convert_to_torch_tensor( - TensorWrapper( - addr_key, - torch_dtype, - [num_slots, scale, *page_shape], - ) - ) - return full_view[:, :QSA_MAIN_KV_ROLES] - - @torch.compiler.disable - def get_qsa_position_buffer(self) -> Optional[torch.Tensor]: - """Return per-token three-axis RoPE/mRoPE position coordinates.""" - if self.qsa_position_layer_id is None: - return None - local_idx = self.layer_offsets[self.qsa_position_layer_id] - addr = self.impl.get_mem_pool_base_address( - local_idx, - QSA_INDEX_POSITION, - PageIndexMode.SHARED, - ) - page_stride = self.impl.get_page_stride(local_idx, QSA_INDEX_POSITION) - expected_stride = ( - QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * self.tokens_per_block - ) - if page_stride != expected_stride: - raise RuntimeError( - "QSA position-cache page stride mismatch: " - f"expected {expected_stride}, got {page_stride}" - ) - converter = self.impl.get_page_index_converter( - local_idx, - QSA_INDEX_POSITION, - ) - page_upper = self.impl.get_page_index_upper_bound( - local_idx, - QSA_INDEX_POSITION, - ) - num_pages_with_offset = page_upper + int(converter.layer_offset) - scale = int(converter.scale) - if scale <= 0: - raise RuntimeError(f"QSA received an invalid position page-index scale: {scale}") - if int(converter.expansion) != _IDENTITY_PAGE_INDEX_EXPANSION: - raise RuntimeError("QSA does not support expanded V2 position page indices") - if num_pages_with_offset % scale != 0: - raise RuntimeError("QSA position-cache page mapping is inconsistent") - num_slots = num_pages_with_offset // scale - full_view = convert_to_torch_tensor( - TensorWrapper( - addr, - QSA_POSITION_CACHE_DTYPE, - [ - num_slots, - scale, - self.tokens_per_block, - QSA_POSITION_COORDINATE_AXES, - ], - ) - ) - return full_view[:, _SHARED_POSITION_ROLE_INDEX] - - def get_qsa_attention_pool_layout(self) -> tuple[int, int]: - """Return the shared attention-pool index and V2 page-index scale. - - One graph-stable QSA block table is reused by every local sparse layer, - so their main-K roles must resolve lifecycle slots through the same - attention pool and converter. - """ - if self.qsa_position_layer_id is None: - raise RuntimeError("QSA cache manager has no local sparse layer") - local_idx = self.layer_offsets[self.qsa_position_layer_id] - pool_id = self.layer_to_pool_mapping_dict[local_idx] - if pool_id >= self.num_attention_op_pools: - raise RuntimeError( - "QSA K/V pool is not represented in the attention block table: " - f"pool {pool_id}, attention pools {self.num_attention_op_pools}" - ) - converter = self.impl.get_page_index_converter(local_idx, Role.KEY) - scale = int(converter.scale) - for layer_id in self.qsa_sparse_layer_ids: - if layer_id not in self.layer_offsets: - continue - candidate_idx = self.layer_offsets[layer_id] - candidate_pool = self.layer_to_pool_mapping_dict[candidate_idx] - candidate_converter = self.impl.get_page_index_converter(candidate_idx, Role.KEY) - if candidate_pool != pool_id or int(candidate_converter.scale) != scale: - raise RuntimeError( - "QSA local layers do not share one attention page mapping: " - f"layer {layer_id} uses pool/scale " - f"{candidate_pool}/{int(candidate_converter.scale)}, expected " - f"{pool_id}/{scale}" - ) - if int(candidate_converter.expansion) != _IDENTITY_PAGE_INDEX_EXPANSION: - raise RuntimeError("QSA does not support expanded V2 page indices") - if scale <= 0: - raise RuntimeError(f"QSA received an invalid V2 page-index scale: {scale}") - return pool_id, scale - -__all__ = ["QSAMambaHybridCacheManagerV2"] +__all__ = ["QSA_INDEX_POSITION", "QSAMambaHybridCacheManagerV2"] diff --git a/tensorrt_llm/_torch/models/modeling_qwen4_exp.py b/tensorrt_llm/_torch/models/modeling_qwen4_exp.py index 9f0a24ff8ef0..60fee5ca2acb 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen4_exp.py +++ b/tensorrt_llm/_torch/models/modeling_qwen4_exp.py @@ -42,10 +42,11 @@ from ..modules.mamba.layernorm_gated import rms_norm_gated_token_major from ..modules.mamba.mamba2_metadata import Mamba2Metadata from ..modules.multi_stream_utils import maybe_execute_in_parallel +from ..modules.qwen4_exp.cache_manager import get_qwen4_exp_ple_layer_mask from ..modules.qwen4_exp.hyper_connection import HCResidual, Qwen4ExpHyperConnection from ..modules.qwen4_exp.ple import PLEMetadata, Qwen4ExpPLE from ..modules.rms_norm import RMSNorm -from ..pyexecutor.config_utils import get_qwen3_hybrid_layer_types, get_qwen4_exp_ple_layer_mask +from ..pyexecutor.config_utils import get_qwen3_hybrid_layer_types from ..speculative import SpecMetadata from ..utils import AuxStreamType, EventType, create_lm_head_tp_mapping from .checkpoints.base_weight_mapper import BaseWeightMapper @@ -473,7 +474,7 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): # allocated to the cache slot count and carried in place across the # prefill->decode boundary, mirroring the GDN conv/ssm per-slot pools and # indexed by the SAME mamba ``state_indices`` (see ``_prepare_ple_state``). - # Shapes match ``config_utils.extract_qwen4_exp_ple_cache_params`` + # Shapes match the model-owned ``extract_qwen4_exp_ple_cache_params``. # (conv_state_shape / ngram_context_len). self._ple_conv_state: Optional[torch.Tensor] = None self._ple_ngram_context: Optional[torch.Tensor] = None diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py new file mode 100644 index 000000000000..f7b82a6f5923 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -0,0 +1,357 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Gated-DeltaNet replay integration for Mamba cache managers.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass + +import torch + +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override + +from tensorrt_llm._torch.modules.mamba.cache_manager import Mamba2ReplayLayerCache, ReplayHistory +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + IntermediateState, + MambaAcceptanceBatch, + MambaLayerCache, + MambaStateLayout, + ReplayStateUpdateMetadata, + _mamba_effective_tp_size, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( + MambaHybridCacheManagerV2, +) +from tensorrt_llm._torch.utils import is_gdn_replay_enabled +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger + +from .cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers, +) + + +@dataclass(frozen=True, kw_only=True) +class GDNReplayLayerCache(Mamba2ReplayLayerCache): + """GDN cached-replay tensors for one recurrent layer.""" + + +class GDNReplayState(ReplayHistory): + """Own GDN all-layer cached-replay checkpoint commits.""" + + @override + def __init__(self, tokens_per_gen_step: int) -> None: + super().__init__(tokens_per_gen_step) + self._state_descriptors: torch.Tensor | None = None + self._state_strides: tuple[int, int, int] | None = None + + @override + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + super().bind(context, ssm_states, conv_states) + states = self._ssm_states + if not states: + return + state_shape = list(context.ssm_state_shape) + expected_inner_strides = ( + state_shape[1] * state_shape[2], + state_shape[2], + 1, + ) + reference = states[0] + for state in states: + if ( + state.dtype != reference.dtype + or state.device != reference.device + or list(state.shape[1:]) != state_shape + or state.stride()[1:] != expected_inner_strides + ): + raise RuntimeError( + "GDN cached replay requires V2 SSM layers with matching " + "dtype, device, shape, and dense inner dimensions." + ) + pointers = [state.data_ptr() for state in states] + slot_stride = reference.stride(0) + layer_stride_bytes = pointers[1] - pointers[0] if len(pointers) > 1 else 0 + has_affine_layout = ( + layer_stride_bytes % reference.element_size() == 0 + and all(state.stride(0) == slot_stride for state in states) + and all( + pointer == pointers[0] + layer * layer_stride_bytes + for layer, pointer in enumerate(pointers) + ) + ) + if has_affine_layout: + self._state_strides = ( + len(pointers), + layer_stride_bytes // reference.element_size(), + slot_stride, + ) + else: + logger.warning_once( + "V2 GDN state views are not affine; using indirect replay checkpoint addressing", + key="gdn_cached_replay_v2_indirect_state_layout", + ) + self._state_descriptors = torch.tensor( + [(state.data_ptr(), state.stride(0)) for state in states], + dtype=torch.int64, + device=reference.device, + ) + state_layout = "affine" if self._state_strides is not None else "indirect" + logger.info_once( + "Configured GDN cached replay commit mode for V2: small-batch fused, " + f"large-batch all-layer; state layout: {state_layout}", + key="gdn_cached_replay_v2_commit_mode_fused", + ) + + @property + def has_bound_states(self) -> bool: + return self._state_descriptors is not None or self._state_strides is not None + + @override + def update(self, batch: MambaAcceptanceBatch) -> None: + num_decodes = batch.num_generations + if num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: + self.commit_all_layers(batch.attention_metadata, num_decodes) + super().update(batch) + + @override + def make_layer_cache( + self, + layer_offset: int, + conv: torch.Tensor, + temporal: torch.Tensor, + ) -> MambaLayerCache: + return GDNReplayLayerCache( + conv=conv, + temporal=temporal, + **self._layer_cache_fields(layer_offset), + ) + + def commit_all_layers(self, attention_metadata: object, num_decodes: int) -> None: + """Commit every local GDN checkpoint using one partitioned launch.""" + states = self._ssm_states + if ( + not states + or (self._state_descriptors is None and self._state_strides is None) + or self.old_x is None + or self.old_B is None + or self.old_dt is None + ): + raise RuntimeError( + "GDN cached replay all-layer commit requires V2 replay state buffers." + ) + + mamba_metadata = attention_metadata.mamba_metadata + if mamba_metadata.replay_num_decodes != num_decodes: + raise RuntimeError( + "GDN replay metadata contains " + f"{mamba_metadata.replay_num_decodes} decode requests, " + f"but state update received {num_decodes}." + ) + state_strides = self._state_strides + commit_gdn_cached_replay_history_layers( + ssm_states=states[0], + ssm_state_descriptors=self._state_descriptors, + ssm_state_num_layers=(state_strides[0] if state_strides is not None else None), + ssm_state_layer_stride=(state_strides[1] if state_strides is not None else None), + ssm_state_slot_stride=(state_strides[2] if state_strides is not None else None), + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=mamba_metadata.replay_work_items[:num_decodes], + n_writes=mamba_metadata.replay_n_writes, + history_size=self.replay_history_size, + ) + + @override + def shutdown(self) -> None: + super().shutdown() + self._state_descriptors = None + self._state_strides = None + + +def select_gdn_replay_state( + *, + spec_config: object | None, + ssm_cache_dtype: torch.dtype, + manager_cls: type, +) -> GDNReplayState | None: + """Apply GDN replay feature gates and return the policy when eligible.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + CppMambaHybridCacheManager, + MambaHybridCacheManagerV2, + ) + + use_replay = spec_config is not None and get_sm_version() >= 80 + if spec_config is None: + logger.info("GDN replay kernel requires speculative decoding; using non-replay path") + elif spec_config.tokens_per_gen_step > 8: + logger.info( + "GDN cached replay supports at most 8 tokens per generation step; using non-replay path" + ) + use_replay = False + if spec_config is not None and ( + getattr(spec_config, "eagle_choices", None) is not None + or getattr(spec_config, "use_dynamic_tree", False) + ): + logger.info("GDN replay kernel incompatible with tree attention; using legacy MTP path") + use_replay = False + if ssm_cache_dtype not in (torch.float32, torch.bfloat16, torch.float16): + logger.info( + "GDN replay kernel does not support quantized SSM cache dtype " + f"{ssm_cache_dtype}; using legacy MTP path" + ) + use_replay = False + if not is_gdn_replay_enabled(): + use_replay = False + if use_replay and not issubclass( + manager_cls, + (CppMambaHybridCacheManager, MambaHybridCacheManagerV2), + ): + logger.info( + "GDN replay requires C++ V1 or V2 Mamba cache manager; " + f"{manager_cls.__name__} was selected, so the non-replay MTP path will be used" + ) + use_replay = False + logger.info(f"GDN replay state update: {'ENABLED' if use_replay else 'DISABLED'}") + return GDNReplayState(spec_config.tokens_per_gen_step) if use_replay else None + + +def create_gdn_state( + manager: MambaHybridCacheManagerV2, use_replay: bool | None +) -> IntermediateState | GDNReplayState: + """Shared selection for the Qwen3.5 and Qwen4 cache managers.""" + if use_replay is False: + return IntermediateState() + if use_replay is True: + if manager.spec_config is None: + raise ValueError("GDN replay requires speculative decoding") + return GDNReplayState(manager.spec_config.tokens_per_gen_step) + return ( + select_gdn_replay_state( + spec_config=manager.spec_config, + ssm_cache_dtype=manager.ssm_state_dtype, + manager_cls=type(manager), + ) + or IntermediateState() + ) + + +def validate_gdn_layout( + manager: MambaHybridCacheManagerV2, state: IntermediateState | GDNReplayState +) -> None: + if isinstance(state, GDNReplayState): + if manager.local_num_mamba_layers and manager._global_n_groups % _mamba_effective_tp_size( + manager._state_layout.mapping + ): + raise ValueError("GDN replay groups must be divisible by the effective TP size") + state.validate(manager._state_layout) + + +class Qwen35HybridCacheManagerV2(MambaHybridCacheManagerV2): + """Qwen3.5/Qwen3Next recurrent state with shared GDN replay.""" + + @override + def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) -> None: + self._requested_replay = use_replay_state_update + kwargs.setdefault("conv_state_layout", "q_k_v") + super().__init__(*args, **kwargs) + + @override + def _initialize_model_state(self) -> None: + self._speculative_state = create_gdn_state(self, self._requested_replay) + + validate_gdn_layout(self, self._speculative_state) + + @override + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + state = self._speculative_state + return state.get_replay_metadata() if isinstance(state, GDNReplayState) else None + + @property + def intermediate_state_indices(self) -> torch.Tensor | None: + return self._speculative_state.intermediate_indices + + @property + def intermediate_ssm_states(self) -> torch.Tensor | None: + state = self._speculative_state + return None if isinstance(state, GDNReplayState) else state.intermediate_ssm + + @property + def intermediate_conv_states(self) -> torch.Tensor | None: + return self._speculative_state.intermediate_conv + + @property + def use_gdn_cached_replay_all_layer_commit(self) -> bool: + state = self._speculative_state + return isinstance(state, GDNReplayState) and state.has_bound_states + + @override + def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: + if self.spec_config is None: + return super().mamba_layer_cache(layer_idx) + return self._speculative_state.make_layer_cache( + self.mamba_layer_offsets[layer_idx], + self.get_conv_states(layer_idx), + self.get_ssm_states(layer_idx), + ) + + @override + def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + self._speculative_state.reset_slots(slots, host_slots) + + @override + def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: + self._speculative_state.update(batch) + + @override + def _setup_model_state(self) -> None: + self._speculative_state.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) + + @override + def _shutdown_model_state(self) -> None: + self._speculative_state.shutdown() + + +def get_gdn_cache_params(config, *, spec_config=None, quant_config=None): + """Shared Qwen3.5/Qwen4 GDN geometry; attention and PLE stay model-owned.""" + from tensorrt_llm._torch.pyexecutor.config_utils import ( + build_mamba_kv_cache_params, + get_qwen3_hybrid_layer_masks, + ) + + attention, recurrent = get_qwen3_hybrid_layer_masks(config) + return build_mamba_kv_cache_params( + config, + state_size=config.linear_key_head_dim, + conv_kernel=config.linear_conv_kernel_dim, + num_heads=config.linear_num_value_heads, + n_groups=config.linear_num_key_heads, + head_dim=config.linear_value_head_dim, + mamba_mask=recurrent, + target_full_attn_mask=attention, + spec_config=spec_config, + quant_config=quant_config, + ) + + +__all__ = [ + "GDNReplayState", + "GDNReplayLayerCache", + "Qwen35HybridCacheManagerV2", + "create_gdn_state", + "select_gdn_replay_state", + "validate_gdn_layout", + "get_gdn_cache_params", +] diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py new file mode 100644 index 000000000000..6d252641a606 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -0,0 +1,680 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""K3-owned fused-verify cache payload and replay state.""" + +from __future__ import annotations + +import sys +from dataclasses import dataclass, field + +import torch + +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override + +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + allocate_mamba_seed_buffer, + compute_deterministic_mamba_seed, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + IntermediateState, + MambaAcceptanceBatch, + MambaLayerCache, + MambaStateLayout, + SpeculativeMambaLayerCache, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( + MambaHybridCacheManagerV2, +) +from tensorrt_llm.logger import logger + + +@dataclass(frozen=True, kw_only=True) +class KDAReplayLayerCache(SpeculativeMambaLayerCache): + """Per-layer KDA replay tensors consumed by fused multi-token verify.""" + + prev_num_accepted_tokens: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + kda_conv_q: torch.Tensor | None = None + kda_conv_k: torch.Tensor | None = None + kda_conv_v: torch.Tensor | None = None + kda_qkg_cache: torch.Tensor | None = None + kda_v_cache: torch.Tensor | None = None + kda_beta_cache: torch.Tensor | None = None + + @property + def has_kda_replay_caches(self) -> bool: + return self.kda_qkg_cache is not None + + def commit_conv_window(self, slot_indices: torch.Tensor, conv_pool: torch.Tensor) -> None: + """Seed the replay convolution windows from the live pool.""" + from ._kda_kernels import copy_kda_replay_conv_window + + copy_kda_replay_conv_window( + conv_pool, + self.kda_conv_q, + self.kda_conv_k, + self.kda_conv_v, + slot_indices, + ) + + +def allocate_kda_replay_fields( + *, + num_local_layers: int, + cache_size: int, + num_speculative_tokens: int, + conv_dim: int, + conv_kernel_size: int, + num_heads: int, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], list[torch.Tensor]]: + """Allocate the complete KDA fused-verify replay payload.""" + if conv_dim % 3 != 0: + raise ValueError("KDA replay caches require q/k/v convolution sections of equal width") + section_dim = conv_dim // 3 + extended_window = conv_kernel_size - 1 + num_speculative_tokens + + def allocate_conv_cache() -> torch.Tensor: + return torch.zeros( + num_local_layers, + cache_size, + extended_window, + section_dim, + dtype=torch.float32, + device=device, + ).transpose(-1, -2) + + fields = { + "prev_num_accepted_tokens": torch.zeros(cache_size, dtype=torch.int32, device=device), + "kda_conv_q": allocate_conv_cache(), + "kda_conv_k": allocate_conv_cache(), + "kda_conv_v": allocate_conv_cache(), + "kda_qkg_cache": torch.zeros( + num_local_layers, + cache_size, + num_speculative_tokens, + 3, + section_dim, + dtype=torch.float32, + device=device, + ), + "kda_v_cache": torch.zeros( + num_local_layers, + cache_size, + num_speculative_tokens, + section_dim, + dtype=torch.float32, + device=device, + ), + "kda_beta_cache": torch.zeros( + num_local_layers, + cache_size, + num_speculative_tokens, + num_heads, + dtype=torch.float32, + device=device, + ), + } + scratch = [ + fields["kda_conv_q"], + fields["kda_conv_k"], + fields["kda_conv_v"], + fields["kda_qkg_cache"], + fields["kda_v_cache"], + fields["kda_beta_cache"], + ] + return fields, scratch + + +class KDAReplayState: + """Own all V2 state used by the KDA fused multi-token verifier.""" + + def __init__( + self, + num_speculative_tokens: int, + ) -> None: + self.num_speculative_tokens = num_speculative_tokens + self._conv_states: tuple[torch.Tensor, ...] = () + self._committed_window = 0 + self.prev_num_accepted_tokens: torch.Tensor | None = None + self.kda_conv_q: torch.Tensor | None = None + self.kda_conv_k: torch.Tensor | None = None + self.kda_conv_v: torch.Tensor | None = None + self.kda_qkg_cache: torch.Tensor | None = None + self.kda_v_cache: torch.Tensor | None = None + self.kda_beta_cache: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + self.rand_seed: torch.Tensor | None = None + self._seed_request_counter = 0 + self._seed_rank_offset = 0 + self._conv_section_dims: tuple[int, ...] = () + + def validate(self, context: MambaStateLayout) -> None: + if context.spec_config is None: + raise ValueError("KDA replay caches require speculative decoding") + expected_num_spec = context.spec_config.tokens_per_gen_step - 1 + if self.num_speculative_tokens != expected_num_spec: + raise ValueError( + "KDA replay cache width " + f"({self.num_speculative_tokens}) must match the draft length " + f"({expected_num_spec})" + ) + if not context.mamba_pp_layers: + return + if context.conv_state_layout != "q_k_v": + raise ValueError("KDA replay requires conv_state_layout='q_k_v'") + if len(context.conv_section_dims) != 3 or len(set(context.conv_section_dims)) != 1: + raise ValueError( + "KDA replay caches require equal [Q | K | V] convolution-state sections" + ) + + def bytes_per_slot(self, context: MambaStateLayout, layer_id: int) -> int: + if not context.mamba_pp_layers: + return 0 + section_dim = context.conv_section_dims[0] + committed_window = context.conv_state_shape[1] + extended_window = committed_window + self.num_speculative_tokens + num_heads = context.ssm_state_shape[0] + float_elements = ( + 3 * section_dim * extended_window + + self.num_speculative_tokens * 3 * section_dim + + self.num_speculative_tokens * section_dim + + self.num_speculative_tokens * num_heads + ) + shared_bytes = 0 + if layer_id == context.mamba_pp_layers[0]: + shared_bytes = 4 + if context.stochastic_rounding: + shared_bytes += 8 + return float_elements * 4 + shared_bytes + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + """Borrow per-layer conv views [slot, conv_dim, window] and allocate scratch.""" + self._seed_rank_offset = context.seed_rank_offset + self._conv_section_dims = context.conv_section_dims + self._conv_states = tuple(conv_states) + self._committed_window = context.conv_state_shape[1] if conv_states else 0 + states = ssm_states + if not states: + return + + cache_size = context.slot_capacity + if cache_size is None: + raise RuntimeError("KDA replay requires the allocated V2 slot capacity") + device = states[0].device + fields, _ = allocate_kda_replay_fields( + num_local_layers=len(context.mamba_pp_layers), + cache_size=cache_size, + num_speculative_tokens=self.num_speculative_tokens, + conv_dim=context.conv_state_shape[0], + conv_kernel_size=context.conv_state_shape[1] + 1, + num_heads=context.ssm_state_shape[0], + device=device, + ) + for name, value in fields.items(): + setattr(self, name, value) + self.intermediate_indices = torch.arange( + context.max_batch_size, + dtype=torch.int32, + device=device, + ) + if context.stochastic_rounding: + self.rand_seed = allocate_mamba_seed_buffer( + cache_size, + context.seed_rank_offset, + device, + ) + logger.info(f"Mamba Cache (kda-replay) is allocated for {cache_size} state slots") + + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + return { + "intermediate_conv_window": None, + "intermediate_ssm": None, + "mamba_ssm_rand_seed": self.rand_seed, + "prev_num_accepted_tokens": self.prev_num_accepted_tokens, + "kda_conv_q": None if self.kda_conv_q is None else self.kda_conv_q[layer_offset], + "kda_conv_k": None if self.kda_conv_k is None else self.kda_conv_k[layer_offset], + "kda_conv_v": None if self.kda_conv_v is None else self.kda_conv_v[layer_offset], + "kda_qkg_cache": ( + None if self.kda_qkg_cache is None else self.kda_qkg_cache[layer_offset] + ), + "kda_v_cache": None if self.kda_v_cache is None else self.kda_v_cache[layer_offset], + "kda_beta_cache": ( + None if self.kda_beta_cache is None else self.kda_beta_cache[layer_offset] + ), + } + + def make_layer_cache( + self, + layer_offset: int, + conv: torch.Tensor, + temporal: torch.Tensor, + ) -> MambaLayerCache: + return KDAReplayLayerCache( + conv=conv, + temporal=temporal, + **self._layer_cache_fields(layer_offset), + ) + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + if self.prev_num_accepted_tokens is None or slots.numel() == 0: + return + self.prev_num_accepted_tokens[slots] = 0 + if self.rand_seed is not None: + self._seed_request_counter += 1 + seeds = [ + compute_deterministic_mamba_seed( + self._seed_request_counter, + slot, + self._seed_rank_offset, + ) + for slot in host_slots + ] + self.rand_seed[slots] = torch.tensor( + seeds, + dtype=torch.int64, + device=self.rand_seed.device, + ) + + def relocate_slots(self, old_slots: list[int], new_slots: list[int]) -> None: + if self.prev_num_accepted_tokens is None: + return + moves = [(old, new) for old, new in zip(old_slots, new_slots) if old >= 0 and old != new] + if not moves: + return + device = self.prev_num_accepted_tokens.device + source_slots = torch.tensor( + [old for old, _ in moves], + dtype=torch.long, + device=device, + ) + destination_slots = torch.tensor( + [new for _, new in moves], + dtype=torch.long, + device=device, + ) + for replay_buffer in self._layer_replay_buffers(): + replay_buffer.index_copy_( + 1, + destination_slots, + replay_buffer.index_select(1, source_slots), + ) + self.prev_num_accepted_tokens.index_copy_( + 0, + destination_slots, + self.prev_num_accepted_tokens.index_select(0, source_slots), + ) + if self.rand_seed is not None: + self.rand_seed.index_copy_( + 0, + destination_slots, + self.rand_seed.index_select(0, source_slots), + ) + + def update(self, batch: MambaAcceptanceBatch) -> None: + self.record_acceptance( + batch.destination_state_indices, + batch.num_accepted_tokens - 1, + batch.is_dummy_request, + ) + + @torch.inference_mode() + def seed_transferred_slots(self, state_indices: torch.Tensor) -> None: + """Seed completed-transfer slots from borrowed persistent conv views.""" + if self.prev_num_accepted_tokens is None or state_indices.numel() == 0: + return + committed_window = self._committed_window + section_offsets = [0] + for section_dim in self._conv_section_dims: + section_offsets.append(section_offsets[-1] + section_dim) + + replay_conv_buffers = (self.kda_conv_q, self.kda_conv_k, self.kda_conv_v) + for layer_offset, conv_state in enumerate(self._conv_states): + selected_conv = conv_state.index_select(0, state_indices) + for section_idx, replay_buffer in enumerate(replay_conv_buffers): + if replay_buffer is None: + raise RuntimeError("KDA replay convolution buffers are not bound") + replay_layer = replay_buffer[layer_offset] + replay_layer.index_fill_(0, state_indices, 0) + section = selected_conv[ + :, + section_offsets[section_idx] : section_offsets[section_idx + 1], + :, + ] + replay_layer[:, :, :committed_window].index_copy_( + 0, + state_indices, + section.to(replay_layer.dtype), + ) + for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, self.kda_beta_cache): + if replay_buffer is None: + raise RuntimeError("KDA replay history buffers are not bound") + replay_buffer.index_fill_(1, state_indices, 0) + self.prev_num_accepted_tokens[state_indices] = 0 + + def record_acceptance( + self, + state_indices: torch.Tensor, + num_accepted_draft_tokens: torch.Tensor, + is_dummy_request: torch.Tensor | None, + ) -> None: + """Record accepted drafts, preserving dummy rows; each input has shape [batch].""" + if self.prev_num_accepted_tokens is None: + raise RuntimeError("KDA replay acceptance requires replay bookkeeping") + slots = state_indices.to(torch.long) + accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) + if is_dummy_request is not None: + current = self.prev_num_accepted_tokens[slots] + accepted = torch.where(is_dummy_request, current, accepted) + self.prev_num_accepted_tokens[slots] = accepted + + def _layer_replay_buffers(self) -> tuple[torch.Tensor, ...]: + buffers = ( + self.kda_conv_q, + self.kda_conv_k, + self.kda_conv_v, + self.kda_qkg_cache, + self.kda_v_cache, + self.kda_beta_cache, + ) + if any(buffer is None for buffer in buffers): + raise RuntimeError("KDA replay buffers are not bound") + return tuple(buffer for buffer in buffers if buffer is not None) + + def shutdown(self) -> None: + self._conv_states = () + self.prev_num_accepted_tokens = None + self.kda_conv_q = None + self.kda_conv_k = None + self.kda_conv_v = None + self.kda_qkg_cache = None + self.kda_v_cache = None + self.kda_beta_cache = None + self.intermediate_indices = None + self.rand_seed = None + + +def select_kda_replay_state( + *, + spec_config: object | None, + manager_cls: type, +) -> KDAReplayState | None: + """Apply KDA fused-verify and manager gates at the feature owner.""" + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + MambaHybridCacheManagerV2, + MixedMambaHybridCacheManager, + ) + + manager_supports_replay = issubclass( + manager_cls, + (MambaHybridCacheManagerV2, MixedMambaHybridCacheManager), + ) + num_spec = get_kda_replay_num_spec( + spec_config, + manager_supports_replay=manager_supports_replay, + ) + if num_spec is None: + return None + return KDAReplayState(num_spec) + + +def get_kda_replay_num_spec( + spec_config: object | None, *, manager_supports_replay: bool +) -> int | None: + """Return the fused-verify replay width when its complete gate passes.""" + if spec_config is None or not manager_supports_replay: + return None + from ._kda_kernels import is_kda_mtp_verify_available + + if not is_kda_mtp_verify_available(): + return None + return spec_config.tokens_per_gen_step - 1 + + +class KimiK3HybridCacheManagerV2(MambaHybridCacheManagerV2): + """KDA/MLA cache layout and K3-only replay lifecycle handling.""" + + @override + def __init__( + self, + *args, + kda_replay_num_spec: int | None = None, + use_replay_state_update: bool = False, + **kwargs, + ) -> None: + if use_replay_state_update: + raise ValueError("KDA replay and Mamba2 replay are mutually exclusive") + self._requested_num_spec = kda_replay_num_spec + self._kda_replay: KDAReplayState | None = None + kwargs.setdefault("conv_state_layout", "q_k_v") + super().__init__(*args, **kwargs) + + @override + def _initialize_model_state(self) -> None: + num_spec = self._requested_num_spec + if num_spec is None: + num_spec = get_kda_replay_num_spec(self.spec_config, manager_supports_replay=True) + if num_spec is not None: + self._kda_replay = KDAReplayState(num_spec) + else: + self._speculative_state = IntermediateState() + + if self._kda_replay is not None: + self._kda_replay.validate(self._state_layout) + + @property + def use_kda_replay_update(self) -> bool: + return self._kda_replay is not None + + @override + def _extra_scratch_bytes_per_slot(self) -> int: + if self._kda_replay is None: + return 0 + return sum( + self._kda_replay.bytes_per_slot(self._state_layout, layer_id) + for layer_id in self.mamba_pp_layers + ) + + @override + def _setup_model_state(self) -> None: + if self._kda_replay is None: + self._speculative_state.bind( + self._state_layout, self.all_ssm_states, self.all_conv_states + ) + else: + self._kda_replay.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) + + @property + def intermediate_state_indices(self) -> torch.Tensor | None: + if self._kda_replay is not None: + return self._kda_replay.intermediate_indices + return self._speculative_state.intermediate_indices + + @property + def intermediate_ssm_states(self) -> torch.Tensor | None: + if self._kda_replay is not None: + return None + return self._speculative_state.intermediate_ssm + + @property + def intermediate_conv_states(self) -> torch.Tensor | None: + if self._kda_replay is not None: + return None + return self._speculative_state.intermediate_conv + + @override + def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: + if self._kda_replay is None: + if self.spec_config is None: + return super().mamba_layer_cache(layer_idx) + return self._speculative_state.make_layer_cache( + self.mamba_layer_offsets[layer_idx], + self.get_conv_states(layer_idx), + self.get_ssm_states(layer_idx), + ) + return self._kda_replay.make_layer_cache( + self.mamba_layer_offsets[layer_idx], + self.get_conv_states(layer_idx), + self.get_ssm_states(layer_idx), + ) + + @override + def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + if self._kda_replay is None: + self._speculative_state.reset_slots(slots, host_slots) + else: + self._kda_replay.reset_slots(slots, host_slots) + + @override + def _on_state_slots_relocated(self, old_slots: list[int], new_slots: list[int]) -> None: + if self._kda_replay is None: + return + self._kda_replay.relocate_slots(old_slots, new_slots) + fresh_slots = [new for old, new in zip(old_slots, new_slots) if old < 0] + if fresh_slots: + slots = torch.tensor( + fresh_slots, dtype=torch.long, device=self.cuda_state_indices.device + ) + self._kda_replay.reset_slots(slots, fresh_slots) + + @override + def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: + if self._kda_replay is None: + self._speculative_state.update(batch) + else: + self._kda_replay.update(batch) + + @override + def update_resources( + self, scheduled_batch, attn_metadata=None, kv_cache_dtype_byte_size=None + ) -> None: + super().update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size) + if ( + self.local_num_mamba_layers + and self._kda_replay is not None + and getattr(self.spec_config, "decoding_type", None) == "NGram" + ): + self._record_replay_request_acceptance(scheduled_batch) + + def _record_replay_request_acceptance(self, scheduled_batch: object) -> None: + replay = self._kda_replay + if replay.prev_num_accepted_tokens is None: + return + generation_requests = scheduled_batch.generation_requests + drafted_requests = [ + request + for request in generation_requests + if request.py_draft_tokens is not None and len(request.py_draft_tokens) > 0 + ] + if not drafted_requests: + return + if len(drafted_requests) != len(generation_requests): + raise RuntimeError( + "Mixed drafted/undrafted generation batch is not supported " + "for KDA replay bookkeeping" + ) + state_index_map = self._request_id_to_state_index + dummy_map = self._request_id_to_is_dummy + active_requests = [ + request for request in generation_requests if request.py_request_id in state_index_map + ] + if not active_requests: + return + device = replay.prev_num_accepted_tokens.device + state_indices = torch.tensor( + [state_index_map[request.py_request_id] for request in active_requests], + dtype=torch.int32, + device=device, + ) + accepted_drafts = torch.tensor( + [request.py_num_accepted_draft_tokens for request in active_requests], + dtype=torch.int32, + device=device, + ) + is_dummy_request = torch.tensor( + [dummy_map.get(request.py_request_id, False) for request in active_requests], + dtype=torch.bool, + device=device, + ) + replay.record_acceptance(state_indices, accepted_drafts, is_dummy_request) + + @override + def on_state_transfer_complete(self, request_ids: list[int]) -> None: + replay = self._kda_replay + if replay is None or replay.prev_num_accepted_tokens is None: + return + slots = sorted( + { + self._request_id_to_state_index[request_id] + for request_id in request_ids + if request_id in self._request_id_to_state_index + } + ) + if slots: + replay.seed_transferred_slots( + torch.tensor(slots, dtype=torch.long, device=replay.prev_num_accepted_tokens.device) + ) + + def seed_kda_replay_caches_for_disagg_gen(self, request_ids: list[int]) -> None: + """Compatibility entry point; the executor uses the generic transfer hook.""" + self.on_state_transfer_complete(request_ids) + + @override + def _shutdown_model_state(self) -> None: + if self._kda_replay is not None: + self._kda_replay.shutdown() + else: + self._speculative_state.shutdown() + + +def get_kimi_cache_params(config, *, spec_config=None, quant_config=None): + """KDA uses equal Q/K/V sections and FP32 recurrent state alongside MLA.""" + from tensorrt_llm._torch.pyexecutor.config_utils import ( + build_mamba_kv_cache_params, + get_kimi_linear_layer_masks, + unwrap_kimi_text_config, + ) + + linear = unwrap_kimi_text_config(config).linear_attn_config + attention, recurrent = get_kimi_linear_layer_masks(config) + params = build_mamba_kv_cache_params( + config, + state_size=linear["head_dim"], + conv_kernel=linear["short_conv_kernel_size"], + num_heads=linear["num_heads"], + n_groups=linear["num_heads"], + head_dim=linear["head_dim"], + mamba_mask=recurrent, + target_full_attn_mask=attention, + spec_config=spec_config, + quant_config=quant_config, + ) + if params.mamba_ssm_cache_dtype != torch.float32: + logger.info( + f"Kimi K3: overriding mamba_ssm_cache_dtype {params.mamba_ssm_cache_dtype} " + "-> torch.float32 (KDA recurrent state must be fp32)" + ) + params.mamba_ssm_cache_dtype = torch.float32 + return params + + +__all__ = [ + "KDAReplayState", + "KDAReplayLayerCache", + "KimiK3HybridCacheManagerV2", + "allocate_kda_replay_fields", + "get_kda_replay_num_spec", + "get_kimi_cache_params", + "select_kda_replay_state", +] diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 59b4800bf8ea..1f71ba24ccdf 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -376,7 +376,11 @@ def _forward_impl( layer_cache = attn_metadata.kv_cache_manager.mamba_layer_cache(self.layer_idx) conv_pool = layer_cache.conv # [slots, 3D, W - 1] bf16 ssm_pool = layer_cache.temporal # [slots, H, V, K] fp32 - generation_state_indices = getattr(mamba_metadata, "generation_state_indices", None) + generation_state_indices = getattr( + mamba_metadata, + "generation_state_indices", + None, + ) if generation_state_indices is None: generation_state_indices = state_indices[num_prefills:] @@ -448,7 +452,7 @@ def _project_output(self, core: torch.Tensor) -> torch.Tensor: def _has_kda_replay_caches(self, layer_cache) -> bool: """True when the manager allocated the fused-verify replay caches.""" - return layer_cache is not None and layer_cache.has_kda_replay_caches + return bool(getattr(layer_cache, "has_kda_replay_caches", False)) def _sync_kda_replay_conv_window(self, layer_cache, slot_indices, conv_pool) -> None: """Seed replay conv caches from the live committed conv pool. diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py new file mode 100644 index 000000000000..66e0ebf5459b --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -0,0 +1,504 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Mamba2-owned speculative and replay cache state.""" + +from __future__ import annotations + +import os +import sys +from dataclasses import dataclass, field + +import torch + +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override + +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + IntermediateState, + MambaAcceptanceBatch, + MambaLayerCache, + MambaStateLayout, + ReplayStateUpdateMetadata, + SpeculativeMambaLayerCache, + _promote_intermediate_states, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( + MambaHybridCacheManagerV2, +) +from tensorrt_llm._utils import get_sm_version +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping + +MIN_REPLAY_HISTORY_SIZE = 16 + +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +@dataclass(frozen=True, kw_only=True) +class Mamba2ReplayLayerCache(SpeculativeMambaLayerCache): + """Mamba2 compact-replay tensors for one recurrent layer.""" + + prev_num_accepted_tokens: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + cache_buf_idx: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + old_x: torch.Tensor | None = None + old_B: torch.Tensor | None = None + old_dt: torch.Tensor | None = None + old_dA_cumsum: torch.Tensor | None = None + + +def _splitmix64(value: int) -> int: + """Return the SplitMix64 finalizer for ``value``.""" + value = (value + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 30 + value = (value * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 27 + value = (value * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 31 + return value & _MAMBA_SSM_UINT64_MASK + + +def compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: + """Return a reproducible, nonzero int64 Philox seed.""" + folded = ( + _MAMBA_SSM_SEED_BASE + + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK + ) & _MAMBA_SSM_UINT64_MASK + seed = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + return seed or 1 + + +def allocate_mamba_seed_buffer( + cache_size: int, rank_offset: int, device: torch.device +) -> torch.Tensor: + """Allocate the graph-stable per-slot Philox seed buffer.""" + seeds = [compute_deterministic_mamba_seed(0, slot, rank_offset) for slot in range(cache_size)] + return torch.tensor(seeds, dtype=torch.int64, device=device) + + +def mamba_seed_rank_offset(mapping: Mapping) -> int: + """Return a distinct deterministic seed offset for this parallel rank.""" + return mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + mapping.rank * 1_009 + + +def advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: torch.Tensor | None = None, +) -> None: + """Advance double-buffered replay bookkeeping after a verify step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to(replay_metadata.prev_num_accepted_tokens.dtype) + previous = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = ( + previous + replay_metadata.replay_step_width > replay_metadata.replay_history_size + ) + next_accepted = torch.where(wrote_checkpoint, accepted_tokens, previous + accepted_tokens) + cache_buffer = replay_metadata.cache_buf_idx[slots] + next_cache_buffer = torch.where(wrote_checkpoint, 1 - cache_buffer, cache_buffer) + if is_dummy_request is not None: + next_accepted = torch.where(is_dummy_request, previous, next_accepted) + next_cache_buffer = torch.where(is_dummy_request, cache_buffer, next_cache_buffer) + replay_metadata.prev_num_accepted_tokens[slots] = next_accepted + replay_metadata.cache_buf_idx[slots] = next_cache_buffer + + +def _reset_mamba_seed_buffer( + seed_buffer: torch.Tensor, + slots: torch.Tensor, + host_slots: list[int], + *, + counter: int, + rank_offset: int, +) -> None: + """Refresh selected slots in the graph-stable int64 Philox buffer.""" + seeds = [compute_deterministic_mamba_seed(counter, slot, rank_offset) for slot in host_slots] + seed_buffer[slots] = torch.tensor(seeds, dtype=torch.int64, device=seed_buffer.device) + + +class Mamba2State(IntermediateState): + """Intermediate promotion plus Mamba2 stochastic-rounding seeds.""" + + uses_replay = False + + @override + def __init__(self) -> None: + super().__init__() + self.rand_seed: torch.Tensor | None = None + self._seed_request_counter = 0 + self._seed_rank_offset = 0 + + @override + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + super().bind(context, ssm_states, conv_states) + self._seed_rank_offset = context.seed_rank_offset + if context.stochastic_rounding and ssm_states: + states = ssm_states[0] + self.rand_seed = allocate_mamba_seed_buffer( + states.shape[0], context.seed_rank_offset, states.device + ) + + @override + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + if self.rand_seed is None: + return + self._seed_request_counter += 1 + _reset_mamba_seed_buffer( + self.rand_seed, + slots, + host_slots, + counter=self._seed_request_counter, + rank_offset=self._seed_rank_offset, + ) + + @override + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + fields = super()._layer_cache_fields(layer_offset) + if self.rand_seed is not None: + fields["mamba_ssm_rand_seed"] = self.rand_seed + return fields + + @override + def shutdown(self) -> None: + super().shutdown() + self.rand_seed = None + + +class ReplayHistory: + """Own compact Mamba2 replay history, seeds, and bookkeeping.""" + + def __init__(self, tokens_per_gen_step: int) -> None: + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + self.rand_seed: torch.Tensor | None = None + self._seed_request_counter = 0 + self._seed_rank_offset = 0 + self.replay_step_width = tokens_per_gen_step + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step) + self.prev_num_accepted_tokens: torch.Tensor | None = None + self.cache_buf_idx: torch.Tensor | None = None + self.old_x: torch.Tensor | None = None + self.old_B: torch.Tensor | None = None + self.old_dt: torch.Tensor | None = None + self.old_dA_cumsum: torch.Tensor | None = None + + @property + def uses_replay(self) -> bool: + return True + + def validate(self, context: MambaStateLayout) -> None: + if context.spec_config is None: + raise ValueError("Mamba2 replay requires speculative decoding") + if context.n_groups_per_rank <= 0 and context.mamba_pp_layers: + raise ValueError("Mamba2 replay requires at least one state group per rank") + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + self._seed_rank_offset = context.seed_rank_offset + if not context.mamba_pp_layers: + return + + states = ssm_states + cache_size = states[0].shape[0] + device = states[0].device + tokens_per_step = context.spec_config.tokens_per_gen_step + common_intermediate_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_conv = torch.zeros( + common_intermediate_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) + + num_heads, head_dim, d_state = context.ssm_state_shape + common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.old_x = torch.zeros( + common_replay_shape + [self.replay_history_size, num_heads, head_dim], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_replay_shape + [self.replay_history_size, context.n_groups_per_rank, d_state], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_replay_shape + [num_heads, self.replay_history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros_like(self.old_dt) + + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + return { + "mamba_ssm_rand_seed": self.rand_seed, + "intermediate_conv_window": self.intermediate_conv[layer_offset], + "prev_num_accepted_tokens": self.prev_num_accepted_tokens, + "cache_buf_idx": self.cache_buf_idx, + "old_x": self.old_x[layer_offset], + "old_B": self.old_B[layer_offset], + "old_dt": self.old_dt[layer_offset], + "old_dA_cumsum": self.old_dA_cumsum[layer_offset], + } + + def make_layer_cache( + self, + layer_offset: int, + conv: torch.Tensor, + temporal: torch.Tensor, + ) -> MambaLayerCache: + return Mamba2ReplayLayerCache( + conv=conv, + temporal=temporal, + **self._layer_cache_fields(layer_offset), + ) + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + self.prev_num_accepted_tokens[slots] = 0 + self.cache_buf_idx[slots] = 0 + for buffer in (self.old_x, self.old_B, self.old_dt, self.old_dA_cumsum): + buffer[:, slots] = 0 + if self.rand_seed is not None: + self._seed_request_counter += 1 + _reset_mamba_seed_buffer( + self.rand_seed, + slots, + host_slots, + counter=self._seed_request_counter, + rank_offset=self._seed_rank_offset, + ) + + def update(self, batch: MambaAcceptanceBatch) -> None: + metadata = self.get_replay_metadata() + if metadata is None: + raise RuntimeError("Mamba2 replay buffers are not bound") + # Import through the facade to retain the established monkeypatch point. + from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager + + mamba_cache_manager._advance_replay_state( + metadata, + batch.destination_state_indices, + batch.num_accepted_tokens, + batch.is_dummy_request, + ) + _promote_intermediate_states(self._conv_states, self.intermediate_conv, batch) + + def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: + if self.prev_num_accepted_tokens is None or self.cache_buf_idx is None: + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size, + ) + + def shutdown(self) -> None: + self._ssm_states = () + self._conv_states = () + self.intermediate_conv = None + self.intermediate_indices = None + self.rand_seed = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + +def create_mamba2_state( + *, spec_config: object | None, use_replay: bool +) -> Mamba2State | ReplayHistory: + """Build Mamba2 intermediate state or compact replay history.""" + if use_replay: + if spec_config is None: + raise ValueError("Mamba replay requires speculative decoding") + return ReplayHistory(spec_config.tokens_per_gen_step) + return Mamba2State() + + +def select_mamba2_state( + *, + spec_config: object | None, + ssm_cache_dtype: torch.dtype, + stochastic_rounding: bool, +) -> Mamba2State | ReplayHistory: + """Apply Mamba2 replay gates and return its selected update policy.""" + sm = get_sm_version() + use_replay = spec_config is not None and sm >= 80 + if spec_config is None: + logger.info("Replay kernel requires speculative decoding; using non-replay path") + if spec_config is not None and ( + getattr(spec_config, "eagle_choices", None) is not None + or getattr(spec_config, "use_dynamic_tree", False) + ): + logger.info("Replay kernel incompatible with tree attention; using legacy MTP path") + use_replay = False + if stochastic_rounding and ssm_cache_dtype == torch.float16 and (sm < 100 or sm in (120, 121)): + logger.info( + "Replay kernel Philox requires 100 <= sm < 120; " + "using legacy MTP path for stochastic rounding support" + ) + use_replay = False + if os.environ.get("TRTLLM_USE_MAMBA_REPLAY", "1") == "0": + logger.info("Replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0") + use_replay = False + else: + logger.info("Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1") + return create_mamba2_state(spec_config=spec_config, use_replay=use_replay) + + +class NemotronHybridCacheManagerV2(MambaHybridCacheManagerV2): + """Nemotron/Mamba2 geometry, replay selection and seed ownership.""" + + @override + def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) -> None: + self._requested_replay = use_replay_state_update + kwargs.setdefault("conv_state_layout", "x_b_c") + super().__init__(*args, **kwargs) + + @override + def _initialize_model_state(self) -> None: + if self._requested_replay is not None: + self._speculative_state = create_mamba2_state( + spec_config=self.spec_config, use_replay=self._requested_replay + ) + else: + self._speculative_state = select_mamba2_state( + spec_config=self.spec_config, + ssm_cache_dtype=self.ssm_state_dtype, + stochastic_rounding=self._mamba_ssm_stochastic_rounding, + ) + + if isinstance(self._speculative_state, ReplayHistory): + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + _mamba_effective_tp_size, + ) + + if self.local_num_mamba_layers and self._global_n_groups % _mamba_effective_tp_size( + self._state_layout.mapping + ): + raise ValueError("Replay state groups must be divisible by the effective TP size") + self._speculative_state.validate(self._state_layout) + + @override + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + state = self._speculative_state + return state.get_replay_metadata() if isinstance(state, ReplayHistory) else None + + @override + def get_mamba_ssm_rand_seed(self) -> torch.Tensor | None: + return self._speculative_state.rand_seed + + @property + def intermediate_state_indices(self) -> torch.Tensor | None: + return self._speculative_state.intermediate_indices + + @property + def intermediate_ssm_states(self) -> torch.Tensor | None: + state = self._speculative_state + return None if isinstance(state, ReplayHistory) else state.intermediate_ssm + + @property + def intermediate_conv_states(self) -> torch.Tensor | None: + return self._speculative_state.intermediate_conv + + @override + def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: + if self.spec_config is None: + return super().mamba_layer_cache(layer_idx) + return self._speculative_state.make_layer_cache( + self.mamba_layer_offsets[layer_idx], + self.get_conv_states(layer_idx), + self.get_ssm_states(layer_idx), + ) + + @override + def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + self._speculative_state.reset_slots(slots, host_slots) + + @override + def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: + self._speculative_state.update(batch) + + @override + def _setup_model_state(self) -> None: + self._speculative_state.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) + + @override + def _shutdown_model_state(self) -> None: + self._speculative_state.shutdown() + + +def get_nemotron_cache_params(config, *, spec_config=None, quant_config=None): + """Derive Mamba2 state geometry from the Nemotron hybrid pattern.""" + from tensorrt_llm._torch.pyexecutor.config_utils import build_mamba_kv_cache_params + + pattern = config.hybrid_override_pattern + return build_mamba_kv_cache_params( + config, + state_size=config.ssm_state_size, + conv_kernel=config.conv_kernel, + num_heads=config.mamba_num_heads, + n_groups=config.n_groups, + head_dim=config.mamba_head_dim, + mamba_mask=[layer == "M" for layer in pattern], + target_full_attn_mask=[layer == "*" for layer in pattern], + spec_config=spec_config, + quant_config=quant_config, + ) + + +__all__ = [ + "Mamba2State", + "ReplayHistory", + "Mamba2ReplayLayerCache", + "NemotronHybridCacheManagerV2", + "create_mamba2_state", + "select_mamba2_state", + "get_nemotron_cache_params", + "allocate_mamba_seed_buffer", + "compute_deterministic_mamba_seed", + "advance_replay_state", + "mamba_seed_rank_offset", + "MIN_REPLAY_HISTORY_SIZE", +] diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 43ae9b449df6..2d54842215b4 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -1061,14 +1061,12 @@ def forward_core( layer_cache.intermediate_conv_window if is_target_verify else None ) intermediate_ssm_states = layer_cache.intermediate_ssm if is_target_verify else None - use_replay = is_target_verify and getattr( - attn_metadata.kv_cache_manager, "use_replay_state_update", False - ) replay_metadata = ( attn_metadata.kv_cache_manager.get_replay_state_update_metadata() - if use_replay + if is_target_verify else None ) + use_replay = replay_metadata is not None use_cached_replay_all_layer_commit = ( num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 5079a6b8443c..e095f3479dd7 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -397,7 +397,11 @@ def __init__( def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, num_contexts: int): self.replay_num_decodes = 0 - if not getattr(kv_cache_manager, 'use_replay_state_update', False): + get_replay_metadata = getattr(kv_cache_manager, + 'get_replay_state_update_metadata', None) + replay_metadata = get_replay_metadata( + ) if get_replay_metadata is not None else None + if replay_metadata is None: return num_decodes = batch_size - num_contexts self.replay_num_decodes = num_decodes @@ -415,17 +419,6 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, if num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: return - if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): - raise RuntimeError( - "Replay state update is enabled, but the KV cache manager " - "does not expose replay state update metadata.") - - replay_metadata = kv_cache_manager.get_replay_state_update_metadata() - if replay_metadata is None: - raise RuntimeError( - "Replay state update is enabled for a decode batch, but the " - "KV cache manager returned no replay state update metadata.") - prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens cache_buf_idx = replay_metadata.cache_buf_idx replay_step_width = replay_metadata.replay_step_width diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index d88c0ed23bc4..28259e10a42d 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -474,15 +474,11 @@ def forward( assert layer_cache is not None, \ "Speculative decoding requires Python MambaCacheManager" intermediate_conv_states = layer_cache.intermediate_conv_window - use_replay = getattr(attn_metadata.kv_cache_manager, - 'use_replay_state_update', False) + replay_metadata = (attn_metadata.kv_cache_manager. + get_replay_state_update_metadata()) + use_replay = replay_metadata is not None draft_token_num = spec_metadata.runtime_draft_len + 1 if use_replay: - replay_metadata = (attn_metadata.kv_cache_manager. - get_replay_state_update_metadata()) - assert replay_metadata is not None, ( - "Mamba replay state update is enabled but replay " - "metadata was not allocated.") replay_step_width = replay_metadata.replay_step_width assert draft_token_num == replay_step_width, ( "Mamba replay state update does not support dynamic " diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py new file mode 100644 index 000000000000..8affe0b5247d --- /dev/null +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -0,0 +1,556 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import dataclasses +import math +import sys +from dataclasses import replace +from typing import TYPE_CHECKING, Optional + +import torch + +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override + +from tensorrt_llm._torch.attention.backends.sparse.qsa.constants import ( + QSA_INDEX_K_CACHE_DTYPE, + QSA_POSITION_CACHE_DTYPE, + QSA_POSITION_COORDINATE_AXES, + QSA_SPARSE_KV_CACHE_DTYPES, +) +from tensorrt_llm._torch.attention.backends.sparse.qsa.params import QSASparseParams +from tensorrt_llm._torch.modules.fla.cache_manager import ( + GDNReplayState, + create_gdn_state, + validate_gdn_layout, +) +from tensorrt_llm._torch.pyexecutor.config_utils import resolve_hf_torch_dtype +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + MambaAcceptanceBatch, + MambaLayerCache, + ReplayStateUpdateMetadata, + _get_local_mamba_cache_layout, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import get_pp_layers +from tensorrt_llm._utils import binding_to_torch_dtype +from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp +from tensorrt_llm.logger import logger +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig +from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole + +if TYPE_CHECKING: + import transformers + +# Storage-layout constants; model geometry comes from QSASparseParams. +_INDEX_K_ELEMENT_BYTES = torch.empty((), dtype=QSA_INDEX_K_CACHE_DTYPE).element_size() +_POSITION_ELEMENT_BYTES = torch.empty((), dtype=QSA_POSITION_CACHE_DTYPE).element_size() + +# Per-token RoPE/mRoPE coordinates used when raw index keys are compressed. +# They are request state shared by all local QSA layers. +QSA_INDEX_POSITION = DataRole("qsa_index_position") + + +PLE_CONV_STATE = DataRole("ple_conv_state") +PLE_NGRAM_CONTEXT = DataRole("ple_ngram_context") + + +def get_qwen4_exp_ple_layer_mask(config: transformers.PretrainedConfig) -> list[bool]: + """Return the decoder-layer mask for the one-based PLE layer IDs.""" + ple_layer_ids = list(getattr(config, "ple_layer_ids", None) or []) + invalid_ids = [ + layer_id + for layer_id in ple_layer_ids + if not isinstance(layer_id, int) + or isinstance(layer_id, bool) + or not 1 <= layer_id <= config.num_hidden_layers + ] + if invalid_ids: + raise ValueError( + "ple_layer_ids must contain one-based decoder-layer IDs in " + f"[1, {config.num_hidden_layers}], got {invalid_ids}" + ) + if len(ple_layer_ids) != len(set(ple_layer_ids)): + raise ValueError("ple_layer_ids must not contain duplicate layer IDs") + if len(ple_layer_ids) > 1: + # Qwen4ExpModel._prepare_ple_state currently resolves one shared state + # tuple. Supporting multiple PLE layers requires separate per-layer + # metadata and pools; removing only this guard would silently reuse state. + raise ValueError( + "Qwen4-Exp currently supports at most one PLE decoder layer, " + f"got ple_layer_ids={ple_layer_ids}" + ) + ple_layer_id_set = set(ple_layer_ids) + return [(layer_id + 1) in ple_layer_id_set for layer_id in range(config.num_hidden_layers)] + + +@dataclasses.dataclass +class Qwen4ExpPLECacheParams: + """Shapes and dtypes for PLE recurrent-state pools.""" + + ple_layer_mask: list[bool] + num_ple_layers: int + short_conv_channels: int + short_conv_state_len: int + ngram_context_len: int + conv_state_dtype: torch.dtype + + +def extract_qwen4_exp_ple_cache_params( + config: transformers.PretrainedConfig, +) -> Qwen4ExpPLECacheParams: + """Derive PLE recurrent-state pool dimensions from the model config.""" + ple_layer_mask = get_qwen4_exp_ple_layer_mask(config) + hc_count = getattr(config, "hc_count", 1) or 1 + return Qwen4ExpPLECacheParams( + ple_layer_mask=ple_layer_mask, + num_ple_layers=sum(ple_layer_mask), + short_conv_channels=hc_count * config.hidden_size, + short_conv_state_len=(config.ple_conv_kernel_size - 1) * config.ngram_size, + ngram_context_len=config.ngram_size - 1, + conv_state_dtype=resolve_hf_torch_dtype(config) or torch.bfloat16, + ) + + +def _qwen4_exp_ple_state_bytes_per_rank( + model_config, + params, + mapping, + *, + spec_config, + use_separate_draft_kv_cache: bool, +) -> int: + """Bytes this rank needs for Qwen4-Exp PLE recurrent state; 0 otherwise.""" + from tensorrt_llm._torch.pyexecutor.config_utils import is_qwen4_exp + + # The estimator is also called with minimal model-config stubs that carry no + # `pretrained_config`; such a config is definitionally not Qwen4-Exp. + pretrained_config = getattr(model_config, "pretrained_config", None) + if pretrained_config is None or not is_qwen4_exp(pretrained_config): + return 0 + ple_params = extract_qwen4_exp_ple_cache_params(pretrained_config) + mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( + is_draft=False, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + combined_layer_mask = [ + is_mamba or is_attention + for is_mamba, is_attention in zip(mamba_layer_mask, full_attention_layer_mask) + ] + local_layer_indices, _ = get_pp_layers( + sum(combined_layer_mask), + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + local_ple_layers = sum( + layer_id < len(ple_params.ple_layer_mask) and ple_params.ple_layer_mask[layer_id] + for layer_id in local_layer_indices + ) + ple_bytes_per_layer = ( + ple_params.short_conv_channels + * ple_params.short_conv_state_len + * ple_params.conv_state_dtype.itemsize + + ple_params.ngram_context_len * torch.int64.itemsize + ) + return local_ple_layers * ple_bytes_per_layer + + +class Qwen4ExpHybridCacheManagerV2(MambaHybridCacheManagerV2): + """Hybrid GDN/KV manager with lifecycle-coupled sparse side buffers. + + QSA adds a per-layer index-K cache and request-wide position coordinates. + These buffers share the main K/V allocation, eviction, and prefix-reuse + lifecycle. The parent hybrid manager owns GDN/Mamba recurrent state. + + Index geometry is resolved from the same serving/checkpoint configuration + as the indexer. A local fallback could otherwise give the indexer and its + paged cache different shapes or page strides. + """ + + @override + def __init__( + self, + *args, + sparse_attention_config=None, + pretrained_config=None, + layer_mask=None, + use_replay_state_update: bool | None = None, + qwen4_exp_ple_cache_params: Qwen4ExpPLECacheParams | None = None, + **kwargs, + ) -> None: + self._requested_replay = use_replay_state_update + self._ple_params = qwen4_exp_ple_cache_params + self._pretrained_config = pretrained_config + self.qsa_index_dim = self.qsa_index_kv_heads = 0 + self.qsa_sparse_layer_ids = [] + self.qsa_position_layer_id = None + self._qsa_enabled = sparse_attention_config is not None + if self._qsa_enabled: + if layer_mask is None: + raise ValueError("QSA cache allocation requires the full-attention layer mask") + params = sparse_attention_config.to_sparse_params(pretrained_config=pretrained_config) + if not isinstance(params, QSASparseParams): + raise ValueError("Qwen4 QSA requires QSA sparse parameters") + self.qsa_index_dim, self.qsa_index_kv_heads = ( + params.index_head_dim, + params.index_kv_heads, + ) + self.qsa_sparse_layer_ids = [i for i, active in enumerate(layer_mask) if active] + self._is_ple_draft = kwargs.get("is_draft", False) + kwargs.setdefault("conv_state_layout", "q_k_v") + super().__init__(*args, layer_mask=layer_mask, **kwargs) + + @override + def _extra_buffers_per_layer( + self, + *, + tokens_per_block: int, + ) -> dict[int, list[BufferConfig]]: + """Register per-layer index K and lifecycle-aligned position pages.""" + if not self._qsa_enabled: + return {} + index_size = ( + self.qsa_index_kv_heads * self.qsa_index_dim * _INDEX_K_ELEMENT_BYTES * tokens_per_block + ) + local_sparse_layers = [ + layer_id for layer_id in self.qsa_sparse_layer_ids if layer_id in self.layer_offsets + ] + # Coordinates are request-wide, so all local indexers use one view. + # Register the position role once; duplicating it on every sparse layer + # wastes one three-axis int32 page per layer without adding state. + self.qsa_position_layer_id = next(iter(local_sparse_layers), None) + result = { + self.layer_offsets[layer_id]: [BufferConfig(role=Role.INDEX_KEY, size=index_size)] + for layer_id in local_sparse_layers + } + if self.qsa_position_layer_id is not None: + local_idx = self.layer_offsets[self.qsa_position_layer_id] + result[local_idx].append( + BufferConfig( + role=QSA_INDEX_POSITION, + size=( + QSA_POSITION_COORDINATE_AXES * _POSITION_ELEMENT_BYTES * tokens_per_block + ), + ) + ) + return result + + @override + def get_index_k_buffer( + self, + layer_idx: int, + kv_layout: str = "NHD", + ) -> Optional[torch.Tensor]: + """Return the index-K view using the indexer's resolved geometry.""" + return super().get_index_k_buffer( + layer_idx, + num_heads=self.qsa_index_kv_heads, + head_dim=self.qsa_index_dim, + dtype=QSA_INDEX_K_CACHE_DTYPE, + kv_layout=kv_layout, + ) + + @override + def get_buffers( + self, + layer_idx: int, + kv_layout: str = "NHD", + ) -> Optional[torch.Tensor]: + """Return this layer's adjacent K/V roles with the physical slot stride. + + QSA tables store lifecycle slot IDs. Preserving V2's coalesced stride is + therefore required even when several layers share one physical pool. + """ + if not self._qsa_enabled or self.dtype not in QSA_SPARSE_KV_CACHE_DTYPES: + # The parent owns packed data and scale-page layouts such as + # NVFP4. QSA returns to the regular backend for these formats. + return super().get_buffers(layer_idx, kv_layout) + if kv_layout not in ("NHD", "HND"): + raise ValueError(f"Unsupported kv_layout: {kv_layout}") + if layer_idx not in self.layer_offsets: + return None + if self.kv_cache_type == CacheTypeCpp.SELFKONLY: + raise NotImplementedError("QSA sparse attention requires both K and V cache buffers") + + layer_offset = self.layer_offsets[layer_idx] + torch_dtype = binding_to_torch_dtype(self.dtype) + head_dim = self.head_dim_per_layer[layer_offset] + num_heads = self.num_kv_heads_per_layer[layer_offset] + if kv_layout == "NHD": + page_shape = [self.tokens_per_block, num_heads, head_dim] + else: + page_shape = [num_heads, self.tokens_per_block, head_dim] + return self._get_slot_role_view( + layer_idx, (Role.KEY, Role.VALUE), dtype=torch_dtype, page_shape=page_shape + ) + + @torch.compiler.disable + def get_qsa_position_buffer(self) -> Optional[torch.Tensor]: + """Return per-token three-axis RoPE/mRoPE position coordinates.""" + if self.qsa_position_layer_id is None: + return None + return self._get_slot_role_view( + self.qsa_position_layer_id, + (QSA_INDEX_POSITION,), + dtype=QSA_POSITION_CACHE_DTYPE, + page_shape=[self.tokens_per_block, QSA_POSITION_COORDINATE_AXES], + )[:, 0] + + def get_qsa_attention_pool_layout(self) -> tuple[int, int]: + """Require one block-table mapping for all local sparse layers.""" + if self.qsa_position_layer_id is None: + raise RuntimeError("QSA cache manager has no local sparse layer") + mapping = self._get_attention_slot_mapping(self.qsa_position_layer_id) + for layer_id in self.qsa_sparse_layer_ids: + if layer_id not in self.layer_offsets: + continue + candidate = self._get_attention_slot_mapping(layer_id) + if candidate != mapping: + raise RuntimeError( + "QSA local layers do not share one attention page mapping: " + f"layer {layer_id} uses pool/scale {candidate}, expected {mapping}" + ) + return mapping + + def _init_qwen4_exp_ple_geometry( + self, + params: Optional["Qwen4ExpPLECacheParams"], + total_layers: int, + default_state_dtype: torch.dtype, + ) -> None: + """Validate and record the Qwen4-Exp PLE pool geometry. + + Disabled PLE contributes no recurrent buffers. + """ + self._ple_layer_ids: list[int] = [] + self._ple_conv_state_shape: list[int] = [] + self._ple_ngram_context_shape: list[int] = [] + self._ple_conv_state_dtype = default_state_dtype + self._ple_conv_states: dict[int, torch.Tensor] = {} + self._ple_ngram_contexts: dict[int, torch.Tensor] = {} + if params is None: + return + if len(params.ple_layer_mask) != total_layers: + raise ValueError( + "PLE layer mask length must match the hybrid layer mask: " + f"got {len(params.ple_layer_mask)}, expected {total_layers}" + ) + self._ple_layer_ids = [ + layer_id for layer_id, active in enumerate(params.ple_layer_mask) if active + ] + if len(self._ple_layer_ids) != params.num_ple_layers: + raise ValueError( + "PLE layer mask count does not match num_ple_layers: " + f"got {len(self._ple_layer_ids)}, expected " + f"{params.num_ple_layers}" + ) + if any(not self._mamba_layer_mask[layer_id] for layer_id in self._ple_layer_ids): + raise ValueError("PLE lifecycle state must belong to Mamba layers") + if ( + params.short_conv_channels <= 0 + or params.short_conv_state_len <= 0 + or params.ngram_context_len <= 0 + ): + raise ValueError("PLE recurrent-state dimensions must be positive") + self._ple_conv_state_shape = [ + params.short_conv_channels, + params.short_conv_state_len, + ] + self._ple_ngram_context_shape = [params.ngram_context_len] + self._ple_conv_state_dtype = params.conv_state_dtype + + def _setup_ple_states(self, num_state_slots: int) -> None: + """Bind each local PLE layer to V2-managed recurrent-state views.""" + if not self._ple_layer_ids: + return + for layer_id in self._ple_layer_ids: + local_layer_idx = self.layer_offsets.get(layer_id) + if local_layer_idx is None: + continue + conv_state = self._get_state_buffer( + local_layer_idx, + PLE_CONV_STATE, + self._ple_conv_state_dtype, + self._ple_conv_state_shape, + ) + ngram_context = self._get_state_buffer( + local_layer_idx, + PLE_NGRAM_CONTEXT, + dtype=torch.long, + state_shape=self._ple_ngram_context_shape, + ) + if conv_state.shape[0] != num_state_slots or ngram_context.shape[0] != num_state_slots: + raise RuntimeError( + "PLE and GDN lifecycle buffers must have the same number " + f"of slots: layer={layer_id}, GDN={num_state_slots}, " + f"conv={conv_state.shape[0]}, ngram={ngram_context.shape[0]}" + ) + self._ple_conv_states[layer_id] = conv_state + self._ple_ngram_contexts[layer_id] = ngram_context + logger.info( + "PLE state views bound to V2 lifecycle buffers for local layers " + f"{sorted(self._ple_conv_states)}: conv " + f"[{num_state_slots}, {', '.join(map(str, self._ple_conv_state_shape))}] " + f"({self._ple_conv_state_dtype}), n-gram context " + f"[{num_state_slots}, {self._ple_ngram_context_shape[0]}] (int64)" + ) + + def ple_layer_cache(self, layer_idx: int) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + """Return persistent ``(conv_state, ngram_context)`` pools for a layer.""" + conv = self._ple_conv_states.get(layer_idx) + ngram = self._ple_ngram_contexts.get(layer_idx) + if conv is None or ngram is None: + return None + return conv, ngram + + def _ple_buffer_configs(self) -> list[BufferConfig]: + if not self._ple_layer_ids: + return [] + return [ + BufferConfig( + role=PLE_NGRAM_CONTEXT, + size=(math.prod(self._ple_ngram_context_shape) * torch.int64.itemsize), + ), + BufferConfig( + role=PLE_CONV_STATE, + size=(math.prod(self._ple_conv_state_shape) * self._ple_conv_state_dtype.itemsize), + ), + ] + + @override + def _initialize_model_state(self) -> None: + self._speculative_state = create_gdn_state(self, self._requested_replay) + + validate_gdn_layout(self, self._speculative_state) + params = self._ple_params + if self._is_ple_draft: + params = None + elif params is None and getattr(self._pretrained_config, "ple_layer_ids", None): + params = get_qwen4_exp_ple_cache_params( + self._pretrained_config, + total_layers=len(self._mamba_layer_mask), + is_draft=False, + ) + self._init_qwen4_exp_ple_geometry( + params, len(self._mamba_layer_mask), self.conv_state_dtype + ) + + @override + def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConfig]: + buffers = super()._get_recurrent_buffer_configs(global_layer_id) + if global_layer_id in self._ple_layer_ids: + buffers.extend(self._ple_buffer_configs()) + return buffers + + @override + def _setup_model_state(self) -> None: + self._speculative_state.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) + if self.local_num_mamba_layers: + self._setup_ple_states(self._state_layout.slot_capacity) + + @override + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + state = self._speculative_state + return state.get_replay_metadata() if isinstance(state, GDNReplayState) else None + + @classmethod + @override + def _extra_state_bytes_per_rank(cls, model_config, mapping, **kwargs) -> int: + if kwargs.get("is_draft", False): + return 0 + params, _, _ = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=kwargs.get("spec_config"), + is_draft=False, + use_separate_draft_kv_cache=kwargs.get("use_separate_draft_kv_cache", False), + ) + return _qwen4_exp_ple_state_bytes_per_rank( + model_config, + params, + mapping, + spec_config=kwargs.get("spec_config"), + use_separate_draft_kv_cache=kwargs.get("use_separate_draft_kv_cache", False), + ) + + @override + def _shutdown_model_state(self) -> None: + self._ple_conv_states.clear() + self._ple_ngram_contexts.clear() + self._speculative_state.shutdown() + + @property + def intermediate_state_indices(self) -> torch.Tensor | None: + return self._speculative_state.intermediate_indices + + @property + def intermediate_ssm_states(self) -> torch.Tensor | None: + state = self._speculative_state + return None if isinstance(state, GDNReplayState) else state.intermediate_ssm + + @property + def intermediate_conv_states(self) -> torch.Tensor | None: + return self._speculative_state.intermediate_conv + + @property + def use_gdn_cached_replay_all_layer_commit(self) -> bool: + state = self._speculative_state + return isinstance(state, GDNReplayState) and state.has_bound_states + + @override + def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: + if self.spec_config is None: + return super().mamba_layer_cache(layer_idx) + return self._speculative_state.make_layer_cache( + self.mamba_layer_offsets[layer_idx], + self.get_conv_states(layer_idx), + self.get_ssm_states(layer_idx), + ) + + @override + def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + self._speculative_state.reset_slots(slots, host_slots) + + @override + def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: + self._speculative_state.update(batch) + + +def get_qwen4_exp_ple_cache_params(config, *, total_layers: int, is_draft: bool): + """Align target-only PLE state with a target/draft cache layout.""" + if is_draft: + return None + + params = extract_qwen4_exp_ple_cache_params(config) + num_target_layers = len(params.ple_layer_mask) + if num_target_layers > total_layers: + raise ValueError( + "PLE layer mask cannot exceed the hybrid cache layout: " + f"got {num_target_layers}, expected at most {total_layers}" + ) + if num_target_layers == total_layers: + return params + + # Unified one-model caches append attention-only MTP layers. + return replace( + params, + ple_layer_mask=params.ple_layer_mask + [False] * (total_layers - num_target_layers), + ) + + +__all__ = [ + "Qwen4ExpHybridCacheManagerV2", + "Qwen4ExpPLECacheParams", + "extract_qwen4_exp_ple_cache_params", + "get_qwen4_exp_ple_layer_mask", + "get_qwen4_exp_ple_cache_params", + "PLE_CONV_STATE", + "PLE_NGRAM_CONTEXT", + "QSA_INDEX_POSITION", +] diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index faa8b064caf6..584a4f047d18 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -21,9 +21,9 @@ import tensorrt_llm import tensorrt_llm.bindings.executor as trtllm -from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, - is_sm_100f, prefer_pinned, - str_dtype_to_binding, torch_dtype_to_str) +from tensorrt_llm._utils import (confidential_compute_enabled, is_sm_100f, + prefer_pinned, str_dtype_to_binding, + torch_dtype_to_str) from tensorrt_llm.inputs.multimodal import MultimodalParams # isort: off @@ -51,9 +51,7 @@ from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, get_num_spec_layers, get_spec_decoder, should_use_separate_draft_kv_cache) -from ..utils import is_gdn_replay_enabled from .config_utils import (MambaKVCacheParams, extract_mamba_kv_cache_params, - extract_qwen4_exp_ple_cache_params, get_layer_attention_window, is_gemma4_hybrid, is_hybrid_linear, is_kimi_linear, is_mla, is_nemotron_hybrid, is_qwen3_hybrid, is_qwen4_exp, @@ -67,6 +65,7 @@ MambaHybridCacheManagerV2, MixedMambaHybridCacheManager, use_py_mamba_cache_manager) +from .kv_cache.mamba_cache_manager.legacy import warn_legacy_mamba_cache_manager from .llm_request import ExecutorResponse, LlmRequestState from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor @@ -121,6 +120,13 @@ def _resolve_disagg_transceiver_route( return backend, runtime +def _select_legacy_mamba_manager(manager_cls): + implementation = ("python" if issubclass( + manager_cls, MixedMambaHybridCacheManager) else "cpp") + warn_legacy_mamba_cache_manager(implementation, manager_cls.__name__) + return manager_cls + + def get_kv_cache_manager_cls( model_config: ModelConfig, kv_cache_config: KvCacheConfig, @@ -193,10 +199,10 @@ def get_kv_cache_manager_cls( logger.info( "Using CppMambaHybridCacheManager for Kimi K3 hybrid " "model (block reuse enabled)") - return CppMambaHybridCacheManager + return _select_legacy_mamba_manager(CppMambaHybridCacheManager) logger.info( "Using MixedMambaHybridCacheManager for Kimi K3 hybrid model") - return MixedMambaHybridCacheManager + return _select_legacy_mamba_manager(MixedMambaHybridCacheManager) # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. @@ -232,12 +238,14 @@ def get_kv_cache_manager_cls( "and transceiver_runtime='PYTHON' requires " "use_kv_cache_manager_v2=True.") if kv_cache_config.enable_block_reuse: - return CppMambaHybridCacheManager + return _select_legacy_mamba_manager( + CppMambaHybridCacheManager) if runtime == "PYTHON" and backend == "NIXL": logger.info("Python transceiver detected; using " "MixedMambaHybridCacheManager for hybrid model") - return MixedMambaHybridCacheManager - return CppMambaHybridCacheManager + return _select_legacy_mamba_manager( + MixedMambaHybridCacheManager) + return _select_legacy_mamba_manager(CppMambaHybridCacheManager) if use_py_mamba_cache_manager() and not is_disagg: if use_v2: @@ -252,7 +260,7 @@ def get_kv_cache_manager_cls( "TRTLLM_USE_PY_MAMBA to use the configured cache manager.") logger.info( "Using MixedMambaHybridCacheManager for hybrid mamba model") - return MixedMambaHybridCacheManager + return _select_legacy_mamba_manager(MixedMambaHybridCacheManager) env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None) if env_override is not None: env_override = env_override.upper() @@ -272,7 +280,8 @@ def get_kv_cache_manager_cls( "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED " "overrides the default Mamba cache manager to " "MixedMambaHybridCacheManager.") - return MixedMambaHybridCacheManager + return _select_legacy_mamba_manager( + MixedMambaHybridCacheManager) if env_override == 'CPP': if use_v2: raise ValueError( @@ -282,14 +291,14 @@ def get_kv_cache_manager_cls( "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP " "overrides the default Mamba cache manager to " "CppMambaHybridCacheManager.") - return CppMambaHybridCacheManager + return _select_legacy_mamba_manager(CppMambaHybridCacheManager) logger.warning( f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " "Expected 'CPP' or 'MIXED'. Using the configured " "KV cache manager default.") if not use_v2: - return CppMambaHybridCacheManager + return _select_legacy_mamba_manager(CppMambaHybridCacheManager) if (kv_cache_config.enable_block_reuse and kv_cache_config.enable_kv_pool_rebalance): @@ -300,7 +309,19 @@ def get_kv_cache_manager_cls( if sparse_attn_algorithm == "qsa": return get_sparse_attn_kv_cache_manager( sparse_attn_config, use_kv_cache_manager_v2=True) - return MambaHybridCacheManagerV2 + if is_kimi_linear(config): + from ..modules.kimi_kda.cache_manager import \ + KimiK3HybridCacheManagerV2 + return KimiK3HybridCacheManagerV2 + if is_qwen4_exp(config): + from ..modules.qwen4_exp.cache_manager import \ + Qwen4ExpHybridCacheManagerV2 + return Qwen4ExpHybridCacheManagerV2 + if is_qwen3_hybrid(config): + from ..modules.fla.cache_manager import Qwen35HybridCacheManagerV2 + return Qwen35HybridCacheManagerV2 + from ..modules.mamba.cache_manager import NemotronHybridCacheManagerV2 + return NemotronHybridCacheManagerV2 elif sparse_attn_config is not None: if sparse_attn_algorithm == "qsa": # The QSA manager extends the hybrid manager because it must retain @@ -2325,29 +2346,6 @@ def _mamba_conv_layout_kwargs(kv_cache_manager_cls: type, return {"model_type": model_type} -def _get_qwen4_exp_ple_cache_params(config, *, total_layers: int, - is_draft: bool): - """Align target-only PLE state with a target/draft cache layout.""" - if is_draft: - return None - - params = extract_qwen4_exp_ple_cache_params(config) - num_target_layers = len(params.ple_layer_mask) - if num_target_layers > total_layers: - raise ValueError( - "PLE layer mask cannot exceed the hybrid cache layout: " - f"got {num_target_layers}, expected at most {total_layers}") - if num_target_layers == total_layers: - return params - - # Unified one-model caches append attention-only MTP layers. - return dataclasses.replace( - params, - ple_layer_mask=params.ple_layer_mask + [False] * - (total_layers - num_target_layers), - ) - - def _create_kv_cache_manager( model_engine: Optional[PyTorchModelEngine], kv_cache_manager_cls, @@ -2552,22 +2550,15 @@ def _create_kv_cache_manager( # manager's own internal gate (`tp_size = 1 if enable_attention_dp # else tp_size`, then num_heads / n_groups / conv_dim divide by # it), so the params pass through unscaled. - # KDA fused multi-token verify (trtllm::kda_mtp_decode): when the - # kernel can run here, allocate the per-slot replay caches instead - # of the legacy per-step intermediate verification buffers. The - # kernel replays accepted drafts from these caches and commits - # states in place, replacing the intermediate-buffer + promotion - # flow for KDA layers. + from ..modules.kimi_kda.cache_manager import select_kda_replay_state + kimi_extra_kwargs = {} - kda_replay_manager_types = (MixedMambaHybridCacheManager, - MambaHybridCacheManagerV2) - if (spec_config is not None - and issubclass(kv_cache_manager_cls, kda_replay_manager_types)): - from ..modules.kimi_kda._kda_kernels import \ - is_kda_mtp_verify_available - if is_kda_mtp_verify_available(): - kimi_extra_kwargs["kda_replay_num_spec"] = ( - spec_config.tokens_per_gen_step - 1) + if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + replay = select_kda_replay_state(spec_config=spec_config, + manager_cls=kv_cache_manager_cls) + if replay is not None: + kimi_extra_kwargs[ + "kda_replay_num_spec"] = replay.num_speculative_tokens # KDA's conv state is a [Q | K | V] concatenation whose three sections # have identical width, i.e. the qwen3_next section layout. kimi_extra_kwargs.update( @@ -2655,54 +2646,17 @@ def _create_kv_cache_manager( num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 else mamba_params.num_mamba_layers) - # Replay state update kernel for MTP: default on for sm >= 80; gates - # below disable it for incompatible feature combinations. Cpp cache - # manager doesn't expose use_replay_state_update, so the wrapper - # property's getattr default keeps replay off there automatically. - sm = get_sm_version() stochastic_rounding = getattr( quant_config, 'mamba_ssm_stochastic_rounding', False) if quant_config is not None else False + from ..modules.mamba.cache_manager import select_mamba2_state - use_replay = spec_config is not None and sm >= 80 - if spec_config is None: - logger.info( - "Replay kernel requires speculative decoding; using non-replay path" - ) - - # Block reuse (prefix caching): replay leaves SSM state at a - # checkpoint after speculation. The next decode step replays forward - # to correct it. If block reuse feeds that stale state into a new - # prefill, the correction never happens. - # Currently we only save and reuse context tokens so this does not affect. - - # Tree attention: replay assumes linear token sequence. - if (spec_config is not None - and (getattr(spec_config, 'eagle_choices', None) is not None - or getattr(spec_config, 'use_dynamic_tree', False))): - logger.info("Replay kernel incompatible with tree attention; " - "using legacy MTP path") - use_replay = False - - # Replay Philox uses PTX cvt.rs.f16x2.f32 which needs 100 <= sm < 120. - # Flashinfer has a SW fallback at any SM. - if (stochastic_rounding - and mamba_params.mamba_ssm_cache_dtype == torch.float16 - and (sm < 100 or sm in (120, 121))): - logger.info("Replay kernel Philox requires 100 <= sm < 120; " - "using legacy MTP path for stochastic rounding support") - use_replay = False - - # Use replay algorithm for mamba (default is on). - enforce_disable_replay = os.environ.get('TRTLLM_USE_MAMBA_REPLAY', - '1') == '0' - if enforce_disable_replay: - logger.info( - "Replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0") - use_replay = False - else: - logger.info( - "Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1") + use_replay = None + if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + use_replay = select_mamba2_state( + spec_config=spec_config, + ssm_cache_dtype=mamba_params.mamba_ssm_cache_dtype, + stochastic_rounding=stochastic_rounding).uses_replay # Stochastic-rounding seeds must live on the cache manager (not be # re-created with torch.randint per forward) whenever SR can fire @@ -2770,76 +2724,23 @@ def _create_kv_cache_manager( )) num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 else mamba_params.num_mamba_layers) - # Replay state update for GDN MTP: mirrors the nemotron_hybrid gating - # above, minus the Mamba2-specific stochastic-rounding/Philox gate. - # The GDN replay kernel does a plain cast on checkpoint commit, so - # quantized SSM cache dtypes stay on the legacy path. - sm = get_sm_version() - use_replay = spec_config is not None and sm >= 80 - if spec_config is None: - logger.info( - "GDN replay kernel requires speculative decoding; using " - "non-replay path") - elif spec_config.tokens_per_gen_step > 8: - logger.info("GDN cached replay supports at most 8 tokens per " - "generation step; using non-replay path") - use_replay = False - - # Tree attention: replay assumes a linear token sequence. - if (spec_config is not None - and (getattr(spec_config, 'eagle_choices', None) is not None - or getattr(spec_config, 'use_dynamic_tree', False))): - logger.info("GDN replay kernel incompatible with tree attention; " - "using legacy MTP path") - use_replay = False - - if mamba_params.mamba_ssm_cache_dtype not in (torch.float32, - torch.bfloat16, - torch.float16): - logger.info( - "GDN replay kernel does not support quantized SSM cache " - f"dtype {mamba_params.mamba_ssm_cache_dtype}; using legacy " - "MTP path") - use_replay = False - - # Replay is enabled by default for eligible GDN MTP workloads. - if not is_gdn_replay_enabled(): - use_replay = False - - # GDN replay supports the contiguous C++ V1 state pool and the indirect - # per-layer state views exposed by V2. Mixed/Python does not expose an - # all-layer commit, so keep that manager but use non-replay MTP. - replay_manager_types = (CppMambaHybridCacheManager, - MambaHybridCacheManagerV2) - if use_replay and not issubclass(kv_cache_manager_cls, - replay_manager_types): - logger.info("GDN replay requires C++ V1 or V2 Mamba cache manager; " - f"{kv_cache_manager_cls.__name__} was selected, so the " - "non-replay MTP path will be used") - use_replay = False - logger.info("GDN replay state update: " + - ("ENABLED" if use_replay else "DISABLED")) + from ..modules.fla.cache_manager import select_gdn_replay_state + + use_replay = None + if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): + use_replay = select_gdn_replay_state( + spec_config=spec_config, + ssm_cache_dtype=mamba_params.mamba_ssm_cache_dtype, + manager_cls=kv_cache_manager_cls) is not None mamba_manager_extra_kwargs = dict(manager_extra_kwargs) mamba_manager_extra_kwargs.update( _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next")) - if getattr(sparse_attention_config, "algorithm", None) == "qsa": - # Resolve the side-cache shape from the same checkpoint geometry - # used to construct the QSA index projection. - mamba_manager_extra_kwargs.update( - sparse_attention_config=sparse_attention_config, - pretrained_config=config, - ) if is_qwen4_exp(config) and issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - ple_cache_params = _get_qwen4_exp_ple_cache_params( - config, - total_layers=len(mamba_layer_mask), - is_draft=is_draft, - ) - if ple_cache_params is not None: - mamba_manager_extra_kwargs[ - "qwen4_exp_ple_cache_params"] = ple_cache_params + mamba_manager_extra_kwargs.update( + pretrained_config=config, + sparse_attention_config=sparse_attention_config) kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 03ede19b9e5c..93230fb09142 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -361,61 +361,6 @@ def get_qwen3_hybrid_num_attention_layers(config): return sum(layer_mask) -def get_qwen4_exp_ple_layer_mask( - config: transformers.PretrainedConfig) -> list[bool]: - """Return the decoder-layer mask for the one-based PLE layer IDs.""" - ple_layer_ids = list(getattr(config, "ple_layer_ids", None) or []) - invalid_ids = [ - layer_id for layer_id in ple_layer_ids - if not isinstance(layer_id, int) or isinstance(layer_id, bool) - or not 1 <= layer_id <= config.num_hidden_layers - ] - if invalid_ids: - raise ValueError( - "ple_layer_ids must contain one-based decoder-layer IDs in " - f"[1, {config.num_hidden_layers}], got {invalid_ids}") - if len(ple_layer_ids) != len(set(ple_layer_ids)): - raise ValueError("ple_layer_ids must not contain duplicate layer IDs") - if len(ple_layer_ids) > 1: - # Qwen4ExpModel._prepare_ple_state currently resolves one shared state - # tuple. Supporting multiple PLE layers requires separate per-layer - # metadata and pools; removing only this guard would silently reuse state. - raise ValueError( - "Qwen4-Exp currently supports at most one PLE decoder layer, " - f"got ple_layer_ids={ple_layer_ids}") - ple_layer_id_set = set(ple_layer_ids) - return [(layer_id + 1) in ple_layer_id_set - for layer_id in range(config.num_hidden_layers)] - - -@dataclasses.dataclass -class Qwen4ExpPLECacheParams: - """Shapes and dtypes for PLE recurrent-state pools.""" - - ple_layer_mask: list[bool] - num_ple_layers: int - short_conv_channels: int - short_conv_state_len: int - ngram_context_len: int - conv_state_dtype: torch.dtype - - -def extract_qwen4_exp_ple_cache_params( - config: transformers.PretrainedConfig) -> Qwen4ExpPLECacheParams: - """Derive PLE recurrent-state pool dimensions from the model config.""" - ple_layer_mask = get_qwen4_exp_ple_layer_mask(config) - hc_count = getattr(config, "hc_count", 1) or 1 - return Qwen4ExpPLECacheParams( - ple_layer_mask=ple_layer_mask, - num_ple_layers=sum(ple_layer_mask), - short_conv_channels=hc_count * config.hidden_size, - short_conv_state_len=(config.ple_conv_kernel_size - 1) * - config.ngram_size, - ngram_context_len=config.ngram_size - 1, - conv_state_dtype=resolve_hf_torch_dtype(config) or torch.bfloat16, - ) - - @dataclasses.dataclass class MambaKVCacheParams: """Normalized mamba-related inputs for kv_cache_manager_cls. @@ -490,62 +435,36 @@ def get_states_bytes_per_layer(self, mapping) -> int: return state_bytes_per_layer -def extract_mamba_kv_cache_params( - config, - spec_config=None, - quant_config=None, -) -> MambaKVCacheParams: - """Build the mamba-related inputs for kv_cache_manager_cls. - - Supports Nemotron-hybrid, Qwen3-hybrid (Qwen3-Next + Qwen3.5), - Qwen4-Exp, and Kimi K3 (kimi_linear). - - Args: - config: HuggingFace model config of a hybrid Mamba model. - spec_config: Optional speculative-decoding config used to describe - appended attention-only MTP/draft layers separately from target - layers. - quant_config: Optional, used only to surface `mamba_ssm_cache_dtype`. - - Returns: - MambaKVCacheParams with normalized field names. - """ +def extract_mamba_kv_cache_params(config, + spec_config=None, + quant_config=None) -> MambaKVCacheParams: + """Dispatch geometry extraction to the model's cache implementation.""" if is_nemotron_hybrid(config): - state_size = config.ssm_state_size - conv_kernel = config.conv_kernel - num_heads = config.mamba_num_heads - n_groups = config.n_groups - head_dim = config.mamba_head_dim - pattern = config.hybrid_override_pattern - target_full_attn_mask = [layer_type == "*" for layer_type in pattern] - mamba_mask = [layer_type == "M" for layer_type in pattern] + from ..modules.mamba.cache_manager import \ + get_nemotron_cache_params as extract elif is_qwen3_hybrid(config) or is_qwen4_exp(config): - state_size = config.linear_key_head_dim - conv_kernel = config.linear_conv_kernel_dim - num_heads = config.linear_num_value_heads - n_groups = config.linear_num_key_heads - head_dim = config.linear_value_head_dim - target_full_attn_mask, mamba_mask = get_qwen3_hybrid_layer_masks(config) + from ..modules.fla.cache_manager import get_gdn_cache_params as extract elif is_kimi_linear(config): - # Kimi K3 KDA (Kimi Delta Attention) state, mapped onto the Mamba - # cache-manager parametrization (see PythonMambaCacheManager): - # conv_dim = head_dim*num_heads + 2*n_groups*state_size - # = 3 * num_heads * head_dim -> [q | k | v] short-conv - # ssm state shape = [num_heads, head_dim, state_size] - # = [H, V, K] fp32 delta-rule recurrent state. - # The pool stores the W - 1 raw inputs needed by the production - # causal-convolution kernels, where W is short_conv_kernel_size. - lin = unwrap_kimi_text_config(config).linear_attn_config - state_size = lin["head_dim"] - conv_kernel = lin["short_conv_kernel_size"] - num_heads = lin["num_heads"] - n_groups = lin["num_heads"] - head_dim = lin["head_dim"] - target_full_attn_mask, mamba_mask = get_kimi_linear_layer_masks(config) + from ..modules.kimi_kda.cache_manager import \ + get_kimi_cache_params as extract else: raise ValueError( f"{type(config).__name__} is not a supported hybrid Mamba config") - + return extract(config, spec_config=spec_config, quant_config=quant_config) + + +def build_mamba_kv_cache_params(config, + *, + state_size, + conv_kernel, + num_heads, + n_groups, + head_dim, + mamba_mask, + target_full_attn_mask, + spec_config=None, + quant_config=None) -> MambaKVCacheParams: + """Normalize dtype and draft layout after model-owned geometry extraction.""" num_draft_layers = 0 if spec_config is not None: # Imported lazily to avoid a circular dependency between @@ -561,15 +480,6 @@ def extract_mamba_kv_cache_params( mamba_ssm_cache_dtype = (resolve_ssm_cache_dtype(config) or resolve_hf_torch_dtype(config) or torch.bfloat16) - if is_kimi_linear(config) and mamba_ssm_cache_dtype != torch.float32: - # The KDA delta-rule recurrent state must be kept in fp32 for - # numerical parity with the HF reference (fla chunk/fused_recurrent - # KDA kernels carry the state in fp32). - logger.info( - f"Kimi K3: overriding mamba_ssm_cache_dtype " - f"{mamba_ssm_cache_dtype} -> torch.float32 (KDA recurrent state " - "must be fp32)") - mamba_ssm_cache_dtype = torch.float32 return MambaKVCacheParams( state_size=state_size, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 6ec6e3898db6..a8fdb4554bb2 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2302,6 +2302,75 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. ) ) + def _get_slot_role_view( + self, + layer_idx: int, + roles: tuple[DataRole, ...], + *, + dtype: torch.dtype, + page_shape: list[int], + ) -> torch.Tensor: + """View adjacent equal-sized roles as [slots, roles, *page_shape]. + + Unlike a packed page view, this retains the physical lifecycle-slot + stride when other layers or roles share the same pool. The caller + supplies the registered role shape and dtype; expanded page mappings + cannot be represented by this view. + """ + if not roles: + raise ValueError("A slot view requires at least one role") + layer_offset = self.layer_offsets[layer_idx] + first_role = roles[0] + address = self.impl.get_mem_pool_base_address( + layer_offset, first_role, PageIndexMode.SHARED + ) + page_stride = self.impl.get_page_stride(layer_offset, first_role) + expected_stride = math.prod(page_shape) * dtype.itemsize + if page_stride != expected_stride: + raise RuntimeError( + f"Slot-view page stride mismatch: expected {expected_stride}, got {page_stride}" + ) + for offset, role in enumerate(roles[1:], start=1): + role_address = self.impl.get_mem_pool_base_address( + layer_offset, role, PageIndexMode.SHARED + ) + if ( + role_address != address + offset * page_stride + or self.impl.get_page_stride(layer_offset, role) != page_stride + ): + raise RuntimeError(f"Slot-view roles are not adjacent equal-sized pages: {roles}") + + converter = self.impl.get_page_index_converter(layer_offset, first_role) + scale = int(converter.scale) + if scale < len(roles): + raise RuntimeError( + f"Invalid slot-view page-index scale: {scale} for {len(roles)} roles" + ) + if int(converter.expansion) != 1: + raise RuntimeError("Slot views do not support expanded page indices") + page_upper = self.impl.get_page_index_upper_bound(layer_offset, first_role) + num_pages_with_offset = page_upper + int(converter.layer_offset) + if num_pages_with_offset % scale != 0: + raise RuntimeError("Slot-view page mapping is inconsistent") + view = convert_to_torch_tensor( + TensorWrapper(address, dtype, [num_pages_with_offset // scale, scale, *page_shape]) + ) + return view[:, : len(roles)] + + def _get_attention_slot_mapping(self, layer_idx: int) -> tuple[int, int]: + """Return an attention layer's block-table pool and unexpanded slot scale.""" + local_idx = self.layer_offsets[layer_idx] + pool_id = self.layer_to_pool_mapping_dict[local_idx] + if pool_id >= self.num_attention_op_pools: + raise RuntimeError(f"KV pool {pool_id} is not represented in the attention block table") + converter = self.impl.get_page_index_converter(local_idx, Role.KEY) + scale = int(converter.scale) + if scale <= 0: + raise RuntimeError(f"Invalid attention page-index scale: {scale}") + if int(converter.expansion) != 1: + raise RuntimeError("Attention slot mapping does not support expanded page indices") + return pool_id, scale + def get_index_k_buffer( self, layer_idx: int, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py deleted file mode 100644 index e242959e5117..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ /dev/null @@ -1,4649 +0,0 @@ -# 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"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import math -import os -from abc import ABC, abstractmethod -from dataclasses import dataclass, replace -from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, - Optional, Protocol, Sequence, Tuple, Union, cast) - -import torch -import triton -import triton.language as tl - -if TYPE_CHECKING: - from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata - from tensorrt_llm._torch.pyexecutor.config_utils import Qwen4ExpPLECacheParams - from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig - from tensorrt_llm.sampling_params import SamplingParams - -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - BlockReusePolicy, KVCacheManagerV2, Role) -from tensorrt_llm._torch.pyexecutor.kv_cache_stats import \ - KVCacheV2IterationStatsReport -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 (TensorWrapper, convert_to_torch_tensor, - nvtx_range, prefer_pinned, - torch_dtype_to_binding) -from tensorrt_llm.bindings.internal.batch_manager import ( - LinearAttentionMetadata, LinearCacheType) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, - BatchDesc, BufferConfig, - DataRole, KVCacheDesc) -from tensorrt_llm.runtime.kv_cache_manager_v2 import \ - KVCacheManagerConfig as KVCacheManagerConfigPy -from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, - SsmLayerConfig, - TokenIdExt, _KVCache) - -GB = 1 << 30 - - -class _PrefixReuseDiagnostics(Protocol): - - def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: - ... - - def _get_num_reusable_tokens_before_pruning(self) -> int: - ... - - -# Replay kernels pad the token/window dimension to at least 16 for tensor-core -# tiles, so history sizes below 16 are no faster when not writing and do -# expensive checkpoint writes more often. This minimum tensor-core tile size is -# present in all tensor-core generations, with newer GPUs adding larger tile -# operations, not smaller ones. In the other direction, history sizes above 16 -# currently pad to 32, which is substantially slower with the current kernel -# design. Keep the default floor at 16 while still allowing larger T, which -# degenerates to checkpointing every step instead of rejecting the request. If -# larger T becomes common, we could consider padding larger histories to powers -# of 2 or multiples of 16, but this is untested. For smaller T, we could explore -# combining larger histories with new kernel designs that stay efficient when -# the window is only partly full. -MIN_REPLAY_HISTORY_SIZE = 16 - - -def _get_num_cuda_graph_padding_dummy_slots( - spec_config: Optional["DecodingBaseConfig"], - max_batch_size: int, -) -> int: - """Return the number of persistent CUDA-graph padding dummy IDs. - - This is computed before ``ModelEngine`` exists and covers draft lengths - reachable at every batch size, including the zero-length acceptance-rate - fallback. ``ModelEngine._compute_dynamic_draft_len_mapping`` is created - later and covers only configured CUDA-graph batch sizes, so it cannot size - this persistent ID set. - """ - if spec_config is None: - return 1 - - draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) - spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) - supports_dynamic_draft_len = (spec_dec_mode is not None and hasattr( - spec_dec_mode, "support_dynamic_draft_len") - and spec_dec_mode.support_dynamic_draft_len()) - if draft_len_schedule and supports_dynamic_draft_len: - runtime_draft_lengths = set() - first_uncovered_batch_size = 1 - for batch_size_threshold, draft_len in draft_len_schedule.items(): - if first_uncovered_batch_size > max_batch_size: - break - if batch_size_threshold >= first_uncovered_batch_size: - runtime_draft_lengths.add(draft_len) - first_uncovered_batch_size = batch_size_threshold + 1 - if first_uncovered_batch_size <= max_batch_size: - runtime_draft_lengths.add(0) - else: - max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 - max_total_draft_tokens = (getattr(spec_config, "max_total_draft_tokens", - 0) or 0) - is_linear_tree = getattr( - spec_config, - "is_linear_tree", - max_draft_len == max_total_draft_tokens, - ) - static_draft_len = (max_draft_len - if is_linear_tree else max_total_draft_tokens) - runtime_draft_lengths = {static_draft_len or 0} - - if ((getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and - (getattr(spec_config, "acceptance_rate_threshold", 0) or 0) > 0): - runtime_draft_lengths.add(0) - return len(runtime_draft_lengths) - - -class MambaRole: - """V2 buffer roles owned only by the hybrid Mamba manager.""" - - SSM_STATE = DataRole("ssm_state") - CONV_STATE = DataRole("conv_state") - PLE_NGRAM_CONTEXT = DataRole("ple_ngram_context") - PLE_CONV_STATE = DataRole("ple_conv_state") - - -def _mamba_effective_tp_size(mapping: Mapping) -> int: - """TP degree for sizing per-rank mamba/KDA state pools. - - Attention-DP replicates the state and takes precedence; helix - repurposes CP ranks as plain TP for recurrent-state layers. - """ - if mapping.enable_attention_dp: - return 1 - if mapping.has_cp_helix(): - return mapping.tp_size * mapping.cp_size - return mapping.tp_size - - -def get_tensor_size_bytes(tensor): - """Calculate tensor size in bytes.""" - if isinstance(tensor, torch.Tensor): - return tensor.element_size() * tensor.nelement() - elif isinstance(tensor, list): - return sum(get_tensor_size_bytes(t) for t in tensor) - return 0 - - -# Mamba SSM stochastic-rounding Philox seed plumbing. -# -# Both the replay kernel and flashinfer's selective_state_update consume a -# `rand_seed` int64 tensor. Historically the non-replay paths created this -# tensor via `torch.randint(..., (1,))` on every decode step, which (a) is -# non-deterministic across runs, (b) allocates a fresh CUDA tensor per call -# and is therefore unfriendly to CUDA-graph capture, and (c) has no notion -# of cache-slot identity (so slot reuse cannot rotate the stream). -# -# The functions below produce per-slot int64 seeds via SplitMix64 finalization -# of (counter, slot, rank). Adjacent inputs yield uncorrelated outputs so -# consecutive cache slots and consecutive request counters do not leave -# structural fingerprints in the Philox input stream. All outputs live in -# (0, 2**62) so they avoid Philox's degenerate seed=0 case while staying -# within int64. -_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 -_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 -_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 -_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D -_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 -_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 - - -def _splitmix64(x: int) -> int: - """SplitMix64 finalizer; pure function, no torch.""" - x = (x + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 30) - x = (x * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 27) - x = (x * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK - x ^= (x >> 31) - return x & _MAMBA_SSM_UINT64_MASK - - -def _compute_deterministic_mamba_seed(counter: int, slot: int, - rank_offset: int) -> int: - """Deterministic int64 seed in (0, 2**62) from (counter, slot, rank). - - Pure function (no RNG, no torch.randint). Identical inputs across - process invocations produce identical outputs, which is what the - acceptance criteria require for cross-run reproducibility. - """ - folded = (_MAMBA_SSM_SEED_BASE + counter * _MAMBA_SSM_SEED_MIX_COUNTER + - slot * _MAMBA_SSM_SEED_MIX_SLOT + - rank_offset * _MAMBA_SSM_SEED_MIX_RANK) - folded &= _MAMBA_SSM_UINT64_MASK - value = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK - if value == 0: - value = 1 - return value - - -def _allocate_mamba_seed_buffer(cache_size: int, rank_offset: int, - device: torch.device) -> torch.Tensor: - """Allocate a (cache_size,) int64 CUDA buffer of deterministic seeds. - - counter=0 at allocation time; per-slot reset on fresh request assignment - bumps a host-side counter and rewrites only the freshly-assigned slot. - """ - slot_seeds = [ - _compute_deterministic_mamba_seed(0, i, rank_offset) - for i in range(cache_size) - ] - return torch.tensor(slot_seeds, dtype=torch.int64, device=device) - - -def _mamba_rank_offset(mapping: Mapping) -> int: - """Distinct seed offset per (tp_rank, pp_rank) so independent ranks - don't draw identical streams when not coordinated.""" - return (mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + - mapping.rank * 1_009) - - -def use_py_mamba_cache_manager() -> bool: - """Check if PythonMambaCacheManager should be forced (agg mode override). - - Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. - - Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager - with PythonMambaCacheManager inside instead of the configured manager. - Disagg mode is unaffected — its compatibility routing selects Mixed or - Cpp based on the transceiver configuration. - """ - return os.environ.get('TRTLLM_USE_PY_MAMBA', '0') == '1' - - -class ReplayStateUpdateMetadata(NamedTuple): - """Shared tensors and fixed sizes for replay state updates.""" - prev_num_accepted_tokens: torch.Tensor - cache_buf_idx: torch.Tensor - replay_step_width: int - replay_history_size: int - - -def _advance_replay_state( - replay_metadata: ReplayStateUpdateMetadata, - state_indices: torch.Tensor, - accepted_tokens: torch.Tensor, - is_dummy_request: Optional[torch.Tensor] = None, -) -> None: - """Advance replay bookkeeping after a speculative generation step.""" - slots = state_indices.long() - accepted_tokens = accepted_tokens.to( - replay_metadata.prev_num_accepted_tokens.dtype) - prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens[slots] - wrote_checkpoint = (prev_num_accepted_tokens + - replay_metadata.replay_step_width - > replay_metadata.replay_history_size) - next_num_accepted_tokens = torch.where( - wrote_checkpoint, - accepted_tokens, - prev_num_accepted_tokens + accepted_tokens, - ) - cache_buf_idx = replay_metadata.cache_buf_idx[slots] - next_cache_buf_idx = torch.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx) - if is_dummy_request is not None: - next_num_accepted_tokens = torch.where( - is_dummy_request, - prev_num_accepted_tokens, - next_num_accepted_tokens, - ) - next_cache_buf_idx = torch.where(is_dummy_request, cache_buf_idx, - next_cache_buf_idx) - replay_metadata.prev_num_accepted_tokens[slots] = next_num_accepted_tokens - replay_metadata.cache_buf_idx[slots] = next_cache_buf_idx - - -class BaseMambaCacheManager(ABC): - """Abstract interface for accessing mamba/recurrent state caches.""" - - @abstractmethod - def get_state_indices(self, *args, **kwargs) -> torch.Tensor: - """Return slot indices of each request. - - Shape: [max_batch_size] - """ - ... - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - """Return replay metadata tensors and fixed replay sizes.""" - return None - - @property - def use_kda_replay_update(self) -> bool: - """Whether KDA fused verification owns per-slot replay caches.""" - return getattr(self, "_use_kda_replay_update", False) - - @abstractmethod - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - """Return conv states for specific layer. - - Shape: [slot_size, conv_dim, d_conv - 1] - """ - ... - - @abstractmethod - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - """Return SSM states for specific layer. - - Shape: [slot_size, num_heads, head_dim, d_state] - """ - ... - - @abstractmethod - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - ... - - @abstractmethod - def is_speculative(self) -> bool: - ... - - @abstractmethod - def mamba_layer_cache( - self, layer_idx: int - ) -> Union['PythonMambaCacheManager.State', - 'PythonMambaCacheManager.SpeculativeState', None]: - ... - - -class PythonMambaCacheManager(BaseResourceManager): - """Pure-Python mamba state manager with speculative decoding support. - - Manages only mamba states (conv + SSM) using PyTorch tensors on GPU. - Supports speculative decoding and disaggregated serving. - """ - - @dataclass(frozen=True, kw_only=True) - class State: - """Base state container for Mamba cache.""" - conv: torch.Tensor - # Represents last state, or "two back" if prev_num_accepted_tokens is > 0 - temporal: torch.Tensor - - # Subclasses override to list fields shared across layers (not indexed by layer). - _SHARED_FIELDS = frozenset() - - def at_layer_idx(self, layer: int): - kwargs = {} - for k, v in vars(self).items(): - if v is None: - kwargs[k] = None - elif k in self._SHARED_FIELDS: - kwargs[k] = v - else: - kwargs[k] = v[layer] - return type(self)(**kwargs) - - @property - def has_kda_replay_caches(self) -> bool: - """Whether the Kimi K3 KDA fused-verify replay caches are - allocated. The base (non-speculative) state never has them; the - speculative state overrides this.""" - return False - - @dataclass(frozen=True, kw_only=True) - class SpeculativeState(State): - """Speculative state with intermediate states for draft tokens. - - Supports three SSM update paths (only one set of tensors is - allocated): - - Legacy: caches full intermediate SSM states (intermediate_ssm) - - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) - - KDA replay: per-slot draft-token caches consumed by the fused - ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted - drafts and commits states in place (kda_conv_*, kda_*_cache) - """ - _SHARED_FIELDS = frozenset({ - "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" - }) - - # Allocated for the legacy and Mamba2-replay paths; None for the - # KDA replay path (the kernel commits conv windows in place). - intermediate_conv_window: torch.Tensor | None = None - - # Legacy path: full intermediate SSM states at each step - intermediate_ssm: torch.Tensor | None = None - - # KDA replay path (fused multi-token verify, kimi_linear). - # Pool invariant under this path: `temporal` holds the state after - # the LAST GOLDEN token; the accepted drafts recorded in - # prev_num_accepted_tokens are pending in these caches and are - # replayed by the kernel at the start of the next verify round. - # Extended conv caches [layers, slots, dim, (W-1) + num_spec] fp32, - # dim-contiguous (stride(dim) == 1): columns [0, W-1) are the - # committed raw-input window, tail columns the pending drafts' raw - # inputs. - kda_conv_q: torch.Tensor | None = None - kda_conv_k: torch.Tensor | None = None - kda_conv_v: torch.Tensor | None = None - # Post-processed per-draft quantities for replay: - # [layers, slots, num_spec, 3, H*K] (q/k/gate), [.., num_spec, H*V], - # [.., num_spec, H] — all fp32. - kda_qkg_cache: torch.Tensor | None = None - kda_v_cache: torch.Tensor | None = None - kda_beta_cache: torch.Tensor | None = None - - # Replay path: compact double-buffered cache - # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). - # 0 means temporal saved state is actually the last state, not two back. - prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers - cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers - # Per-cache-slot Philox seeds. Replay bumps them in-place per launch so - # CUDA graph replay uses fresh SR draws without allocating RNG tensors. - # (cache,) int64 - shared across layers - mamba_ssm_rand_seed: torch.Tensor | None = None - old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) - old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) - # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. - old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 - old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 - - @property - def has_kda_replay_caches(self) -> bool: - """True when the KDA fused-verify replay caches were allocated - (the fused ``trtllm::kda_mtp_decode`` verify path).""" - return self.kda_qkg_cache is not None - - def commit_conv_window(self, slot_indices: torch.Tensor, - conv_pool: torch.Tensor) -> None: - """Seed the KDA replay caches from the live ``W - 1`` conv pool.""" - from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ - copy_kda_replay_conv_window - - copy_kda_replay_conv_window( - conv_pool, - self.kda_conv_q, - self.kda_conv_k, - self.kda_conv_v, - slot_indices, - ) - - def __init__( - self, - d_state: int, - d_conv: int, - num_heads: int, - n_groups: int, - head_dim: int, - num_layers: int, - max_batch_size: int, - spec_state_size: int, - mapping: Mapping, - dtype: torch.dtype, - ssm_cache_dtype: torch.dtype, - layer_mask: Optional[List[bool]] = None, - speculative_num_draft_tokens: Optional[int] = None, - model_type: str = "nemotron_hybrid", - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - ) -> None: - - self.mamba_ssm_cache_dtype = ssm_cache_dtype - self.speculative_num_draft_tokens = speculative_num_draft_tokens - self.spec_state_size = spec_state_size - self._use_replay_state_update = use_replay_state_update - # KDA replay path (kimi_linear fused multi-token verify). Mutually - # exclusive with use_replay_state_update; requires speculative mode. - self._kda_replay_num_spec = kda_replay_num_spec - self._use_kda_replay_update = kda_replay_num_spec is not None - if self._use_kda_replay_update: - assert not use_replay_state_update, ( - "kda_replay_num_spec and use_replay_state_update are " - "mutually exclusive") - assert speculative_num_draft_tokens is not None, ( - "KDA replay caches require speculative decoding") - assert kda_replay_num_spec == speculative_num_draft_tokens, ( - f"KDA replay cache width ({kda_replay_num_spec}) must match " - f"the draft length ({speculative_num_draft_tokens}): the " - "fused verify kernel is compiled with a static NUM_SPEC") - self.replay_history_size: Optional[int] = None - self.replay_step_width: Optional[int] = None - # When True, allocate the per-slot Philox seed buffer even outside - # the replay path so the non-replay flashinfer SR kernel reads a - # persistent deterministic seed instead of a per-call torch.randint. - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - # Host-side counter bumped per fresh cache-slot assignment. Combined - # with slot index and rank offset to produce reproducible per-slot - # seed values. Starts at 0 so the post-init "reset" stream is - # disjoint from the counter=0 stream used at allocation time. - self._seed_request_counter = 0 - - # get tp size - tp_size = _mamba_effective_tp_size(mapping) - - # derive mamba parameters for conv and ssm states - d_inner = head_dim * num_heads - conv_dim = d_inner + 2 * n_groups * d_state - nheads = num_heads - - # check that can be partitioned - assert nheads % tp_size == 0, "nheads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - - # partition conv_dim and nheads - d_inner_local = d_inner // tp_size - ng_ds_local = n_groups * d_state // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - d_inner = d_inner // tp_size - - # Per-section dims for conv_state. - # Qwen3-Next: [Q | K | V] = [ng*ds, ng*ds, d_inner] - # Nemotron_hybrid: [x | B | C] = [d_inner, ng*ds, ng*ds] - 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}") - - # conv and ssm states device - device = torch.device("cuda") - - pp_layers, num_layers = get_pp_layers( - num_layers, - mapping, - layer_mask=layer_mask, - ) - num_local_layers = len(pp_layers) - self.mamba_layer_offsets = { - idx: offset - for offset, idx in enumerate(pp_layers) - } - - conv_state_shape = (conv_dim, d_conv - 1) - ssm_state_shape = (nheads, head_dim, d_state) - - # create mamba conv and ssm states - conv_states = torch.zeros( - size=(num_local_layers, max_batch_size) + conv_state_shape, - dtype=dtype, - device=device, - ) - - ssm_states = torch.zeros( - size=(num_local_layers, max_batch_size) + ssm_state_shape, - dtype=self.mamba_ssm_cache_dtype, - device=device, - ) - - # Per-slot Philox seeds. Allocated whenever stochastic rounding can - # fire, even outside the replay path and even when MTP/spec is off, - # so all consumers (replay kernel, MTP non-replay flashinfer, - # non-MTP flashinfer) read a persistent deterministic seed from the - # cache manager instead of calling torch.randint per forward. - self._mamba_ssm_rand_seed: Optional[torch.Tensor] = None - if (self._use_replay_state_update - or self._mamba_ssm_stochastic_rounding): - self._mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - max_batch_size, self._seed_rank_offset, device) - - # create state container - if speculative_num_draft_tokens is not None: - T = speculative_num_draft_tokens + 1 - self.replay_step_width = T - - # Conv intermediate cache — legacy and Mamba2-replay paths only. - # The KDA replay kernel commits conv windows in place, so the - # per-step window scratch is not needed. - intermediate_conv_window_cache = None - if not self._use_kda_replay_update: - intermediate_conv_window_cache = torch.zeros( - size=(num_local_layers, self.spec_state_size, T) + - conv_state_shape, - dtype=dtype, - device=device, - ) - - # SSM speculative cache — path-specific tensors - spec_kwargs = {} - # Share the manager-level seed buffer through SpeculativeState - # so the MTP path can still read it via layer_cache. - if self._mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self._mamba_ssm_rand_seed - if self._use_kda_replay_update: - # KDA replay caches for the fused multi-token verify kernel - # (trtllm::kda_mtp_decode). Per cache slot (persistent - # across rounds — the next round replays from them), unlike - # the batch-row-indexed intermediate buffers. - M = self._kda_replay_num_spec - # conv_dim covers the [q | k | v] short-conv sections; the - # kernel consumes them as three dim-contiguous caches. - assert conv_dim % 3 == 0, ( - "KDA replay caches expect the [q | k | v] conv-state " - "sectioning (3 equal sections)") - section_dim = conv_dim // 3 - # d_conv is KDA's convolution width; the live pool stores - # its W - 1 committed input columns. - w_kernel = d_conv - extended_s = w_kernel - 1 + M - - def _dim_contiguous_conv_cache(): - return torch.zeros(num_local_layers, - max_batch_size, - extended_s, - section_dim, - dtype=torch.float32, - device=device).transpose(-1, -2) - - spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( - max_batch_size, dtype=torch.int32, device=device) - spec_kwargs['kda_conv_q'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_conv_k'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_conv_v'] = _dim_contiguous_conv_cache() - spec_kwargs['kda_qkg_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - 3, - section_dim, - dtype=torch.float32, - device=device) - spec_kwargs['kda_v_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.float32, - device=device) - spec_kwargs['kda_beta_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - nheads, - dtype=torch.float32, - device=device) - ssm_spec_cache = [ - spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], - spec_kwargs['kda_conv_v'], spec_kwargs['kda_qkg_cache'], - spec_kwargs['kda_v_cache'], spec_kwargs['kda_beta_cache'] - ] - spec_path_label = "kda-replay" - elif self._use_replay_state_update: - assert n_groups % tp_size == 0, \ - "replay state update requires n_groups divisible by tp_size" - n_groups_per_rank = n_groups // tp_size - self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) - - # Compact replay cache. - spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( - max_batch_size, dtype=torch.int32, device=device) - spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size, - dtype=torch.int32, - device=device) - spec_kwargs['old_x'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - self.replay_history_size, - nheads, - head_dim, - dtype=dtype, - device=device) - spec_kwargs['old_B'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - self.replay_history_size, - n_groups_per_rank, - d_state, - dtype=dtype, - device=device) - spec_kwargs['old_dt'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - nheads, - self.replay_history_size, - dtype=torch.float32, - device=device) - spec_kwargs['old_dA_cumsum'] = torch.zeros( - num_local_layers, - max_batch_size, - 2, - nheads, - self.replay_history_size, - dtype=torch.float32, - device=device) - ssm_spec_cache = [ - spec_kwargs['old_x'], spec_kwargs['old_B'], - spec_kwargs['old_dt'], spec_kwargs['old_dA_cumsum'] - ] - spec_path_label = "replay" - else: - # Legacy: full intermediate SSM states at each step - spec_kwargs['intermediate_ssm'] = torch.zeros( - size=(num_local_layers, self.spec_state_size, T) + - ssm_state_shape, - dtype=self.mamba_ssm_cache_dtype, - device=device) - ssm_spec_cache = [spec_kwargs['intermediate_ssm']] - spec_path_label = "legacy" - - self.mamba_cache = self.SpeculativeState( - conv=conv_states, - temporal=ssm_states, - intermediate_conv_window=intermediate_conv_window_cache, - **spec_kwargs, - ) - - logger.info( - f"Mamba Cache ({spec_path_label}) is allocated. " - f"max_mamba_cache_size: {max_batch_size}, " - f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " - f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB, " - f"ssm_spec_cache size: {get_tensor_size_bytes(ssm_spec_cache) / GB:.2f}GB, " - "intermediate_conv_window_cache size: " - f"{get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" - ) - else: - self.mamba_cache = self.State( - conv=conv_states, - temporal=ssm_states, - ) - - logger.info( - f"Mamba Cache is allocated. " - f"max_mamba_cache_size: {max_batch_size}, " - f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " - f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB" - ) - - # mamba cache available blocks - self.mamba_cache_free_blocks = [i for i in range(max_batch_size)] - - # mamba cache index, maps request_id -> state indices - self.mamba_cache_index: Dict[int, int] = {} - self._dummy_request_ids: set[int] = set() - # Batch-order mask aligned with state_indices; duplicate dummy request - # IDs mark every batch row even when they share one cache slot. - self._dummy_request_mask = torch.zeros(max_batch_size, - dtype=torch.bool, - device=device) - self._dummy_request_mask_host = torch.zeros(max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned()) - - # Permanent slot shared by every CUDA-graph padding sentinel id - # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per - # draft length). Pool sizing must include +1 headroom for this; - # see MixedMambaHybridCacheManager. - self._padding_slot: int = self.mamba_cache_free_blocks.pop() - - # Reserved slot for the attention-DP padding dummy. Keep it out of the - # free pool so dummy insertion never consumes real-request capacity. - self._attention_dp_dummy_slot: Optional[int] = ( - self.mamba_cache_free_blocks.pop() - if mapping.enable_attention_dp else None) - - # save intermediate state indices for requests - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device=device) - - # Physical tensor rows include reserved dummy slots. Resource capacity - # is the number of real request slots remaining after those - # reservations. - self._max_batch_size = max_batch_size - self._max_resource_count = len(self.mamba_cache_free_blocks) - - def get_max_resource_count(self) -> int: - """Return the maximum number of real requests that can be cached.""" - return self._max_resource_count - - def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: - """Return the prefix of *context_requests* that fits in the - available Mamba cache blocks. Requests that already have a - cached block do not consume a free block.""" - free = len(self.mamba_cache_free_blocks) - result = [] - for r in context_requests: - if r.py_request_id in self.mamba_cache_index: - result.append(r) - elif free > 0: - result.append(r) - free -= 1 - else: - break - return result - - def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - """For Mamba cache manager, we always need one slot per request.""" - return 1 - - @torch.inference_mode() - def _prepare_mamba_cache_blocks(self, request_ids: List[int]): - for r in request_ids: - if r in self.mamba_cache_index: - continue - if len(self.mamba_cache_free_blocks) == 0: - raise RuntimeError("run out of mamba cache blocks") - block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - # Fresh request: no drafts pending in the replay caches. - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - if self._mamba_ssm_rand_seed is not None: - # Deterministic per-slot rotation on fresh assignment. - # `block` is pulled from mamba_cache_free_blocks, which - # excludes _padding_slot by construction (see __init__), - # so padding sentinels never reach this branch. - self._seed_request_counter += 1 - self._mamba_ssm_rand_seed[block] = ( - _compute_deterministic_mamba_seed( - self._seed_request_counter, block, - self._seed_rank_offset)) - - @torch.inference_mode() - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - """Seed the fused-verify KDA replay conv caches from the conv pool. - - On a disaggregated generation server the ctx->gen transfer populates - only the base ``conv`` / ``temporal`` pools (see - ``disaggregation/resource/page.py``); the per-slot ``kda_conv_*`` - replay caches consumed by ``trtllm::kda_mtp_decode`` are normally - seeded by a *local* prefill/decode via - ``_sync_kda_replay_conv_window`` and would otherwise hold zeros (or a - previous occupant's window) for a transferred request — corrupting - the recurrent state from the first verify step onward. Call this - after the state transfer completes, before the first generation - forward. Mirrors ``_sync_kda_replay_conv_window``: the conv pool row - stores the committed ``W - 1`` raw inputs directly. Pending-draft - scratch is cleared (no drafts are pending for a freshly transferred - request). - """ - if not (self._use_kda_replay_update - and isinstance(self.mamba_cache, self.SpeculativeState)): - return - blocks = [ - self.mamba_cache_index[rid] for rid in request_ids - if rid in self.mamba_cache_index - ] - if not blocks: - return - conv = self.mamba_cache.conv # [L, slots, 3D, W - 1] - idx = torch.tensor(sorted(set(blocks)), - dtype=torch.long, - device=conv.device) - d = conv.shape[2] // 3 - committed = conv.shape[3] - cs = conv.index_select(1, idx) - for cache, section in ( - (self.mamba_cache.kda_conv_q, cs[:, :, :d]), - (self.mamba_cache.kda_conv_k, cs[:, :, d:2 * d]), - (self.mamba_cache.kda_conv_v, cs[:, :, 2 * d:]), - ): - # cache: [L, slots, D, committed + num_spec]; zero the draft - # tail columns and seed the committed window in one copy. - seeded = torch.zeros( - (cache.shape[0], idx.numel()) + cache.shape[2:], - dtype=cache.dtype, - device=cache.device) - seeded[:, :, :, :committed] = section.to(cache.dtype) - cache.index_copy_(1, idx, seeded) - for buf in (self.mamba_cache.kda_qkg_cache, - self.mamba_cache.kda_v_cache, - self.mamba_cache.kda_beta_cache): - buf.index_fill_(1, idx, 0) - self.mamba_cache.prev_num_accepted_tokens[idx] = 0 - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - context_ids = [ - i.py_request_id for i in scheduled_batch.context_requests - ] - generation_ids = [ - i.py_request_id for i in scheduled_batch.generation_requests - ] - request_ids = context_ids + generation_ids - self._prepare_mamba_cache_blocks(request_ids) - self._refresh_dummy_request_mask([req.is_dummy for req in requests]) - - def _is_padding_sentinel(self, request_id: int) -> bool: - # cuda_graph_runner caches one dummy per runtime_draft_len value - # (see _get_padded_batch), so any id in the range of dummy request IDs - # may be live concurrently. - from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ - CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl = self.speculative_num_draft_tokens or 0 - return (CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= - CUDA_GRAPH_DUMMY_REQUEST_ID) - - @torch.inference_mode() - def add_dummy_requests(self, request_ids: List[int], **kwargs): - # Sentinels alias to the shared _padding_slot; non-sentinel - # dummies (warmup, attention-DP idle padding) get their own - # slot and are freed individually. - if not request_ids: - return - self._dummy_request_ids.update(request_ids) - for r in request_ids: - if r in self.mamba_cache_index: - block = self.mamba_cache_index[r] - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - continue - if self._is_padding_sentinel(r): - block = self._padding_slot - elif (r == ATTENTION_DP_DUMMY_REQUEST_ID - and self._attention_dp_dummy_slot is not None): - block = self._attention_dp_dummy_slot - else: - if len(self.mamba_cache_free_blocks) == 0: - raise RuntimeError("run out of mamba cache blocks") - block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block - if (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_replay_state_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - self.mamba_cache.cache_buf_idx[block] = 0 - elif (isinstance(self.mamba_cache, self.SpeculativeState) - and self._use_kda_replay_update): - self.mamba_cache.prev_num_accepted_tokens[block] = 0 - - def free_resources(self, request: LlmRequest): - request_id = request.py_request_id - if request_id not in self.mamba_cache_index: - return - self._dummy_request_ids.discard(request_id) - block = self.mamba_cache_index.pop(request_id) - # Reserved slots must not re-enter the real-request free pool. - if block != self._padding_slot and \ - block != self._attention_dp_dummy_slot: - self.mamba_cache_free_blocks.append(block) - - def get_state_indices(self, request_ids: List[int], - is_padding: List[bool]) -> List[int]: - assert len(request_ids) == len(is_padding) - indices = [self.mamba_cache_index[rid] for rid in request_ids] - is_dummy = [ - rid in self._dummy_request_ids or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.as_tensor(is_dummy, dtype=torch.bool)) - self._dummy_request_mask.copy_(self._dummy_request_mask_host, - non_blocking=True) - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).conv - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).temporal - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if not isinstance(self.mamba_cache, self.SpeculativeState): - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset).intermediate_ssm - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - """Get intermediate conv states for speculative decoding.""" - if not isinstance(self.mamba_cache, self.SpeculativeState): - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx( - layer_offset).intermediate_conv_window - - def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_x[layer_offset] - - def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_B[layer_offset] - - def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_dt[layer_offset] - - def get_replay_old_dA_cumsum(self, - layer_idx: int) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.old_dA_cumsum[layer_offset] - - def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - return self.mamba_cache.cache_buf_idx - - def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: - if not self._use_replay_state_update: - return None - return self.mamba_cache.prev_num_accepted_tokens - - def is_speculative(self) -> bool: - return isinstance(self.mamba_cache, self.SpeculativeState) - - def mamba_layer_cache(self, - layer_idx: int) -> Union[State, SpeculativeState]: - layer_offset = self.mamba_layer_offsets[layer_idx] - return self.mamba_cache.at_layer_idx(layer_offset) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.mamba_ssm_cache_dtype - - 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. - - Used by mamba2_mixer non-MTP paths that don't hold a SpeculativeState. - Callers must bump in-place (`.add_(1)` or slice-and-add) and pass a - view that matches the consuming kernel's expected shape. - """ - return self._mamba_ssm_rand_seed - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - @property - def use_kda_replay_update(self) -> bool: - return self._use_kda_replay_update - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - if (not self._use_replay_state_update - or not isinstance(self.mamba_cache, self.SpeculativeState) - or self.mamba_cache.prev_num_accepted_tokens is None - or self.mamba_cache.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=( - self.mamba_cache.prev_num_accepted_tokens), - cache_buf_idx=self.mamba_cache.cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def shutdown(self): - """Release tensor memory.""" - # Clear mamba cache states - empty = torch.tensor([]) - - def _drop(tensor): - return empty if tensor is not None else None - - if isinstance(self.mamba_cache, self.SpeculativeState): - self.mamba_cache = self.SpeculativeState( - conv=empty, - temporal=empty, - intermediate_conv_window=_drop( - self.mamba_cache.intermediate_conv_window), - intermediate_ssm=_drop(self.mamba_cache.intermediate_ssm), - prev_num_accepted_tokens=_drop( - self.mamba_cache.prev_num_accepted_tokens), - cache_buf_idx=_drop(self.mamba_cache.cache_buf_idx), - mamba_ssm_rand_seed=_drop(self.mamba_cache.mamba_ssm_rand_seed), - old_x=_drop(self.mamba_cache.old_x), - old_B=_drop(self.mamba_cache.old_B), - old_dt=_drop(self.mamba_cache.old_dt), - old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), - kda_conv_q=_drop(self.mamba_cache.kda_conv_q), - kda_conv_k=_drop(self.mamba_cache.kda_conv_k), - kda_conv_v=_drop(self.mamba_cache.kda_conv_v), - kda_qkg_cache=_drop(self.mamba_cache.kda_qkg_cache), - kda_v_cache=_drop(self.mamba_cache.kda_v_cache), - kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), - ) - else: - self.mamba_cache = self.State(conv=empty, temporal=empty) - - torch.cuda.empty_cache() - - @torch.compile(options={"max-autotune": True}) - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor, - accepted_leaf_positions: Optional[torch.Tensor] = None): - 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 - # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. - accepted_positions = (accepted_leaf_positions if accepted_leaf_positions - is not None else num_accepted_draft_tokens) - state_indices_d = state_indices[num_contexts:num_contexts + num_gens] - src_state_indices = self.intermediate_state_indices[:num_gens] - - if self._use_kda_replay_update: - # KDA replay: the fused verify kernel already committed the SSM - # state and conv windows (in place, after the golden token) and - # cached this round's drafts. All that remains is recording how - # many of those drafts the sampler accepted, so the next round's - # kernel launch replays exactly that prefix. - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - prev = self.mamba_cache.prev_num_accepted_tokens - current = prev[state_indices_d] - accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) - prev[state_indices_d] = torch.where(is_dummy_request, current, - accepted) - return - - if self._use_replay_state_update: - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - replay_metadata = self.get_replay_state_update_metadata() - assert replay_metadata is not None - _advance_replay_state( - replay_metadata, - state_indices_d, - num_accepted_tokens[num_contexts:num_contexts + num_gens], - is_dummy_request, - ) - else: - # Legacy: copy accepted SSM state from intermediate cache. - ssm_states = self.mamba_cache.temporal - intermediate_ssm_cache = self.mamba_cache.intermediate_ssm - accepted_ssm_state = intermediate_ssm_cache[:, src_state_indices, - accepted_positions] - ssm_states[:, state_indices_d, :] = accepted_ssm_state - - # Conv: both paths save all intermediate conv windows, carry over the accepted one. - conv_states = self.mamba_cache.conv - intermediate_conv_window_cache = self.mamba_cache.intermediate_conv_window - accepted_conv_state = intermediate_conv_window_cache[:, - src_state_indices, - accepted_positions] - conv_states[:, state_indices_d, :] = accepted_conv_state - - -class MambaCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Facade for standalone mamba state management (no KV cache). - - Delegates to PythonMambaCacheManager. - """ - - def __init__( - self, - d_state: int, - d_conv: int, - num_heads: int, - n_groups: int, - head_dim: int, - num_layers: int, - max_batch_size: int, - spec_state_size: int, - mapping: Mapping, - dtype: torch.dtype, - ssm_cache_dtype: torch.dtype, - layer_mask: Optional[List[bool]] = None, - stream: Optional[torch.cuda.Stream] = None, - speculative_num_draft_tokens: Optional[int] = None, - model_type: str = "nemotron_hybrid", - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - ) -> None: - max_num_sequences = max_batch_size * mapping.pp_size - - self._impl = PythonMambaCacheManager( - d_state=d_state, - d_conv=d_conv, - num_heads=num_heads, - n_groups=n_groups, - head_dim=head_dim, - num_layers=num_layers, - max_batch_size=max_num_sequences, - spec_state_size=spec_state_size, - mapping=mapping, - dtype=dtype, - ssm_cache_dtype=ssm_cache_dtype, - layer_mask=layer_mask, - speculative_num_draft_tokens=speculative_num_draft_tokens, - model_type=model_type, - use_replay_state_update=use_replay_state_update, - mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - kda_replay_num_spec=kda_replay_num_spec, - ) - - def get_max_resource_count(self) -> int: - return self._impl.get_max_resource_count() - - def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: - return self._impl.filter_ctx_requests_by_capacity(context_requests) - - def get_needed_resource_to_completion(self, request: LlmRequest) -> int: - return self._impl.get_needed_resource_to_completion(request) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - self._impl.prepare_resources(scheduled_batch) - - def free_resources(self, request: LlmRequest): - self._impl.free_resources(request) - - def add_dummy_requests(self, request_ids: List[int], **kwargs): - self._impl.add_dummy_requests(request_ids, **kwargs) - - def get_state_indices( - self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None - ) -> Union[torch.Tensor, List[int]]: - return self._impl.get_state_indices(request_ids, is_padding) - - @property - def mamba_cache_free_blocks(self) -> List[int]: - return self._impl.mamba_cache_free_blocks - - @property - def mamba_cache_index(self) -> Dict[int, int]: - return self._impl.mamba_cache_index - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self._impl.get_conv_states(layer_idx) - - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - self._impl.seed_kda_replay_caches_for_disagg_gen(request_ids) - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self._impl.get_ssm_states(layer_idx) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self._impl.get_mamba_ssm_cache_dtype() - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - """Delegate to the underlying Python manager.""" - return self._impl.get_mamba_ssm_rand_seed() - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - @property - def use_kda_replay_update(self) -> bool: - return self._impl.use_kda_replay_update - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - return self._impl.get_replay_state_update_metadata() - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_intermediate_ssm_states(layer_idx) - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_intermediate_conv_states(layer_idx) - - def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_x(layer_idx) - - def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_B(layer_idx) - - def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_dt(layer_idx) - - def get_replay_old_dA_cumsum(self, - layer_idx: int) -> Optional[torch.Tensor]: - return self._impl.get_replay_old_dA_cumsum(layer_idx) - - def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: - return self._impl.get_replay_cache_buf_idx() - - def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: - return self._impl.get_replay_prev_num_accepted_tokens() - - def is_speculative(self) -> bool: - return self._impl.is_speculative() - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - return self._impl.mamba_layer_cache(layer_idx) - - def shutdown(self): - self._impl.shutdown() - - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: torch.Tensor, - accepted_leaf_positions: Optional[torch.Tensor] = None): - # Non-speculative configs don't allocate intermediate state; the - # promotion is a clean no-op. - if not self._impl.is_speculative(): - return - self._impl.update_mamba_states(attn_metadata, num_accepted_tokens, - state_indices, accepted_leaf_positions) - - -class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): - """Shared interface and state plumbing for hybrid Mamba managers. - - Concrete storage, state-index, and resource lifecycles remain owned by the - Cpp and V2 implementations. - """ - - _supports_additional_snapshot_offsets = False - - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size: int) -> None: - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if spec_config is None or self.local_num_mamba_layers == 0: - return - - tokens_per_gen_step = spec_config.tokens_per_gen_step - if not self._use_replay_state_update: - self.intermediate_ssm_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, - tokens_per_gen_step - ] + self.ssm_state_shape, - dtype=self.ssm_state_dtype, - device="cuda", - ) - self.intermediate_conv_states = torch.zeros( - size=[ - self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step - ] + self.conv_state_shape, - dtype=self.conv_state_dtype, - device="cuda", - ) - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - - def _allocate_pool_replay_buffers( - self, - spec_config, - cache_size: int, - device: Optional[torch.device], - ) -> bool: - """Allocate replay tensors shared by the Cpp and V2 state pools.""" - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - - if (self.local_num_mamba_layers == 0 - or (not self._use_replay_state_update - and not self._mamba_ssm_stochastic_rounding)): - return False - - assert device is not None - self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( - cache_size, self._seed_rank_offset, device) - if spec_config is None or not self._use_replay_state_update: - return False - - history_size = self.replay_history_size - assert history_size is not None - nheads, head_dim, d_state = self.ssm_state_shape - common_shape = [self.local_num_mamba_layers, cache_size, 2] - self.prev_num_accepted_tokens = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.cache_buf_idx = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.old_x = torch.zeros( - common_shape + [history_size, nheads, head_dim], - dtype=self.conv_state_dtype, - device=device, - ) - self.old_B = torch.zeros( - common_shape + [history_size, self._n_groups_per_rank, d_state], - dtype=self.conv_state_dtype, - device=device, - ) - self.old_dt = torch.zeros( - common_shape + [nheads, history_size], - dtype=torch.float32, - device=device, - ) - self.old_dA_cumsum = torch.zeros( - common_shape + [nheads, history_size], - dtype=torch.float32, - device=device, - ) - return True - - @torch.inference_mode() - def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: - if self._dummy_request_mask is None: - return - - n = len(is_dummy) - assert n <= self._dummy_request_mask_host.shape[0] - self._dummy_request_mask_host.zero_() - if n > 0: - self._dummy_request_mask_host[:n].copy_( - torch.tensor(is_dummy, dtype=torch.bool)) - self._dummy_request_mask.copy_(self._dummy_request_mask_host, - non_blocking=True) - - def _reset_context_mamba_slots(self, num_contexts: int) -> None: - if num_contexts == 0: - return - - context_slots = self.cuda_state_indices[:num_contexts].long() - if (self._use_replay_state_update - and self.prev_num_accepted_tokens is not None - and self.cache_buf_idx is not None): - self.prev_num_accepted_tokens[context_slots] = 0 - self.cache_buf_idx[context_slots] = 0 - if self.old_x is not None: - self.old_x[:, context_slots] = 0 - if self.old_B is not None: - self.old_B[:, context_slots] = 0 - if self.old_dt is not None: - self.old_dt[:, context_slots] = 0 - if self.old_dA_cumsum is not None: - self.old_dA_cumsum[:, context_slots] = 0 - - if self.mamba_ssm_rand_seed is None: - return - self._seed_request_counter += 1 - counter = self._seed_request_counter - rank_offset = self._seed_rank_offset - host_slots = self._host_state_indices[:num_contexts].tolist() - new_seeds = [ - _compute_deterministic_mamba_seed(counter, slot, rank_offset) - for slot in host_slots - ] - self.mamba_ssm_rand_seed[context_slots] = torch.tensor( - new_seeds, - dtype=torch.int64, - device=self.mamba_ssm_rand_seed.device, - ) - - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - """Set reusable Mamba snapshot boundaries before scheduling.""" - if not self.enable_block_reuse: - for request in requests: - request.expect_snapshot_points = [] - return - - state_config = self.kv_cache_config.mamba_state_config - interval = state_config.periodic_snapshot_interval - for request in requests: - snapshot_points = set() - if interval is not None and interval > 0: - snapshot_points.update( - range(interval, request.prompt_len + 1, interval)) - if self._supports_additional_snapshot_offsets: - for offset in state_config.additional_snapshot_offsets_from_start: - if offset <= request.prompt_len: - snapshot_points.add(offset) - for offset in state_config.additional_snapshot_offsets_from_end: - point = request.prompt_len - offset - if point > 0: - snapshot_points.add(point) - request.expect_snapshot_points = sorted(snapshot_points) - - def is_speculative(self) -> bool: - return self.spec_config is not None - - def get_ssm_states(self, layer_idx: int) -> torch.Tensor: - return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_conv_states(self, layer_idx: int) -> torch.Tensor: - return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_ssm_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_ssm_states is None: - return None - return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] - - def get_intermediate_conv_states(self, - layer_idx: int) -> Optional[torch.Tensor]: - if self.intermediate_conv_states is None: - return None - return self.intermediate_conv_states[ - self.mamba_layer_offsets[layer_idx]] - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - conv = self.get_conv_states(layer_idx) - ssm = self.get_ssm_states(layer_idx) - if self.spec_config is not None: - layer_offset = self.mamba_layer_offsets[layer_idx] - spec_kwargs = {} - if self.mamba_ssm_rand_seed is not None: - spec_kwargs['mamba_ssm_rand_seed'] = self.mamba_ssm_rand_seed - if self._use_replay_state_update: - spec_kwargs['old_x'] = self.old_x[layer_offset] - spec_kwargs['old_B'] = self.old_B[layer_offset] - spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] - spec_kwargs['cache_buf_idx'] = self.cache_buf_idx - spec_kwargs['prev_num_accepted_tokens'] = ( - self.prev_num_accepted_tokens) - else: - spec_kwargs['intermediate_ssm'] = self.intermediate_ssm_states[ - layer_offset] - return PythonMambaCacheManager.SpeculativeState( - conv=conv, - temporal=ssm, - intermediate_conv_window=self. - intermediate_conv_states[layer_offset], - **spec_kwargs, - ) - return PythonMambaCacheManager.State(conv=conv, temporal=ssm) - - @property - def use_replay_state_update(self) -> bool: - return self.get_replay_state_update_metadata() is not None - - def get_replay_state_update_metadata( - self) -> Optional[ReplayStateUpdateMetadata]: - prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', - None) - cache_buf_idx = getattr(self, 'cache_buf_idx', None) - if (not self._use_replay_state_update - or prev_num_accepted_tokens is None or cache_buf_idx is None - or self.replay_step_width is None - or self.replay_history_size is None): - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=prev_num_accepted_tokens, - cache_buf_idx=cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size) - - def get_mamba_ssm_cache_dtype(self) -> torch.dtype: - return self.ssm_state_dtype - - def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: - return getattr(self, 'mamba_ssm_rand_seed', None) - - -def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: - """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" - pool_size = max_batch_size - # One permanent slot is shared by every CUDA-graph padding sentinel. - pool_size += 1 - if mapping.enable_attention_dp: - # Attention-DP can insert a transient dummy request on an otherwise - # idle rank. Keep this headroom internal so scheduler-visible - # max_batch_size still limits real requests. - pool_size += 1 - return pool_size - - -class MixedMambaHybridCacheManager(KVCacheManager, MambaCacheManager, - MambaHybridCacheManager): - """Hybrid cache manager combining separate KVCacheManager and MambaCacheManager. - - Manages KV cache and mamba states in independent pools, with support of - speculative decoding and disaggregated serving. - Does not support block reuse / prefix caching for mamba states. - """ - - 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 parameters - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - layer_mask: List[bool], - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. - # It's derived from the model's BuildConfig for consistency with the C++ backend. - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - max_num_tokens: int = 8192, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - is_estimating_kv_cache: bool = False, - execution_stream: Optional[torch.cuda.Stream] = None, - model_type: str = "nemotron_hybrid", - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - kda_replay_num_spec: Optional[int] = None, - # Per-pool configurations forwarded to the C++ KVCacheManager ctor. - # Lets a single manager host pools with mixed shapes (e.g. Gemma4 - # hybrid attention). See KVCacheManager.__init__. - pool_configurations: Optional[List[PoolConfiguration]] = None, - ) -> None: - - # mamba hybrid cache requires block reuse to be disabled in KV cache config - assert not kv_cache_config.enable_block_reuse, ( - "mamba hybrid cache requires block reuse to be disabled in KV cache config" - ) - - # Host-drafter spec modes (NGram) have no spec worker to call - # update_mamba_states; this manager promotes accepted states itself - # in update_resources. One-model modes (MTP/Eagle/DFlash) and the - # suffix-automaton worker promote from their spec workers and must - # NOT be promoted twice. - self._promote_states_in_update_resources = ( - spec_config is not None - and getattr(spec_config, "decoding_type", None) == "NGram") - - pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) - - MambaCacheManager.__init__( - self, - mamba_d_state, - mamba_d_conv, - mamba_num_heads, - mamba_n_groups, - mamba_head_dim, - mamba_num_layers, - pool_size, - max_batch_size, - mapping, - mamba_cache_dtype, - mamba_ssm_cache_dtype, - mamba_layer_mask, - execution_stream, - speculative_num_draft_tokens=(spec_config.tokens_per_gen_step - 1 - if spec_config is not None else None), - model_type=model_type, - use_replay_state_update=use_replay_state_update, - mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, - kda_replay_num_spec=kda_replay_num_spec, - ) - - # initialize kv cache manager - KVCacheManager.__init__( - self, - kv_cache_config, - kv_cache_type, - num_layers=num_layers, - num_kv_heads=num_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_num_tokens=max_num_tokens, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - execution_stream=execution_stream, - is_draft=is_draft, - pool_configurations=pool_configurations, - ) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - MambaCacheManager.prepare_resources(self, scheduled_batch) - KVCacheManager.prepare_resources(self, scheduled_batch) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - MambaCacheManager.free_resources(self, request) - KVCacheManager.free_resources(self, request, pin_on_release) - - def add_dummy_requests(self, request_ids: List[int], **kwargs): - MambaCacheManager.add_dummy_requests(self, request_ids) - return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) - - def shutdown(self): - MambaCacheManager.shutdown(self) - KVCacheManager.shutdown(self) - - def update_resources(self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None): - KVCacheManager.update_resources(self, scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - self._maybe_promote_drafter_states(scheduled_batch, attn_metadata) - - def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): - """Commit accepted verification states for host-drafter spec modes. - - One-model spec workers (MTP/Eagle/DFlash) and the suffix-automaton - worker call update_mamba_states themselves right after on-device - acceptance. Host-drafter modes (NGram) have no spec worker: - acceptance lands on the requests as - ``py_num_accepted_draft_tokens`` during sampler update, and this - hook — running right after, alongside the KV rewind — promotes the - accepted step's intermediate state into the live pools (or, on the - KDA replay path, records the accepted-draft count for the next - round's replay). Without it, the pools would keep the - pre-verification state and the next forward would resume from a - stale prefix. - """ - if not self._promote_states_in_update_resources: - return - if not self.is_speculative() or attn_metadata is None: - return - gen_requests = scheduled_batch.generation_requests - if not gen_requests: - return - drafted = [ - r for r in gen_requests - if r.py_draft_tokens is not None and len(r.py_draft_tokens) > 0 - ] - if not drafted: - # Drafter skipped this step (e.g. speculation gated off): the - # forward ran the plain in-place decode path; nothing to promote. - return - assert len(drafted) == len(gen_requests), ( - "mixed drafted/undrafted generation batch is not supported for " - "hybrid state promotion (drafts are padded to the static max)") - device = self._impl.mamba_cache.temporal.device - num_contexts = len(scheduled_batch.context_requests) - num_accepted = torch.tensor( - [0] * num_contexts + - [r.py_num_accepted_draft_tokens + 1 for r in gen_requests], - dtype=torch.int32, - device=device) - # Batch-ordered slots (contexts then gens), matching the ordering - # the forward used for the intermediate scratch buffers. Requests - # that finished this step were already freed by response handling - # (which runs before update_resources) and are gone from the index; - # their rows must stay in place for alignment, so redirect them to - # the reserved padding slot (a harmless scratch write). - slot_index = self.mamba_cache_index - padding_slot = self._impl._padding_slot - state_indices = torch.tensor([ - slot_index.get(r.py_request_id, padding_slot) - for r in scheduled_batch.context_requests + gen_requests - ], - dtype=torch.int32, - device=device) - self.update_mamba_states(attn_metadata, num_accepted, state_indices) - - -@triton.jit -def _promote_mamba_state_kernel( - src_ptr, - dst_ptr, - src_idx_ptr, - accepted_position_source_ptr, - num_accepted_tokens_ptr, - blk_ptr, - replay_pnat_ptr, - replay_cache_buf_idx_ptr, - dummy_request_mask_ptr, - num_gens, - count, - src_s_layer, - src_s_row, - src_s_step, - dst_s_layer, - dst_s_block, - POSITION_SOURCE_IS_TOKEN_COUNT: tl.constexpr, - REPLAY_STEP_WIDTH: tl.constexpr, - REPLAY_HISTORY_SIZE: tl.constexpr, - UPDATE_REPLAY_BOOKKEEPING: tl.constexpr, - BLOCK: tl.constexpr, -): - # One program copies a BLOCK-sized tile of the contiguous inner state for a - # single (layer, gen) pair. grid = (num_layers * num_gens, ceil(count/BLOCK)). - pair = tl.program_id(0) - tile = tl.program_id(1) - layer = (pair // num_gens).to(tl.int64) - g = pair % num_gens - row = tl.load(src_idx_ptr + g).to(tl.int64) - accepted_position = tl.load(accepted_position_source_ptr + g) - if POSITION_SOURCE_IS_TOKEN_COUNT: - accepted_position -= 1 - acc = accepted_position.to(tl.int64) - blk = tl.load(blk_ptr + g).to(tl.int64) - - # The replay metadata update is batch-scoped, so layer 0's first copy tile - # owns it. This folds PNAT and active-buffer maintenance into the existing - # accepted-convolution-state promotion launch. - if UPDATE_REPLAY_BOOKKEEPING: - if (layer == 0) & (tile == 0): - num_accepted_tokens = tl.load(num_accepted_tokens_ptr + g) - is_dummy_request = tl.load(dummy_request_mask_ptr + g) - pnat = tl.load(replay_pnat_ptr + blk) - cache_buf_idx = tl.load(replay_cache_buf_idx_ptr + blk) - wrote_checkpoint = pnat + REPLAY_STEP_WIDTH > REPLAY_HISTORY_SIZE - next_pnat = tl.where(wrote_checkpoint, num_accepted_tokens, - pnat + num_accepted_tokens) - next_cache_buf_idx = tl.where(wrote_checkpoint, 1 - cache_buf_idx, - cache_buf_idx) - tl.store(replay_pnat_ptr + blk, next_pnat, mask=~is_dummy_request) - tl.store(replay_cache_buf_idx_ptr + blk, - next_cache_buf_idx, - mask=~is_dummy_request) - - # int64 throughout: per-layer strides are O(1e8), so layer*stride overflows - # int32 and would corrupt addresses / fault. - src_base = layer * src_s_layer + row * src_s_row + acc * src_s_step - dst_base = layer * dst_s_layer + blk * dst_s_block - offs = tile * BLOCK + tl.arange(0, BLOCK) - mask = offs < count - v = tl.load(src_ptr + src_base + offs, mask=mask) - tl.store(dst_ptr + dst_base + offs, v, mask=mask) - - -def _promote_mamba_state_triton( - dst: torch.Tensor, - intermediate: torch.Tensor, - src_state_indices: torch.Tensor, - accepted_position_source: torch.Tensor, - dst_state_indices: torch.Tensor, - num_accepted_tokens: Optional[torch.Tensor] = None, - position_source_is_token_count: bool = False, - replay_pnat: Optional[torch.Tensor] = None, - replay_cache_buf_idx: Optional[torch.Tensor] = None, - dummy_request_mask: Optional[torch.Tensor] = None, - replay_step_width: int = 0, - replay_history_size: int = 0, - BLOCK: int = 2048) -> None: - """Scatter each generation request's accepted draft-step recurrent state - from the per-request intermediate buffer into the unified C++ pool view. - - Args: - dst: ``[num_layers, num_blocks, *state_shape]`` view of the C++ pool. - The block dim (dim 1) is strided (the pool interleaves - ``ssm_bytes | conv_bytes`` per block), so ``dst`` is non-contiguous. - The kernel writes through ``dst``'s real strides, touching only this - state's bytes -- the case that defeats torch.compile / aot_autograd - (dtype-view mutation) and Inductor codegen (``XBLOCK`` on the uint8 - pool). - intermediate: ``[num_layers, max_batch_size, T, *state_shape]`` dense. - src_state_indices, accepted_position_source, dst_state_indices: - ``[num_gens]`` int. - num_accepted_tokens: ``[num_gens]`` accepted-token counts. Required - when replay bookkeeping is enabled. - position_source_is_token_count: Whether ``accepted_position_source`` - contains accepted-token counts instead of zero-based intermediate - state positions. When true, the kernel subtracts one while - computing the source position. - replay_pnat, replay_cache_buf_idx, dummy_request_mask: Optional replay - metadata. When supplied together, the kernel advances PNAT and the - active history buffer while promoting the convolution state. - - For each gen ``g``:: - - accepted_position = ( - accepted_position_source[g] - 1 - if position_source_is_token_count - else accepted_position_source[g] - ) - dst[:, dst_state_indices[g]] = intermediate[ - :, src_state_indices[g], accepted_position - ] - - Pure gather->scatter copy; bandwidth-bound (~85% of HBM peak), one launch. - """ - num_layers = dst.shape[0] - num_gens = src_state_indices.shape[0] - if num_gens == 0: - return - update_replay_bookkeeping = replay_pnat is not None - if update_replay_bookkeeping: - assert num_accepted_tokens is not None - assert replay_cache_buf_idx is not None - assert dummy_request_mask is not None - assert replay_step_width > 0 - assert replay_history_size >= replay_step_width - assert replay_pnat.dtype == torch.int32 - assert replay_cache_buf_idx.dtype == torch.int32 - assert dummy_request_mask.dtype == torch.bool - assert dummy_request_mask.numel() >= num_gens - else: - # Compile-time-disabled pointer arguments still require tensors. - num_accepted_tokens = accepted_position_source - assert replay_cache_buf_idx is None - assert dummy_request_mask is None - replay_pnat = dst_state_indices - replay_cache_buf_idx = dst_state_indices - dummy_request_mask = dst_state_indices - count = 1 - for s in dst.shape[2:]: - count *= s - # Grid dim order matters: dim 0 -> CUDA grid.x (limit 2^31-1), dim 1 -> - # grid.y (limit 65535). Put the (layer, gen) pairs in dim 0 since that is - # what grows with batch size (num_layers*num_gens stays far below 2^31 for - # any real config); the tile count in dim 1 is small (~count/BLOCK). Do NOT - # swap them: pairs would hit the 65535 y-limit at num_layers*num_gens>65535. - grid = (num_layers * num_gens, triton.cdiv(count, BLOCK)) - _promote_mamba_state_kernel[grid]( - intermediate, - dst, - src_state_indices, - accepted_position_source, - num_accepted_tokens, - dst_state_indices, - replay_pnat, - replay_cache_buf_idx, - dummy_request_mask, - num_gens, - count, - intermediate.stride(0), - intermediate.stride(1), - intermediate.stride(2), - dst.stride(0), - dst.stride(1), - POSITION_SOURCE_IS_TOKEN_COUNT=position_source_is_token_count, - REPLAY_STEP_WIDTH=replay_step_width, - REPLAY_HISTORY_SIZE=replay_history_size, - UPDATE_REPLAY_BOOKKEEPING=update_replay_bookkeeping, - BLOCK=BLOCK, - ) - - -def _mamba_snapshot_rule_counts( - kv_cache_config: KvCacheConfig, - max_seq_len: Optional[int], - tokens_per_block: int, -) -> Tuple[int, int]: - """Return reachable fixed rules and their partial-block upper bound.""" - if not kv_cache_config.enable_block_reuse: - return 0, 0 - - num_rules = 0 - num_unaligned_rules = 0 - state_config = kv_cache_config.mamba_state_config - for offset in set(state_config.additional_snapshot_offsets_from_start): - if max_seq_len is not None and offset > max_seq_len: - continue - num_rules += 1 - num_unaligned_rules += int(offset % tokens_per_block != 0) - for offset in set(state_config.additional_snapshot_offsets_from_end): - # A from-end offset is reachable iff some valid prompt is longer than - # the offset. Its absolute alignment depends on that prompt. - if max_seq_len is not None and offset >= max_seq_len: - continue - num_rules += 1 - num_unaligned_rules += 1 - return num_rules, num_unaligned_rules - - -def _mamba_regular_snapshot_interval( - kv_cache_config: KvCacheConfig, - max_seq_len: Optional[int], -) -> Optional[int]: - if not kv_cache_config.enable_block_reuse: - return None - interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval - if interval is None or interval <= 0: - return None - if max_seq_len is not None and interval > max_seq_len: - return None - return interval - - -def _get_local_mamba_cache_layout( - model_config, - mapping: Mapping, - *, - spec_config=None, - is_draft: bool = False, - use_separate_draft_kv_cache: bool = False, -): - """Return normalized params and local Mamba/attention layer counts. - - Cache construction and affine sizing must follow the model's PP layout: - partition base layers first, then place appended speculative layers on the - last PP rank. The normalized params retain target masks and the appended - draft-layer count so estimation selects the same combined or per-manager - layout as runtime. - """ - from tensorrt_llm._torch.pyexecutor.config_utils import \ - extract_mamba_kv_cache_params - - params = extract_mamba_kv_cache_params( - model_config.pretrained_config, - spec_config=spec_config, - quant_config=model_config.quant_config, - ) - mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( - is_draft=is_draft, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - combined_layer_mask = [ - is_mamba or is_attention for is_mamba, is_attention in zip( - mamba_layer_mask, full_attention_layer_mask) - ] - local_layer_indices, _ = get_pp_layers( - sum(combined_layer_mask), - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - local_mamba_layers = sum(mamba_layer_mask[layer_idx] - for layer_idx in local_layer_indices) - local_attention_layers = sum(full_attention_layer_mask[layer_idx] - for layer_idx in local_layer_indices) - return params, local_mamba_layers, local_attention_layers - - -# --- Qwen4-Exp PLE --------------------------------------------------------- -# PLE (the n-gram side path) is a Qwen4-Exp-only feature that rides on the -# hybrid manager's Mamba slots. It is kept in dedicated helpers so it can be -# lifted into a model-owned cache manager wholesale later on. - - -def _qwen4_exp_ple_state_bytes_per_rank( - model_config, - params, - mapping, - *, - spec_config, - use_separate_draft_kv_cache: bool, -) -> int: - """Bytes this rank needs for Qwen4-Exp PLE recurrent state; 0 otherwise.""" - from tensorrt_llm._torch.pyexecutor.config_utils import ( - extract_qwen4_exp_ple_cache_params, is_qwen4_exp) - - # The estimator is also called with minimal model-config stubs that carry no - # `pretrained_config`; such a config is definitionally not Qwen4-Exp. - pretrained_config = getattr(model_config, "pretrained_config", None) - if pretrained_config is None or not is_qwen4_exp(pretrained_config): - return 0 - ple_params = extract_qwen4_exp_ple_cache_params(pretrained_config) - mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( - is_draft=False, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - combined_layer_mask = [ - is_mamba or is_attention for is_mamba, is_attention in zip( - mamba_layer_mask, full_attention_layer_mask) - ] - local_layer_indices, _ = get_pp_layers( - sum(combined_layer_mask), - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - local_ple_layers = sum(layer_id < len(ple_params.ple_layer_mask) - and ple_params.ple_layer_mask[layer_id] - for layer_id in local_layer_indices) - ple_bytes_per_layer = ( - ple_params.short_conv_channels * ple_params.short_conv_state_len * - ple_params.conv_state_dtype.itemsize + - ple_params.ngram_context_len * torch.int64.itemsize) - return local_ple_layers * ple_bytes_per_layer - - -def _estimate_mamba_hybrid_cache_cost( - model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - tokens_per_block: int, - max_seq_len: Optional[int], - num_reserved_dummy_slots: int, - include_explicit_snapshots: bool, - cap_partial_attention_snapshots: bool, - is_draft: bool = False, - use_separate_draft_kv_cache: bool = False, - **kwargs, -) -> Tuple[int, int]: - spec_config = kwargs.get("spec_config") - params, local_mamba_layers, local_attention_layers = ( - _get_local_mamba_cache_layout( - model_config, - mapping, - spec_config=spec_config, - is_draft=is_draft, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - )) - attention_slope = (KVCacheManager.get_cache_size_per_token( - model_config, - mapping, - num_layers=local_attention_layers, - **kwargs, - ) if local_attention_layers > 0 else 0) - state_bytes_per_rank = (local_mamba_layers * - params.get_states_bytes_per_layer(mapping)) - if not is_draft: - state_bytes_per_rank += _qwen4_exp_ple_state_bytes_per_rank( - model_config, - params, - mapping, - spec_config=spec_config, - use_separate_draft_kv_cache=use_separate_draft_kv_cache, - ) - max_resident_sequences = max_batch_size * mapping.pp_size - - if include_explicit_snapshots: - fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( - kv_cache_config, max_seq_len, tokens_per_block) - else: - fixed_rules = 0 - unaligned_fixed_rules = 0 - fixed_state_slots = (max_resident_sequences + num_reserved_dummy_slots + - max_resident_sequences * fixed_rules) - attention_block_bytes = attention_slope * tokens_per_block - - interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) - has_unaligned_periodic_snapshot = (interval is not None - and interval % tokens_per_block != 0) - if cap_partial_attention_snapshots: - # Snapshot alignment is unknown while estimating cache cost. Once a - # snapshot is possible, reserve one retained partial attention page - # per resident lineage. Dummy requests carry no attention capacity. - has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None - partial_attention_slots = (max_resident_sequences - if has_non_live_ssm_capacity else 0) - else: - partial_attention_slots = (max_resident_sequences * - unaligned_fixed_rules) - intercept = (fixed_state_slots * state_bytes_per_rank + - partial_attention_slots * attention_block_bytes) - - if interval is None: - regular_slope = 0 - else: - regular_slope = math.ceil(state_bytes_per_rank / interval) - if (has_unaligned_periodic_snapshot - and not cap_partial_attention_snapshots): - regular_slope += math.ceil(attention_block_bytes / interval) - return attention_slope + regular_slope, intercept - - -class CppMambaHybridCacheManager(KVCacheManager, MambaHybridCacheManager): - """Hybrid cache manager storing mamba states inside the KVCacheManager pool. - - Both KV cache blocks and recurrent state blocks are managed by the unified - C++ KVCacheManager, enabling block reuse / prefix caching across attention - and mamba layers. This compatibility manager remains available through the - manager preference override and legacy disaggregated routing. - - """ - - 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, - # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. - # It's derived from the model's BuildConfig for consistency with the C++ backend. - 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, # this is the full attention layer mask - 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: - # 3 kinds of layers: - # 1) Mamba layers (mamba_layer_mask is True) - # 2) Full attention layers (full_attention_layer_mask is True) - # 3) Not managed layers (both masks are False) - 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) - layer_mask = [ - mamba_layer_mask[i] or full_attention_layer_mask[i] - for i in range(total_layers) - ] - # PP sharding is done across all layers. - # This is called again in the super().__init__, but we want it to run first - # to set up mtp states before the C++ backend is initialized. - self.pp_layers, _ = get_pp_layers( - mamba_num_layers + num_layers, - mapping, - spec_config=spec_config, - layer_mask=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) - self.requests = [] - # Seed externally visible mamba fields before any early return so that - # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work - # on ranks with no local mamba layers. - self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = ( - use_replay_state_update and model_type == "qwen3_next" - and self.local_num_mamba_layers > 0) - self.replay_step_width: Optional[int] = ( - spec_config.tokens_per_gen_step - if spec_config is not None and use_replay_state_update else None) - self.replay_history_size: Optional[int] = ( - max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) - if self.replay_step_width is not None else None) - # Same allocation gate as PythonMambaCacheManager: the rand_seed - # buffer must exist whenever SR can fire, not only on the replay path. - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - # Host-side counter bumped per fresh context-slot assignment; combined - # with the slot index and rank offset to produce reproducible per-slot - # seed values without any torch.randint. - self._seed_request_counter = 0 - self.ssm_state_dtype = mamba_ssm_cache_dtype - # Keep the shared Mamba interface valid on PP ranks that do not own a - # local Mamba layer. - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - - if self.local_num_mamba_layers == 0: - logger.info( - "No local mamba layers for this rank, skipping mamba cache initialization" - ) - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=num_layers, - num_kv_heads=num_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=full_attention_layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - is_draft=is_draft, - **kwargs, - ) - # PP ranks replay the same scheduling decisions, so a rank without - # local Mamba layers must still publish the configured boundaries. - self.kv_cache_config = kv_cache_config - self.linear_attention_metadata = LinearAttentionMetadata() - self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_config.periodic_snapshot_interval - if kv_cache_config.enable_block_reuse else 0) - return - - # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) - tp_size = _mamba_effective_tp_size(mapping) - 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 - self.conv_state_shape = [conv_dim, mamba_d_conv - 1] - self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] - self.conv_state_dtype = mamba_cache_dtype - - # Store GLOBAL (pre-TP-division) mamba params for disagg RnnModelConfig. - # d_inner is computed at the top of __init__ before TP division. - d_inner_global = d_inner # = mamba_head_dim * mamba_num_heads (GLOBAL) - conv_dim_global = d_inner_global + 2 * mamba_n_groups * mamba_d_state - self._rnn_d_state = mamba_d_state - self._rnn_d_conv = mamba_d_conv - self._rnn_num_heads = mamba_num_heads # GLOBAL - self._rnn_n_groups = mamba_n_groups - self._rnn_head_dim = mamba_head_dim - self._rnn_hidden_size = d_inner_global # GLOBAL = head_dim * num_heads - self._rnn_conv_dim_size = conv_dim_global # GLOBAL conv dim - self._rnn_num_layers = mamba_num_layers - - # Conv section layout for section-aware split/concat in TP mismatch. - # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; - # we just need to tell C++ which ordering to use. - self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" - self.ssm_bytes = (math.prod(self.ssm_state_shape) * - self.ssm_state_dtype.itemsize) - self.conv_bytes = (math.prod(self.conv_state_shape) * - self.conv_state_dtype.itemsize) - - total_bytes = self.ssm_bytes + self.conv_bytes - if total_bytes % self.ssm_state_dtype.itemsize != 0: - raise RuntimeError( - f"Total state bytes ({total_bytes}) not divisible by " - f"ssm_state_dtype size ({self.ssm_state_dtype.itemsize})") - if total_bytes % self.conv_state_dtype.itemsize != 0: - raise RuntimeError( - f"Total state bytes ({total_bytes}) not divisible by " - f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") - if self.ssm_bytes % self.conv_state_dtype.itemsize != 0: - raise RuntimeError( - f"SSM state bytes ({self.ssm_bytes}) not divisible by " - f"conv_state_dtype size ({self.conv_state_dtype.itemsize})") - self.linear_attention_metadata = LinearAttentionMetadata() - self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value - self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes - self.linear_attention_metadata.states_snapshot_interval = ( - kv_cache_config.mamba_state_config.periodic_snapshot_interval - if kv_cache_config.enable_block_reuse else 0) - # RNN model params for disagg TP-mismatch split/concat. - conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} - self.linear_attention_metadata.rnn_num_heads = self._rnn_num_heads - self.linear_attention_metadata.rnn_head_dim = self._rnn_head_dim - self.linear_attention_metadata.rnn_d_state = self._rnn_d_state - self.linear_attention_metadata.rnn_d_conv = self._rnn_d_conv - self.linear_attention_metadata.rnn_n_groups = self._rnn_n_groups - self.linear_attention_metadata.rnn_conv_section_layout = conv_section_map.get( - self._rnn_conv_section_layout, 0) - self.linear_attention_metadata.rnn_ssm_bytes = self.ssm_bytes - self.linear_attention_metadata.rnn_ssm_dtype_size = self.ssm_state_dtype.itemsize - self.linear_attention_metadata.rnn_conv_dtype_size = self.conv_state_dtype.itemsize - kv_cache_config = kv_cache_config.model_copy(deep=True) - if kv_cache_config.enable_partial_reuse: - logger.warning( - "Partial reuse is not supported for mamba hybrid models, disabling partial reuse" - ) - kv_cache_config.enable_partial_reuse = False - - # Keep the vector in physical global-layer order. Disabled entries are - # placeholders that are never projected into this manager. - kv_cache_config.max_attention_window = [ - LinearCacheType.RECURRENT_STATES.value - if mamba_layer_mask[layer_idx] else max_seq_len - for layer_idx in range(len(layer_mask)) - ] - - recurrent_states_window = LinearCacheType.RECURRENT_STATES.value - local_windows = { - recurrent_states_window - if mamba_layer_mask[layer_idx] else max_seq_len - for layer_idx in self.pp_layers - } - kwargs["pool_configurations"] = [ - PoolConfiguration( - window_size=window_size, - head_dim=head_dim, - dtype=torch_dtype_to_binding(self.ssm_state_dtype) - if window_size == recurrent_states_window else dtype, - ) for window_size in sorted(local_windows) - ] - - # Normalize num_kv_heads to a per-layer list and zero out mamba - # layer positions: those layers carry SSM/conv state instead of KV - # heads, so the parent KV cache should not allocate KV head storage - # for them. - 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) - - # pass remaining arguments to super class - 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=layer_mask, - is_estimating_kv_cache=is_estimating_kv_cache, - is_draft=is_draft, - linear_attention_metadata=self.linear_attention_metadata, - **kwargs, - ) - - assert self.local_num_mamba_layers > 0, "At least one mamba layer is required" - self.mamba_layer_offsets = {} - for idx, layer_id in enumerate(self.mamba_pp_layers): - self.mamba_layer_offsets[layer_id] = idx - - self.host_block_offsets = torch.zeros([ - self.impl.num_pools, self.max_batch_size, 2, self.max_blocks_per_seq - ], - dtype=torch.int32, - device="cpu") - self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ - self.layer_offsets[self.mamba_pp_layers[0]]][0] - - 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()) - self._row_indices = torch.arange(self.max_batch_size, - dtype=torch.long, - device="cpu") - self._request_id_to_state_index = {} - self._request_id_to_is_dummy = {} - self.kv_cache_config = kv_cache_config - self.is_estimating_kv_cache = is_estimating_kv_cache - - self._setup_states() - self._setup_replay_buffers(spec_config) - if use_replay_state_update and model_type == "qwen3_next": - logger.info_once( - "Configured GDN cached replay commit mode: small-batch fused, " - "large-batch all-layer", - key="gdn_cached_replay_commit_mode_fused", - ) - - @staticmethod - def get_cache_size_per_token( - model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - num_layers: Optional[int] = None, - tokens_per_block: int = 32, - max_seq_len: Optional[int] = None, - **kwargs, - ): - """Affine memory model for the unified hybrid KV pool. - - Returns ``(slope_bytes_per_token, intercept_bytes)``: - - * ``slope`` = attention KV bytes per token plus conservative periodic - Mamba-state and partial-attention-snapshot costs. - * ``intercept`` = live and CUDA-graph dummy Mamba state. - - Memory budget -> max tokens then becomes - ``T = (budget - intercept) // slope`` instead of plain - ``T = budget // bytes_per_token``. - """ - return _estimate_mamba_hybrid_cache_cost( - model_config, - mapping, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - num_reserved_dummy_slots=1, - include_explicit_snapshots=False, - cap_partial_attention_snapshots=False, - **kwargs, - ) - - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return self._use_gdn_cached_replay_all_layer_commit - - def _commit_gdn_cached_replay_history_layers( - self, - attn_metadata: "AttentionMetadata", - num_decodes: int, - ) -> None: - """Synchronously advance every local GDN checkpoint in one launch.""" - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers) - - if (not self._use_gdn_cached_replay_all_layer_commit - or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): - return - if (self.all_ssm_states is None or self.old_x is None - or self.old_B is None or self.old_dt is None - or self.replay_history_size is None): - raise RuntimeError( - "GDN cached replay all-layer commit requires replay state buffers." - ) - - mamba_metadata = attn_metadata.mamba_metadata - if mamba_metadata.replay_num_decodes != num_decodes: - raise RuntimeError( - "GDN replay metadata contains " - f"{mamba_metadata.replay_num_decodes} decode requests, " - f"but state update received {num_decodes}.") - commit_gdn_cached_replay_history_layers( - ssm_states=self.all_ssm_states, - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=mamba_metadata.replay_work_items[:num_decodes], - n_writes=mamba_metadata.replay_n_writes, - history_size=self.replay_history_size, - ) - - def shutdown(self): - # Release tensor views into the pool before the pool memory is freed, - # so their deleters don't see stale pointers. - self.all_ssm_states = None - self.all_conv_states = None - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - super().shutdown() - - def add_dummy_requests( - self, - request_ids: List[int], - # Note that token_nums should be past_kv_len + input_len (without - # spec decoding). The draft tokens will be added in this function, - # so we don't need to take care of it in the caller. When preparing - # token_nums, we should not take the draft tokens into account, so - # don't use the kv_cache_manager.max_seq_len, which includes both - # extra tokens and draft tokens. - 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, - draft_kv_cache_manager: Optional[KVCacheManager] = None, - capture_sampling_params: Optional["SamplingParams"] = 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, - draft_kv_cache_manager=draft_kv_cache_manager, - capture_sampling_params=capture_sampling_params, - ) - if requests: - self.requests.extend(requests) - # Only process the newly added requests, not all of self.requests. - # self.requests may contain stale entries from the previous - # _prepare_resources call (e.g., disagg transfer-pending requests) - # that would exceed max_batch_size and cause out-of-bounds access. - self._setup_state_indices(requests) - return requests - - 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) - - @nvtx_range("hybrid_prepare_resources") - def _prepare_resources(self, scheduled_batch: ScheduledRequests): - self.requests = scheduled_batch.context_requests + \ - scheduled_batch.generation_requests - # Issue all async block onboards; defer the syncTransfers() (refresh_blocks) - # until just before forward so the CPU-side prep (attn metadata, input - # tensor assembly, draft model prep, etc.) can overlap with the in-flight - # cudaMemcpyAsync calls instead of being serialized behind them. - # copy_linear_attention_block returns True iff it actually issued a copy; - # we skip refresh_blocks entirely when nothing was scheduled. - self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( - self.requests) - self._setup_state_indices() - self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - super().prepare_resources(scheduled_batch) - if self.local_num_mamba_layers == 0: - return - self._prepare_resources(scheduled_batch) - - @nvtx_range("hybrid_flush_state_transfers") - def flush_state_transfers(self) -> None: - """Complete any deferred state-block onboards scheduled by - prepare_resources(). Must be called before forward() reads recurrent - state blocks. Cheap no-op when nothing was scheduled. - """ - if getattr(self, "_pending_state_transfers", False): - self.impl.refresh_blocks() - self._pending_state_transfers = False - - @nvtx_range("hybrid_update_mamba_states") - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None, - accepted_leaf_positions: Optional[torch.Tensor] = None): - if self.local_num_mamba_layers == 0: - return - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + - num_gens] - # Dynamic tree supplies explicit tree-node leaf positions. Linear MTP - # passes accepted-token counts and lets the promotion kernel subtract - # one, avoiding a separate GPU subtraction/cast. - if accepted_leaf_positions is None: - accepted_position_source = accepted_tokens - position_source_is_token_count = True - else: - accepted_position_source = accepted_leaf_positions.to(torch.int32) - position_source_is_token_count = False - # Match the API of MambaCacheManager.update_mamba_states: callers - # may pass per-request state slot indices explicitly (e.g. MTP via - # attn_metadata.mamba_metadata.state_indices). Fall back to this - # manager's own slot mapping when not provided. - 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] - - # The accepted SSM/conv promotion is a bandwidth-bound gather->scatter - # into the dtype-reinterpreted, strided C++ pool view. torch.compile - # can't handle the dtype-view mutation (and Inductor chokes on the - # uint8 pool with "XBLOCK too large"), so a dedicated Triton kernel - # writes through the view's real strides (~85% of HBM peak, one launch - # per state). - if self._use_replay_state_update: - # Every GDN layer has finished reading the old checkpoint and - # writing its candidate history. Advance all local checkpoints - # in one launch before PNAT and the active history buffer change. - self._commit_gdn_cached_replay_history_layers( - attn_metadata, num_gens) - assert self._dummy_request_mask is not None - else: - # Legacy: copy the accepted SSM state from the intermediate buffer. - _promote_mamba_state_triton( - self.all_ssm_states, - self.intermediate_ssm_states, - src_state_indices, - accepted_position_source, - state_indices_d, - position_source_is_token_count=position_source_is_token_count) - - # Conv: both paths save all intermediate conv windows, carry over the - # accepted one. The replay path also advances PNAT and the active cache - # buffer in this same launch. - replay_dummy_request_mask = None - if self._use_replay_state_update: - replay_dummy_request_mask = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - _promote_mamba_state_triton( - self.all_conv_states, - self.intermediate_conv_states, - src_state_indices, - accepted_position_source, - state_indices_d, - num_accepted_tokens=accepted_tokens, - position_source_is_token_count=position_source_is_token_count, - replay_pnat=(self.prev_num_accepted_tokens - if self._use_replay_state_update else None), - replay_cache_buf_idx=(self.cache_buf_idx - if self._use_replay_state_update else None), - dummy_request_mask=replay_dummy_request_mask, - replay_step_width=(self.replay_step_width - if self._use_replay_state_update else 0), - replay_history_size=(self.replay_history_size - if self._use_replay_state_update else 0)) - - def get_num_available_tokens(self, - token_num_upper_bound: int, - max_num_draft_tokens: int = 0, - **kwargs) -> int: - # Base bound from attention KV cache pool (the parent's behaviour). - result = super().get_num_available_tokens(token_num_upper_bound, - max_num_draft_tokens, - **kwargs) - # Also bound by the recurrent-state pool capacity: each request needs - # roughly ceil(N / states_snapshot_interval) recurrent-state blocks - # (+1 for the corner case where N is a multiple of tokens_per_block). - # When block reuse is disabled, only one snapshot is needed per - # request, so no additional capping is required here. - interval = (self.linear_attention_metadata.states_snapshot_interval - if self.linear_attention_metadata is not None else 0) - # Attention-only PP ranks keep the interval so every rank publishes - # identical scheduling boundaries, but they have no recurrent-state - # pool whose capacity should constrain their attention KV cache. - if self.local_num_mamba_layers > 0 and interval and interval > 0: - stats = self.impl.get_kv_cache_stats() - rs_free = stats.num_free_blocks_per_window_size.get( - LinearCacheType.RECURRENT_STATES.value, 0) - # Reserve 1 block for the always-allocated last block (corner case - # / final live state) so we don't promise more tokens than the - # pool can actually back at allocation time. - usable_rs_blocks = max(0, rs_free - 1) - rs_token_cap = usable_rs_blocks * interval - result = min(result, rs_token_cap) - return max(result, 0) - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - if request in self.requests: - self.requests.remove(request) - self._request_id_to_state_index.pop(request.py_request_id, None) - self._request_id_to_is_dummy.pop(request.py_request_id, None) - super().free_resources(request, pin_on_release) - - def _setup_state_indices(self, requests=None) -> None: - if self.local_num_mamba_layers == 0: - return - if requests is None: - requests = self.requests - block_indices = [] - for req in requests: - if req.is_context_finished: - next_step = self.get_num_tokens(req) - 1 - elif self.kv_cache_config.enable_block_reuse: - next_step = (req.context_current_position - 1 + - req.context_chunk_size) - else: - next_step = req.prompt_len - 1 - block_indices.append(next_step // self.tokens_per_block) - self.impl.copy_batch_block_offsets( - self.host_block_offsets, [req.py_request_id for req in requests], 1, - 0) - max_blocks = self.blocks_per_window[ - LinearCacheType.RECURRENT_STATES.value][0] - n = len(requests) - self._host_state_indices.zero_() - if n > 0: - # Vectorized gather: replace per-element Python loop with a single - # tensor index op. block_indices is a small Python list so - # torch.tensor() conversion here is O(n) at C level, much cheaper - # than n round-trips through Python indexing. - bi = torch.tensor(block_indices, dtype=torch.long) - rows = self._row_indices[:n] - # host_block_offsets: [num_pools, max_batch_size, 2, max_blocks_per_seq] - values = self.host_block_offsets[self.recurrent_states_pool_index, - rows, 0, bi] - invalid_mask = (values < 0) | (values >= max_blocks) - # The C++ recurrent-state manager uses null page-table entries for - # logical blocks that are not snapshot boundaries. Usually a - # context chunk ends exactly at an allocated snapshot (or at the - # final live-state block), but the scheduler may shorten a chunk - # further when KV capacity is tight. In that case the Mamba - # kernel must keep accumulating into the next allocated snapshot - # or final block, matching copyLinearAttentionBlock(), which also - # skips placeholders when it advances the live state. - for bad_i in invalid_mask.nonzero( - as_tuple=False).flatten().tolist(): - req = requests[bad_i] - if req.is_context_finished: - continue - last_prompt_block = (req.prompt_len - - 1) // self.tokens_per_block - row = self.host_block_offsets[self.recurrent_states_pool_index, - rows[bad_i], 0] - candidates = row[block_indices[bad_i]:last_prompt_block + 1] - valid_candidates = ((candidates >= 0) - & (candidates < max_blocks)).nonzero( - as_tuple=False) - if valid_candidates.numel() > 0: - values[bad_i] = candidates[valid_candidates[0, 0]] - - invalid_mask = (values < 0) | (values >= max_blocks) - if invalid_mask.any(): - bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) - req = requests[bad_i] - value = int(values[bad_i]) - raise RuntimeError( - f"Invalid recurrent state block index {value} " - f"(expected 0 <= index < {max_blocks}) for request {bad_i}, " - f"prompt_len={req.prompt_len}, " - f"is_context_finished={req.is_context_finished}, " - f"context_current_position={req.context_current_position}, " - f"prepopulated_token_num={req.prepopulated_prompt_len}, " - f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " - f"block_index for next step is {block_indices[bad_i]}, " - "\nblock_ids=" - f"{self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" - ) - self._host_state_indices[:n] = values - - self.cuda_state_indices.copy_(self._host_state_indices, - non_blocking=True) - is_dummy = [req.is_dummy for req in requests] - self._refresh_dummy_request_mask(is_dummy) - - # Build request_id → pool block offset mapping so that - # get_state_indices can return indices in arbitrary request order. - # Bulk tolist avoids a per-request tensor-index + .item() round-trip. - state_values = self._host_state_indices[:n].tolist() - for req, value, dummy in zip(requests, state_values, is_dummy): - self._request_id_to_state_index[req.py_request_id] = value - self._request_id_to_is_dummy[req.py_request_id] = dummy - - def get_state_indices(self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None) -> list: - if self.local_num_mamba_layers == 0: - # Mamba metadata is prepared on every PP rank even when this rank - # owns only attention layers. No local kernel consumes these - # indices, so avoid consulting state that is intentionally absent. - return [0] * len(request_ids) if request_ids is not None else [] - if request_ids is not None: - # Return indices in the order of the caller's request_ids, - # not the internal self.requests order. This is critical when - # the batch is reordered after prepare_resources (e.g. disagg - # serving sorts generation_requests by py_batch_idx). - indices = [ - self._request_id_to_state_index[rid] for rid in request_ids - ] - if is_padding is None: - is_padding = [False] * len(request_ids) - assert len(request_ids) == len(is_padding) - is_dummy = [ - self._request_id_to_is_dummy.get(rid, False) or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - return self.cuda_state_indices - - def _setup_states(self) -> None: - # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) - pool: torch.Tensor = self.impl.get_recurrent_states_pool().view( - torch.uint8).reshape(self.local_num_mamba_layers, -1, - self.ssm_bytes + self.conv_bytes) - num_blocks_in_pool = pool.shape[1] - self.all_ssm_states = pool[:, :, :self.ssm_bytes].view( - self.ssm_state_dtype).view( - [self.local_num_mamba_layers, num_blocks_in_pool] + - self.ssm_state_shape) - self.all_conv_states = pool[:, :, self.ssm_bytes:self.ssm_bytes + - self.conv_bytes].view( - self.conv_state_dtype).view([ - self.local_num_mamba_layers, - num_blocks_in_pool - ] + self.conv_state_shape) - self.all_ssm_states.zero_() - self.all_conv_states.zero_() - - def _setup_replay_buffers(self, spec_config) -> None: - cache_size = self.all_ssm_states.shape[1] - device = self.all_ssm_states.device - if not self._allocate_pool_replay_buffers(spec_config, cache_size, - device): - return - - self._dummy_request_mask = torch.zeros(self.max_batch_size, - dtype=torch.bool, - device=device) - self._dummy_request_mask_host = torch.zeros( - self.max_batch_size, - dtype=torch.bool, - pin_memory=prefer_pinned(), - ) - - -class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): - """Hybrid Mamba cache manager backed by KVCacheManagerV2. - - Attention KV pages and Mamba recurrent-state pages are both owned by the - Python V2 cache manager. Mamba layers are represented as V2 SSM layers, - while this wrapper exposes the state tensors and slot indices expected by - the PyTorch Mamba kernels. - """ - - _supports_additional_snapshot_offsets = True - # Recurrent-state snapshots use a specialized commit/history protocol, so - # keep main-like reuse endpoints and the existing unpaired draft path. - _supports_reuse_match_backoff = False - - # Qwen4-Exp PLE state is opt-in. These class-level defaults keep every other - # model — and any partially-constructed instance that sets only the fields it - # needs — out of the PLE branches without carrying PLE geometry. - # `_init_qwen4_exp_ple_geometry` rebinds all of them per instance. - _ple_layer_ids: list[int] = [] - _ple_conv_state_shape: list[int] = [] - _ple_ngram_context_shape: list[int] = [] - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - max_seq_len: int, - max_batch_size: int, - mapping: Mapping, - dtype: DataType = DataType.HALF, - spec_config: Optional["DecodingBaseConfig"] = None, - layer_mask: Optional[List[bool]] = None, - is_estimating_kv_cache: bool = False, - is_draft: bool = False, - use_replay_state_update: bool = False, - mamba_ssm_stochastic_rounding: bool = False, - conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", - kda_replay_num_spec: Optional[int] = None, - qwen4_exp_ple_cache_params: Optional["Qwen4ExpPLECacheParams"] = None, - **kwargs, - ) -> None: - if conv_state_layout not in ("x_b_c", "q_k_v"): - raise ValueError( - f"Unsupported convolution state layout: {conv_state_layout!r}") - if "model_type" in kwargs: - # The V1 managers select the conv-state layout by model_type; this - # class takes it explicitly. Silently absorbing model_type here - # means a caller's layout request would be dropped on the floor. - raise TypeError( - "MambaHybridCacheManagerV2 does not accept 'model_type' " - f"(got {kwargs['model_type']!r}); pass " - "conv_state_layout='x_b_c' or 'q_k_v' instead") - - self._kda_replay_num_spec = kda_replay_num_spec - self._use_kda_replay_update = kda_replay_num_spec is not None - if self._use_kda_replay_update: - if use_replay_state_update: - raise ValueError( - "kda_replay_num_spec and use_replay_state_update are " - "mutually exclusive") - if spec_config is None: - raise ValueError( - "KDA replay caches require speculative decoding") - expected_num_spec = spec_config.tokens_per_gen_step - 1 - if kda_replay_num_spec != expected_num_spec: - raise ValueError( - "KDA replay cache width " - f"({kda_replay_num_spec}) must match the draft length " - f"({expected_num_spec})") - if conv_state_layout != "q_k_v": - raise ValueError( - "KDA replay requires conv_state_layout='q_k_v'") - self._record_kda_replay_in_update_resources = ( - self.use_kda_replay_update - and getattr(spec_config, "decoding_type", None) == "NGram") - - total_layers = len(mamba_layer_mask) - if layer_mask is None: - full_attention_layer_mask = [False] * total_layers - elif len(layer_mask) != total_layers: - raise ValueError( - f"layer_mask length ({len(layer_mask)}) must match " - f"mamba_layer_mask length ({total_layers})") - else: - full_attention_layer_mask = list(layer_mask) - - combined_layer_mask = [ - mamba_layer_mask[i] or full_attention_layer_mask[i] - for i in range(total_layers) - ] - - self._mamba_layer_mask = list(mamba_layer_mask) - self._use_replay_state_update = use_replay_state_update - self._use_gdn_cached_replay_all_layer_commit = (use_replay_state_update - and conv_state_layout - == "q_k_v") - self._gdn_cached_replay_state_descriptors: Optional[torch.Tensor] = None - self._gdn_cached_replay_state_strides: Optional[Tuple[int, int, - int]] = None - self.replay_step_width: Optional[int] = ( - spec_config.tokens_per_gen_step - if spec_config is not None and use_replay_state_update else None) - self.replay_history_size: Optional[int] = ( - max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) - if self.replay_step_width is not None else None) - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) - self._seed_request_counter = 0 - self._recurrent_evicted_blocks_total = 0 - self._recurrent_onboarded_blocks_total = 0 - self._recurrent_dropped_blocks_total = 0 - self._recurrent_status_logged = False - # Branch points keyed by request id. prepare_expect_snapshot_points() - # overwrites the per-request list on every scheduler pass, so the point - # has to be re-merged from here rather than stored only on the request. - self._branch_snapshot_points: Dict[int, int] = {} - self._snapshot_pruned_tokens_total = 0 - self._page_pruned_tokens_total = 0 - self._branch_snapshots_taken_total = 0 - self._branch_snapshots_skipped_total: Dict[str, int] = {} - num_cuda_graph_padding_dummy_slots = ( - _get_num_cuda_graph_padding_dummy_slots(spec_config, - max_batch_size)) - self._num_reserved_dummy_slots = (num_cuda_graph_padding_dummy_slots + - int(mapping.enable_attention_dp)) - # PLE uses the same request slots and snapshots as the GDN state. - self._init_qwen4_exp_ple_geometry(qwen4_exp_ple_cache_params, - total_layers, mamba_cache_dtype) - self.ssm_state_dtype = (mamba_ssm_cache_dtype if mamba_ssm_cache_dtype - is not None else mamba_cache_dtype) - self.conv_state_dtype = mamba_cache_dtype - - self.pp_layers, _ = get_pp_layers( - mamba_num_layers + num_layers, - mapping, - spec_config=spec_config, - layer_mask=combined_layer_mask, - ) - self.mamba_pp_layers = [ - layer_idx for layer_idx in self.pp_layers - if mamba_layer_mask[layer_idx] - ] - self.local_num_mamba_layers = len(self.mamba_pp_layers) - self._use_gdn_cached_replay_all_layer_commit = ( - self._use_gdn_cached_replay_all_layer_commit - and self.local_num_mamba_layers > 0) - - if self.local_num_mamba_layers > 0: - tp_size = _mamba_effective_tp_size(mapping) - d_inner = mamba_head_dim * mamba_num_heads - grouped_state_dim = mamba_n_groups * mamba_d_state - conv_dim = d_inner + 2 * grouped_state_dim - nheads = mamba_num_heads - assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" - assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" - if kwargs.get("is_disagg", - False) and grouped_state_dim % tp_size != 0: - raise ValueError( - "Disaggregated Mamba transfer requires each convolution " - "state section to be divisible by tp_size") - if use_replay_state_update: - assert mamba_n_groups % tp_size == 0, \ - "replay state update requires mamba_n_groups divisible by tp_size" - self._n_groups_per_rank = mamba_n_groups // tp_size - d_inner_local = d_inner // tp_size - grouped_state_dim_local = grouped_state_dim // tp_size - conv_dim = conv_dim // tp_size - nheads = nheads // tp_size - self.conv_state_shape = [conv_dim, mamba_d_conv - 1] - self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] - # TP-mismatch disaggregated transfers must split the flat - # convolution state at its true semantic boundaries. Mamba2 stores - # [x | B | C], while GDN stores [Q | K | V]. The large section is - # therefore first for Mamba2 and last for GDN. - if conv_state_layout == "x_b_c": - self.conv_section_dims = [ - d_inner_local, - grouped_state_dim_local, - grouped_state_dim_local, - ] - else: - self.conv_section_dims = [ - grouped_state_dim_local, - grouped_state_dim_local, - d_inner_local, - ] - self.ssm_bytes = (math.prod(self.ssm_state_shape) * - self.ssm_state_dtype.itemsize) - self.conv_bytes = (math.prod(self.conv_state_shape) * - self.conv_state_dtype.itemsize) - else: - logger.info( - "No local mamba layers for this rank, skipping mamba state views" - ) - self._n_groups_per_rank = 0 - self.conv_state_shape = [] - self.ssm_state_shape = [] - self.conv_section_dims = [] - self.ssm_bytes = 0 - self.conv_bytes = 0 - - 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: - block_reuse_config = kv_cache_config.block_reuse_config - block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) - if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: - # SSM reuse is valid only at explicit snapshot boundaries. - kv_cache_config.block_reuse_config = block_reuse_config.model_copy( - update={"policy": BlockReusePolicy.PER_REQUEST.value}) - self.kv_cache_config = kv_cache_config - - super().__init__( - kv_cache_config, - kv_cache_type, - num_layers=mamba_num_layers + num_layers, - num_kv_heads=per_layer_kv_heads, - head_dim=head_dim, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - max_batch_size=max_batch_size, - mapping=mapping, - dtype=dtype, - spec_config=spec_config, - layer_mask=combined_layer_mask, - is_draft=is_draft, - is_estimating_kv_cache=is_estimating_kv_cache, - num_reserved_index_slots=self._num_reserved_dummy_slots, - **kwargs, - ) - - self.mamba_layer_offsets = { - layer_id: offset - for offset, layer_id in enumerate(self.mamba_pp_layers) - } - self._request_id_to_state_index = {} - self._request_id_to_is_dummy = {} - - state_index_capacity = (self.max_batch_size + - self._num_reserved_dummy_slots) - self.cuda_state_indices = torch.zeros([state_index_capacity], - dtype=torch.int32, - device="cuda") - self._host_state_indices = torch.zeros([state_index_capacity], - dtype=torch.int32, - pin_memory=prefer_pinned()) - - if self.local_num_mamba_layers > 0: - first_mamba_local_layer = self.layer_offsets[ - self.mamba_pp_layers[0]] - self.ssm_layer_group_id = self.impl.get_layer_group_id( - LayerId(first_mamba_local_layer)) - self._ssm_page_index_scale = self.impl.get_page_index_scale( - LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) - num_ssm_pages = self.impl.get_page_index_upper_bound( - LayerId(first_mamba_local_layer), MambaRole.SSM_STATE) - num_ssm_slots = ((num_ssm_pages + self._ssm_page_index_scale - 1) // - self._ssm_page_index_scale) - required_live_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - if num_ssm_slots < required_live_slots: - KVCacheManagerV2.shutdown(self) - raise ValueError( - "The V2 Mamba state pool has only " - f"{num_ssm_slots} slots but needs at least " - f"{required_live_slots} live/dummy slots. Increase the " - "KV cache budget or allocate a larger Mamba pool_ratio.") - self._setup_states() - # PLE pools use exactly the state-slot capacity resolved above. - self._setup_ple_states(num_ssm_slots) - self._setup_replay_buffers(spec_config) - if self._use_gdn_cached_replay_all_layer_commit: - state_layout = ("affine" - if self._gdn_cached_replay_state_strides - is not None else "indirect") - logger.info_once( - "Configured GDN cached replay commit mode for V2: " - "small-batch fused, large-batch all-layer; " - f"state layout: {state_layout}", - key="gdn_cached_replay_v2_commit_mode_fused", - ) - else: - self.ssm_layer_group_id = None - self._ssm_page_index_scale = 1 - self.all_ssm_states = [] - self.all_conv_states = [] - self._setup_replay_buffers(spec_config) - - def _init_qwen4_exp_ple_geometry( - self, - params: Optional["Qwen4ExpPLECacheParams"], - total_layers: int, - default_state_dtype: torch.dtype, - ) -> None: - """Validate and record the Qwen4-Exp PLE pool geometry. - - A no-op for every other model: the fields stay empty and - `_ple_buffer_configs` then contributes nothing. - """ - self._ple_layer_ids: list[int] = [] - self._ple_conv_state_shape: list[int] = [] - self._ple_ngram_context_shape: list[int] = [] - self._ple_conv_state_dtype = default_state_dtype - self._ple_conv_states: dict[int, torch.Tensor] = {} - self._ple_ngram_contexts: dict[int, torch.Tensor] = {} - if params is None: - return - if len(params.ple_layer_mask) != total_layers: - raise ValueError( - "PLE layer mask length must match the hybrid layer mask: " - f"got {len(params.ple_layer_mask)}, expected {total_layers}") - self._ple_layer_ids = [ - layer_id for layer_id, active in enumerate(params.ple_layer_mask) - if active - ] - if len(self._ple_layer_ids) != params.num_ple_layers: - raise ValueError( - "PLE layer mask count does not match num_ple_layers: " - f"got {len(self._ple_layer_ids)}, expected " - f"{params.num_ple_layers}") - if any(not self._mamba_layer_mask[layer_id] - for layer_id in self._ple_layer_ids): - raise ValueError("PLE lifecycle state must belong to Mamba layers") - if (params.short_conv_channels <= 0 or params.short_conv_state_len <= 0 - or params.ngram_context_len <= 0): - raise ValueError("PLE recurrent-state dimensions must be positive") - self._ple_conv_state_shape = [ - params.short_conv_channels, - params.short_conv_state_len, - ] - self._ple_ngram_context_shape = [params.ngram_context_len] - self._ple_conv_state_dtype = params.conv_state_dtype - - def _setup_ple_states(self, num_state_slots: int) -> None: - """Bind each local PLE layer to V2-managed recurrent-state views.""" - if not self._ple_layer_ids: - return - for layer_id in self._ple_layer_ids: - local_layer_idx = self.layer_offsets.get(layer_id) - if local_layer_idx is None: - continue - conv_state = self._get_state_buffer( - local_layer_idx, - MambaRole.PLE_CONV_STATE, - self._ple_conv_state_dtype, - self._ple_conv_state_shape, - ) - ngram_context = self._get_state_buffer( - local_layer_idx, - MambaRole.PLE_NGRAM_CONTEXT, - dtype=torch.long, - state_shape=self._ple_ngram_context_shape, - ) - if (conv_state.shape[0] != num_state_slots - or ngram_context.shape[0] != num_state_slots): - raise RuntimeError( - "PLE and GDN lifecycle buffers must have the same number " - f"of slots: layer={layer_id}, GDN={num_state_slots}, " - f"conv={conv_state.shape[0]}, ngram={ngram_context.shape[0]}" - ) - self._ple_conv_states[layer_id] = conv_state - self._ple_ngram_contexts[layer_id] = ngram_context - logger.info( - "PLE state views bound to V2 lifecycle buffers for local layers " - f"{sorted(self._ple_conv_states)}: conv " - f"[{num_state_slots}, {', '.join(map(str, self._ple_conv_state_shape))}] " - f"({self._ple_conv_state_dtype}), n-gram context " - f"[{num_state_slots}, {self._ple_ngram_context_shape[0]}] (int64)") - - def ple_layer_cache( - self, - layer_idx: int) -> Optional[tuple[torch.Tensor, torch.Tensor]]: - """Return persistent ``(conv_state, ngram_context)`` pools for a layer.""" - conv = self._ple_conv_states.get(layer_idx) - ngram = self._ple_ngram_contexts.get(layer_idx) - if conv is None or ngram is None: - return None - return conv, ngram - - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) - - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size: int) -> None: - if not self.use_kda_replay_update: - super()._setup_mtp_intermediate_states(spec_config, max_batch_size) - return - - self.spec_config = spec_config - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - if self.local_num_mamba_layers == 0: - return - self.intermediate_state_indices = torch.arange(max_batch_size, - dtype=torch.int32, - device="cuda") - - def _allocate_pool_replay_buffers( - self, - spec_config, - cache_size: int, - device: Optional[torch.device], - ) -> bool: - allocated = super()._allocate_pool_replay_buffers( - spec_config, cache_size, device) - self.kda_conv_q = None - self.kda_conv_k = None - self.kda_conv_v = None - self.kda_qkg_cache = None - self.kda_v_cache = None - self.kda_beta_cache = None - - if (not self.use_kda_replay_update or self.local_num_mamba_layers == 0): - return allocated - assert spec_config is not None - assert device is not None - if (len(self.conv_section_dims) != 3 - or len(set(self.conv_section_dims)) != 1): - raise ValueError("KDA replay caches require equal [Q | K | V] " - "convolution-state sections") - - num_spec = self._kda_replay_num_spec - assert num_spec is not None - section_dim = self.conv_section_dims[0] - committed_window = self.conv_state_shape[1] - extended_window = committed_window + num_spec - - def allocate_dim_contiguous_conv_cache() -> torch.Tensor: - return torch.zeros( - self.local_num_mamba_layers, - cache_size, - extended_window, - section_dim, - dtype=torch.float32, - device=device, - ).transpose(-1, -2) - - self.prev_num_accepted_tokens = torch.zeros(cache_size, - dtype=torch.int32, - device=device) - self.kda_conv_q = allocate_dim_contiguous_conv_cache() - self.kda_conv_k = allocate_dim_contiguous_conv_cache() - self.kda_conv_v = allocate_dim_contiguous_conv_cache() - self.kda_qkg_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - 3, - section_dim, - dtype=torch.float32, - device=device, - ) - self.kda_v_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - section_dim, - dtype=torch.float32, - device=device, - ) - self.kda_beta_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - self.ssm_state_shape[0], - dtype=torch.float32, - device=device, - ) - logger.info("Mamba Cache (kda-replay) is allocated for " - f"{cache_size} state slots") - return True - - def _commit_gdn_cached_replay_history_layers( - self, - attn_metadata: "AttentionMetadata", - num_decodes: int, - ) -> None: - """Advance all V2 GDN checkpoints through their state views.""" - from tensorrt_llm._torch.modules.fla.cached_replay import ( - CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, - commit_gdn_cached_replay_history_layers) - - if (not getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) - or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE): - return - if (not self.all_ssm_states - or (self._gdn_cached_replay_state_descriptors is None - and self._gdn_cached_replay_state_strides is None) - or self.old_x is None or self.old_B is None - or self.old_dt is None or self.replay_history_size is None): - raise RuntimeError( - "GDN cached replay all-layer commit requires V2 replay state buffers." - ) - - mamba_metadata = attn_metadata.mamba_metadata - if mamba_metadata.replay_num_decodes != num_decodes: - raise RuntimeError( - "GDN replay metadata contains " - f"{mamba_metadata.replay_num_decodes} decode requests, " - f"but state update received {num_decodes}.") - state_strides = self._gdn_cached_replay_state_strides - commit_gdn_cached_replay_history_layers( - ssm_states=self.all_ssm_states[0], - ssm_state_descriptors=(self._gdn_cached_replay_state_descriptors), - ssm_state_num_layers=(state_strides[0] - if state_strides is not None else None), - ssm_state_layer_stride=(state_strides[1] - if state_strides is not None else None), - ssm_state_slot_stride=(state_strides[2] - if state_strides is not None else None), - old_u=self.old_x, - old_k=self.old_B, - old_G=self.old_dt, - replay_work_items=mamba_metadata.replay_work_items[:num_decodes], - n_writes=mamba_metadata.replay_n_writes, - history_size=self.replay_history_size, - ) - - def _create_kv_cache( - self, - request_id: int, - lora_task_id: Optional[int], - input_tokens: Optional[Sequence[TokenIdExt]], - *, - cache_salt: Optional[str] = None, - is_dummy: bool = False, - enable_request_stats: bool = False, - expected_prompt_length: Optional[int] = None, - ) -> Optional[_KVCache]: - kv_cache = super()._create_kv_cache( - request_id, - lora_task_id, - input_tokens, - cache_salt=cache_salt, - is_dummy=is_dummy, - enable_request_stats=enable_request_stats, - expected_prompt_length=expected_prompt_length, - ) - if (self.mapping.rank == 0 and kv_cache is not None - and input_tokens is not None and not is_dummy - and self.local_num_mamba_layers > 0): - # Prefix lookup excludes the final prompt token because it must be - # recomputed by prefill. Restore it for the request-length metric. - request_total_tokens = len(input_tokens) + 1 - prefix_reuse_diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) - logger.debug( - f"[MambaHybridCacheManagerV2] prefix reuse rank={self.mapping.rank} " - f"request_id={request_id} " - f"request_total_tokens={request_total_tokens} " - f"longest_attention_match_tokens=" - f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_hybrid_pruning()} " - f"content_divergence_tokens=" - f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_pruning()} " - f"latest_recurrent_snapshot_tokens=" - f"{kv_cache.num_committed_tokens}") - return kv_cache - - def _skip_branch_snapshot(self, reason: str) -> None: - self._branch_snapshots_skipped_total[reason] = ( - self._branch_snapshots_skipped_total.get(reason, 0) + 1) - - def _format_branch_snapshot_counters(self) -> str: - """Attribute reuse loss between missing snapshots and evicted pages. - - A flat cache hit rate means something different depending on whether - branch points were never found or were found and not reused, so the - skip reasons are reported alongside the totals. - """ - skipped = self._branch_snapshots_skipped_total - return (f"snapshot_pruned_tokens={self._snapshot_pruned_tokens_total} " - f"page_pruned_tokens={self._page_pruned_tokens_total} " - f"branch_snapshots_taken={self._branch_snapshots_taken_total} " - f"branch_snapshots_skipped={sum(skipped.values())} " - f"branch_snapshots_skipped_by_reason=" - f"{dict(sorted(skipped.items()))}") - - def _record_branch_snapshot_point(self, req: LlmRequest, kv_cache: _KVCache, - num_lookup_tokens: Optional[int]) -> None: - """Snapshot where this request leaves the reuse tree, for its siblings.""" - if not self.kv_cache_config.mamba_state_config.enable_branch_snapshot: - return - if req.is_dummy_request: - return - if num_lookup_tokens is None: - self._skip_branch_snapshot("no_fresh_match") - # A fresh lookup can record a point before resume() fails. On the - # next attempt, reapply that point against the cache's authoritative - # reuse depth: the request cursor is not rewound until preparation - # succeeds and may still describe a deeper, abandoned attempt. - if req.py_request_id in self._branch_snapshot_points: - self._apply_branch_snapshot_point( - req, context_current_position=kv_cache.num_committed_tokens) - return - - diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) - divergence = diagnostics._get_num_reusable_tokens_before_pruning() - reused = kv_cache.num_committed_tokens - if self.local_num_mamba_layers > 0: - hybrid_depth = ( - diagnostics._get_num_reusable_tokens_before_hybrid_pruning()) - # Loss attributable to a missing recurrent snapshot, versus loss - # attributable to attention pages having been evicted. These call - # for different fixes, so they are counted apart. - self._snapshot_pruned_tokens_total += max(0, hybrid_depth - reused) - self._page_pruned_tokens_total += max(0, divergence - hybrid_depth) - - # The whole lookup range matched, so there is no fork here. - if divergence >= num_lookup_tokens: - self._skip_branch_snapshot("no_divergence") - return - # Align down: tokens past a block boundary come from a partial-match - # child rather than a confirmed shared prefix, and an unaligned commit - # takes the partial-block snapshot path. - point = (divergence // self.tokens_per_block) * self.tokens_per_block - if point == 0: - self._skip_branch_snapshot("aligned_to_zero") - return - # Reuse already reached the fork, so a snapshot there exists. - if point <= reused: - self._skip_branch_snapshot("already_reused") - return - # A branch point at or beyond the prompt end cannot improve sibling - # reuse. - if point >= req.prompt_len: - self._skip_branch_snapshot("at_prompt_end") - return - - self._branch_snapshot_points[req.py_request_id] = point - self._apply_branch_snapshot_point(req, context_current_position=reused) - - def prepare_expect_snapshot_points(self, - requests: List[LlmRequest]) -> None: - super().prepare_expect_snapshot_points(requests) - if (not self.enable_block_reuse or not self.kv_cache_config. - mamba_state_config.enable_branch_snapshot): - return - for request in requests: - self._apply_branch_snapshot_point(request) - - def _apply_branch_snapshot_point( - self, - request: LlmRequest, - context_current_position: Optional[int] = None, - ) -> None: - points = set(request.expect_snapshot_points) - # Keep the flag usable without any configured snapshot placement. If a - # placement is configured, preserve it: adding prompt_len can replace a - # nearby recurrent checkpoint in the same radix-tree block because a - # block stores one recurrent page per life cycle. - if not points: - points.add(request.prompt_len) - point = self._branch_snapshot_points.get(request.py_request_id) - current_position = (request.context_current_position - if context_current_position is None else - context_current_position) - if point is not None and point > current_position: - points.add(point) - request.expect_snapshot_points = sorted(points) - - def get_iteration_stats(self) -> Optional[KVCacheV2IterationStatsReport]: - """Log recurrent-cache movement; this is KDA state for Kimi K3.""" - report = super().get_iteration_stats() - if report is None: - return None - - pool_group_ids = sorted({ - pool_group_id - for pool_group_id, _, kind in - self._stats_life_cycle_metadata().values() if kind == "ssm" - }) - if not pool_group_ids: - return report - pool_group_reports = [ - report.by_pool_group[pool_group_id] - for pool_group_id in pool_group_ids - ] - - stats = [ - pool_group_report.stats for pool_group_report in pool_group_reports - ] - evicted_blocks = sum(stat.iter_offload_blocks for stat in stats) - evicted_bytes = sum(stat.iter_offload_bytes for stat in stats) - onboarded_blocks = sum(stat.iter_onboard_blocks for stat in stats) - onboarded_bytes = sum(stat.iter_onboard_bytes for stat in stats) - dropped_blocks = sum(stat.iter_host_dropped_blocks for stat in stats) - dropped_bytes = sum(stat.iter_host_dropped_bytes for stat in stats) - has_movement = bool(evicted_blocks or onboarded_blocks - or dropped_blocks) - if has_movement: - self._recurrent_evicted_blocks_total += evicted_blocks - self._recurrent_onboarded_blocks_total += onboarded_blocks - self._recurrent_dropped_blocks_total += dropped_blocks - if (self.mapping.rank == 0 - and (has_movement or not self._recurrent_status_logged)): - logger.debug( - f"[MambaHybridCacheManagerV2] recurrent cache status " - f"rank={self.mapping.rank} pool_group_ids={pool_group_ids} " - f"evicted_recurrent_blocks={evicted_blocks} " - f"evicted_recurrent_bytes={evicted_bytes} " - f"onboarded_recurrent_blocks={onboarded_blocks} " - f"onboarded_recurrent_bytes={onboarded_bytes} " - f"dropped_recurrent_blocks={dropped_blocks} " - f"dropped_recurrent_bytes={dropped_bytes} " - f"total_evicted_recurrent_blocks=" - f"{self._recurrent_evicted_blocks_total} " - f"total_onboarded_recurrent_blocks=" - f"{self._recurrent_onboarded_blocks_total} " - f"total_dropped_recurrent_blocks=" - f"{self._recurrent_dropped_blocks_total} " - f"gpu_used_recurrent_blocks=" - f"{sum(stat.primary_used_num_blocks for stat in stats)} " - f"gpu_free_recurrent_blocks=" - f"{sum(stat.primary_free_num_blocks for stat in stats)} " - f"gpu_evictable_recurrent_blocks=" - f"{sum(stat.primary_evictable_num_blocks for stat in stats)} " - f"host_used_recurrent_blocks=" - f"{sum(stat.secondary_used_num_blocks for stat in stats)} " - f"host_free_recurrent_blocks=" - f"{sum(stat.secondary_free_num_blocks for stat in stats)} " - f"{self._format_branch_snapshot_counters()}") - self._recurrent_status_logged = True - return report - - @staticmethod - def get_cache_size_per_token(model_config, - mapping: Mapping, - *, - max_batch_size: int, - kv_cache_config: KvCacheConfig, - num_layers: Optional[int] = None, - tokens_per_block: int = 32, - max_seq_len: Optional[int] = None, - **kwargs): - spec_config = kwargs.get("spec_config") - num_reserved_dummy_slots = (_get_num_cuda_graph_padding_dummy_slots( - spec_config, max_batch_size) + int(mapping.enable_attention_dp)) - return _estimate_mamba_hybrid_cache_cost( - model_config, - mapping, - max_batch_size=max_batch_size, - kv_cache_config=kv_cache_config, - tokens_per_block=tokens_per_block, - max_seq_len=max_seq_len, - num_reserved_dummy_slots=num_reserved_dummy_slots, - include_explicit_snapshots=True, - cap_partial_attention_snapshots=True, - **kwargs, - ) - - def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: - return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] - - def _get_pool_roles(self, - pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: - layer_id = int(self.impl.layer_grouping[pool_id][0]) - if self._is_local_mamba_layer(layer_id): - return MambaRole.SSM_STATE, None - return super()._get_pool_roles(pool_id) - - def _max_resident_sequences(self) -> int: - return self.max_batch_size * self.mapping.pp_size - - def _mamba_state_bytes_per_slot(self) -> int: - base_bytes = self.local_num_mamba_layers * (self.ssm_bytes + - self.conv_bytes) - local_ple_layers = sum(layer_id in self.pp_layers - for layer_id in self._ple_layer_ids) - if local_ple_layers == 0: - return base_bytes - ple_bytes = ( - math.prod(self._ple_conv_state_shape) * - self._ple_conv_state_dtype.itemsize + - math.prod(self._ple_ngram_context_shape) * torch.int64.itemsize) - return base_bytes + local_ple_layers * ple_bytes - - def _ple_buffer_configs(self) -> list[BufferConfig]: - if not self._ple_layer_ids: - return [] - return [ - BufferConfig( - role=MambaRole.PLE_NGRAM_CONTEXT, - size=(math.prod(self._ple_ngram_context_shape) * - torch.int64.itemsize), - ), - BufferConfig( - role=MambaRole.PLE_CONV_STATE, - size=(math.prod(self._ple_conv_state_shape) * - self._ple_conv_state_dtype.itemsize), - ), - ] - - def _num_ssm_snapshots_for_capacity( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> int: - if capacity <= 0 or not kv_cache_config.enable_block_reuse: - return 0 - - fixed_rules, _ = _mamba_snapshot_rule_counts(kv_cache_config, - self.max_seq_len, - self.tokens_per_block) - interval = _mamba_regular_snapshot_interval(kv_cache_config, - self.max_seq_len) - regular_snapshots = capacity // interval if interval is not None else 0 - return (self._max_resident_sequences() * fixed_rules + - regular_snapshots) - - def _num_ssm_states_per_typical_request( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> int: - fixed_rules, _ = _mamba_snapshot_rule_counts( - kv_cache_config, - capacity, - self.tokens_per_block, - ) - # Additional snapshots are stable boundaries that must remain alive. - # Periodic snapshots are evictable cache entries and therefore do not - # increase the guaranteed state count represented by BatchDesc. - return 1 + fixed_rules - - def _typical_request_descs( - self, - capacity: int, - kv_cache_config: KvCacheConfig, - ) -> List[KVCacheDesc]: - """Model one request with one descriptor per live SSM state.""" - num_states = self._num_ssm_states_per_typical_request( - capacity, kv_cache_config) - capacity_per_state, capacity_remainder = divmod(capacity, num_states) - capacities = [ - capacity_per_state + int(i < capacity_remainder) - for i in range(num_states) - ] - return [ - KVCacheDesc( - capacity=state_capacity, - history_length=max(0, state_capacity - 1), - ) for state_capacity in capacities - ] - - def _get_typical_request_capacity( - self, - kv_cache_config: KvCacheConfig, - ) -> int: - if kv_cache_config.avg_seq_len is not None: - return kv_cache_config.avg_seq_len - - fallback_capacity = max(1, self.max_seq_len // 2) - logger.warning( - "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " - "model using KV cache manager V2. Falling back to " - f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " - "'kv_cache_config.avg_seq_len' in the YAML configuration to the " - "workload's average total sequence length for an accurate KV/SSM " - "pool ratio.") - return fallback_capacity - - def _get_quota_from_max_tokens(self, max_tokens: int) -> int: - attention_quota = super()._get_quota_from_max_tokens(max_tokens) - num_request_lineages = self._max_resident_sequences() - snapshot_slots = self._num_ssm_snapshots_for_capacity( - max_tokens, self.kv_cache_config) - state_slots = (num_request_lineages + self._num_reserved_dummy_slots + - snapshot_slots) - state_quota = state_slots * self._mamba_state_bytes_per_slot() - # Once the plan contains any non-live SSM capacity, reserve one partial - # attention page per request lineage. This remains conservative when - # the plan contains fewer than one non-live slot per lineage. - extra_attention_quota = (num_request_lineages * - self._attention_cache_bytes_per_token() * - self.tokens_per_block - if snapshot_slots > 0 else 0) - return attention_quota + state_quota + extra_attention_quota - - def _get_max_tokens_from_quota(self, quota: int) -> float: - if self._get_quota_from_max_tokens(0) > quota: - return 0 - - low = 0 - high = 1 - while self._get_quota_from_max_tokens(high) <= quota: - low = high - high *= 2 - if high >= 1 << 62: - return float("inf") - - while low + 1 < high: - mid = (low + high) // 2 - if self._get_quota_from_max_tokens(mid) <= quota: - low = mid - else: - high = mid - return low - - def _minimum_live_gpu_quota(self) -> int: - """Return the minimum quota for live states and one attention page.""" - attention_block_quota = (self._attention_cache_bytes_per_token() * - self.tokens_per_block) - num_state_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - state_quota = num_state_slots * self._mamba_state_bytes_per_slot() - return max( - self._get_quota_from_max_tokens(0), - state_quota + attention_block_quota, - ) - - def _build_cache_config( - self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: - kv_cache_config = self.kv_cache_config - cache_tiers = config.cache_tiers - gpu_quota = cache_tiers[0].quota - minimum_live_quota = self._minimum_live_gpu_quota() - if minimum_live_quota > gpu_quota: - raise ValueError( - "The V2 Mamba GPU cache quota is too small for live recurrent " - f"states and attention pages: got {gpu_quota} bytes, need at " - f"least {minimum_live_quota} bytes.") - # _build_base_config already constructed every attention layer, - # including dtype-specific scale and subclass-provided side buffers. - # Preserve those configs and replace only the local Mamba layers. - layers = list(config.layers) - for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): - if self._mamba_layer_mask[global_layer_idx]: - layer_id = LayerId(local_layer_idx) - buffers = [ - BufferConfig(role=MambaRole.SSM_STATE, size=self.ssm_bytes), - BufferConfig(role=MambaRole.CONV_STATE, - size=self.conv_bytes), - ] - if global_layer_idx in self._ple_layer_ids: - buffers.extend(self._ple_buffer_configs()) - layers[local_layer_idx] = SsmLayerConfig( - layer_id=layer_id, - buffers=buffers, - ) - - dummy_requests = [ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(self._num_reserved_dummy_slots) - ] - constraints = [ - replace( - batch, - kv_caches=[*batch.kv_caches, *dummy_requests], - ) for batch in config.constraints - ] - - typical_step = config.typical_step - if config.initial_pool_ratio is None: - typical_capacity = self._get_typical_request_capacity( - kv_cache_config) - request_descs = self._typical_request_descs(typical_capacity, - kv_cache_config) - typical_step = BatchDesc(request_descs * - self._max_resident_sequences() + - dummy_requests) - # The recurrent (SSM) state pool must hold one slot per resident - # sequence plus every reserved dummy slot. Unlike attention pages, a - # Mamba state is fixed-size per sequence, so this floor is independent - # of sequence length. The base config only emits constraints when - # ``avg_seq_len`` is set, and speculative decoding inflates the reserved - # dummy slots (CUDA-graph padding), so without an explicit floor the SSM - # pool can be undersized (see the live/dummy-slot check in _setup_states - # / __init__). Add a min-slots constraint of zero-capacity requests: - # these cost no attention pages but reserve one SSM slot each. - if any(isinstance(layer, SsmLayerConfig) for layer in layers): - ssm_floor_slots = (self._max_resident_sequences() + - self._num_reserved_dummy_slots) - constraints = [ - *constraints, - BatchDesc([ - KVCacheDesc(capacity=0, history_length=0) - for _ in range(ssm_floor_slots) - ]), - ] - return replace( - config, - layers=layers, - typical_step=typical_step, - constraints=constraints, - # SSM lifecycles require minimum-snapshot commit semantics. The - # flag is harmless when reuse is disabled because no commits are - # attempted, while the runtime config still needs the invariant. - commit_min_snapshot=True, - ) - - def _get_state_buffer(self, local_layer_idx: int, role, dtype: torch.dtype, - state_shape: List[int]) -> torch.Tensor: - addr = self.impl.get_mem_pool_base_address(LayerId(local_layer_idx), - role, PageIndexMode.SHARED) - num_pages = self.impl.get_page_index_upper_bound( - LayerId(local_layer_idx), role) - raw = convert_to_torch_tensor( - TensorWrapper(addr, dtype, [num_pages] + state_shape)) - page_index_scale = self.impl.get_page_index_scale( - LayerId(local_layer_idx), role) - num_slots = (num_pages + page_index_scale - 1) // page_index_scale - # V2 coalesces same-size per-layer buffers inside each slot. Kernels - # index Mamba states by logical slot id, so expose only this layer's - # sub-page from each coalesced slot instead of the raw page-index view. - return raw.as_strided( - [num_slots] + state_shape, - [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), - ) - - def _setup_states(self) -> None: - local_layer_ids = [ - self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers - ] - self.all_ssm_states = [ - self._get_state_buffer(local_layer_idx, MambaRole.SSM_STATE, - self.ssm_state_dtype, self.ssm_state_shape) - for local_layer_idx in local_layer_ids - ] - self.all_conv_states = [ - self._get_state_buffer(local_layer_idx, MambaRole.CONV_STATE, - self.conv_state_dtype, self.conv_state_shape) - for local_layer_idx in local_layer_ids - ] - if self._use_gdn_cached_replay_all_layer_commit: - expected_inner_strides = ( - self.ssm_state_shape[1] * self.ssm_state_shape[2], - self.ssm_state_shape[2], - 1, - ) - reference = self.all_ssm_states[0] - for state in self.all_ssm_states: - if (state.dtype != reference.dtype - or state.device != reference.device - or list(state.shape[1:]) != self.ssm_state_shape - or state.stride()[1:] != expected_inner_strides): - raise RuntimeError( - "GDN cached replay requires V2 SSM layers with matching " - "dtype, device, shape, and dense inner dimensions.") - pointers = [state.data_ptr() for state in self.all_ssm_states] - slot_stride = reference.stride(0) - layer_stride_bytes = (pointers[1] - - pointers[0] if len(pointers) > 1 else 0) - has_affine_layout = ( - layer_stride_bytes % reference.element_size() == 0 and all( - state.stride(0) == slot_stride - for state in self.all_ssm_states) - and all(pointer == pointers[0] + layer * layer_stride_bytes - for layer, pointer in enumerate(pointers))) - if has_affine_layout: - self._gdn_cached_replay_state_strides = ( - len(pointers), - layer_stride_bytes // reference.element_size(), - slot_stride, - ) - else: - logger.warning_once( - "V2 GDN state views are not affine; using indirect replay " - "checkpoint addressing", - key="gdn_cached_replay_v2_indirect_state_layout", - ) - self._gdn_cached_replay_state_descriptors = torch.tensor( - [(state.data_ptr(), state.stride(0)) - for state in self.all_ssm_states], - dtype=torch.int64, - device=reference.device, - ) - - def _setup_replay_buffers(self, spec_config) -> None: - cache_size = 0 - device = None - if self.local_num_mamba_layers > 0: - cache_size = self.all_ssm_states[0].shape[0] - assert all(t.shape[0] == cache_size for t in self.all_ssm_states) - device = self.all_ssm_states[0].device - if not self._allocate_pool_replay_buffers(spec_config, cache_size, - device): - return - - mask_capacity = self._host_state_indices.shape[0] - self._dummy_request_mask = torch.zeros(mask_capacity, - dtype=torch.bool, - device=device) - self._dummy_request_mask_host = torch.zeros( - mask_capacity, - dtype=torch.bool, - pin_memory=prefer_pinned(), - ) - - def _attention_cache_bytes_per_token(self) -> int: - # Mamba layers have zero KV heads, so the generic calculation naturally - # returns only bytes owned by local attention layers. - return super().get_cache_bytes_per_token() - - def get_cache_bytes_per_token(self) -> int: - cache_bytes = self._attention_cache_bytes_per_token() - - interval = ( - self.kv_cache_config.mamba_state_config.periodic_snapshot_interval) - if (self.kv_cache_config.enable_block_reuse and interval is not None - and interval > 0): - cache_bytes += self._mamba_state_bytes_per_slot() // interval - if cache_bytes == 0 and self.local_num_mamba_layers > 0: - cache_bytes = self._mamba_state_bytes_per_slot() - return max(1, cache_bytes) - - def get_num_free_blocks(self) -> int: - assert len(self.kv_cache_map) == 0, ( - "get_num_free_blocks is only used when the kv cache manager is empty" - ) - attention_pages = [] - ssm_pages = [] - for local_layer_idx in range(self.num_local_layers): - layer_id = LayerId(local_layer_idx) - if self._is_local_mamba_layer(local_layer_idx): - ssm_pages.append( - self.impl.get_page_index_upper_bound( - layer_id, MambaRole.SSM_STATE) // - self._ssm_page_index_scale) - else: - attention_pages.append( - self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // - self.kv_factor) - if attention_pages: - return max(attention_pages) - return max(ssm_pages) if ssm_pages else 0 - - @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 _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: - for global_layer_id, local_layer_id in self.layer_offsets.items(): - if self._is_local_mamba_layer(local_layer_id): - continue - # A layer group is a lifecycle, not a physical memory pool. - # Differently sized attention buffers can share one lifecycle, - # so scan every attention layer in this diagnostic path. - yield KVCacheManagerV2.get_buffers(self, global_layer_id) - - yield from self.all_ssm_states - yield from self.all_conv_states - yield from self._ple_ngram_contexts.values() - yield from self._ple_conv_states.values() - - def _reset_context_mamba_slots(self, num_contexts: int) -> None: - super()._reset_context_mamba_slots(num_contexts) - if num_contexts == 0 or not self.use_kda_replay_update: - return - assert self.prev_num_accepted_tokens is not None - context_slots = self.cuda_state_indices[:num_contexts].long() - self.prev_num_accepted_tokens[context_slots] = 0 - - def _reset_kda_replay_slots( - self, - slots: Union[List[int], torch.Tensor], - ) -> None: - """Reset pending KDA drafts for newly initialized logical slots.""" - assert self.use_kda_replay_update - assert self.prev_num_accepted_tokens is not None - slots = torch.as_tensor(slots, - dtype=torch.long, - device=self.prev_num_accepted_tokens.device) - if slots.numel() > 0: - self.prev_num_accepted_tokens[slots] = 0 - - def _relocate_kda_replay_slots(self, old_slots: List[int], - new_slots: List[int]) -> None: - """Move replay history when V2 remaps a live request's SSM slot.""" - assert self.use_kda_replay_update - assert self.prev_num_accepted_tokens is not None - moves = [(old, new) for old, new in zip(old_slots, new_slots) - if old >= 0 and old != new] - if not moves: - return - - device = self.prev_num_accepted_tokens.device - source_slots = torch.tensor([old for old, _ in moves], - dtype=torch.long, - device=device) - destination_slots = torch.tensor([new for _, new in moves], - dtype=torch.long, - device=device) - replay_buffers = ( - self.kda_conv_q, - self.kda_conv_k, - self.kda_conv_v, - self.kda_qkg_cache, - self.kda_v_cache, - self.kda_beta_cache, - ) - for replay_buffer in replay_buffers: - assert replay_buffer is not None - replay_buffer.index_copy_( - 1, - destination_slots, - replay_buffer.index_select(1, source_slots), - ) - self.prev_num_accepted_tokens.index_copy_( - 0, - destination_slots, - self.prev_num_accepted_tokens.index_select(0, source_slots), - ) - - 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, - draft_kv_cache_manager: Optional[BaseResourceManager] = None, - capture_sampling_params: Optional["SamplingParams"] = 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, - draft_kv_cache_manager=draft_kv_cache_manager, - capture_sampling_params=capture_sampling_params, - ) - if requests and prepare_resource: - self._setup_state_indices( - requests, num_contexts=0 if is_gen else len(requests)) - return requests - - def free_resources(self, request: LlmRequest, pin_on_release: bool = False): - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is not None and kv_cache.is_active: - self.try_commit_blocks(request, kv_cache) - self._request_id_to_state_index.pop(request.py_request_id, None) - self._request_id_to_is_dummy.pop(request.py_request_id, None) - self._branch_snapshot_points.pop(request.py_request_id, None) - super().free_resources(request, pin_on_release) - - def prepare_resources(self, scheduled_batch: ScheduledRequests): - super().prepare_resources(scheduled_batch) - if self.local_num_mamba_layers == 0: - return - requests = (scheduled_batch.context_requests + - scheduled_batch.generation_requests) - num_contexts = len(scheduled_batch.context_requests) - self._setup_state_indices(requests, num_contexts=num_contexts) - self._reset_context_mamba_slots(num_contexts) - - def _setup_state_indices(self, - requests: List[LlmRequest], - num_contexts: int = 0) -> None: - if self.local_num_mamba_layers == 0: - return - replay_requests = (requests[num_contexts:] - if self.use_kda_replay_update else []) - old_state_values = [ - self._request_id_to_state_index.get(request.py_request_id, -1) - for request in replay_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_layer_group_id) - if base_index < 0: - raise RuntimeError( - f"Invalid SSM state block index {base_index} for " - f"request {req.py_request_id}") - self._host_state_indices[i] = base_index - - self.cuda_state_indices.copy_(self._host_state_indices, - non_blocking=True) - is_dummy = [req.is_dummy for req in requests] - self._refresh_dummy_request_mask(is_dummy) - state_values = self._host_state_indices[:n].tolist() - for req, value, dummy in zip(requests, state_values, is_dummy): - self._request_id_to_state_index[req.py_request_id] = value - self._request_id_to_is_dummy[req.py_request_id] = dummy - - if self.use_kda_replay_update: - new_state_values = [ - self._request_id_to_state_index[request.py_request_id] - for request in replay_requests - ] - self._relocate_kda_replay_slots(old_state_values, new_state_values) - self._reset_kda_replay_slots([ - new for old, new in zip(old_state_values, new_state_values) - if old < 0 - ]) - - def _record_kda_replay_acceptance( - self, - state_indices: torch.Tensor, - num_accepted_draft_tokens: torch.Tensor, - is_dummy_request: Optional[torch.Tensor] = None, - ) -> None: - """Record the draft prefix that the next KDA verify must replay.""" - if self.prev_num_accepted_tokens is None: - raise RuntimeError( - "KDA replay acceptance requires replay bookkeeping") - slots = state_indices.to(torch.long) - accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) - if is_dummy_request is not None: - current = self.prev_num_accepted_tokens[slots] - accepted = torch.where(is_dummy_request, current, accepted) - self.prev_num_accepted_tokens[slots] = accepted - - @torch.inference_mode() - def seed_kda_replay_caches_for_disagg_gen(self, - request_ids: List[int]) -> None: - """Seed replay caches from transferred ``W - 1`` conv states. - - This is the V2-storage counterpart of - ``MambaCacheManager.seed_kda_replay_caches_for_disagg_gen``; see that - method for the transfer-ordering rationale and replay contract. - """ - if (not self.use_kda_replay_update or self.local_num_mamba_layers == 0): - return - slots = [ - self._request_id_to_state_index[request_id] - for request_id in request_ids - if request_id in self._request_id_to_state_index - ] - if not slots: - return - - assert self.prev_num_accepted_tokens is not None - device = self.prev_num_accepted_tokens.device - state_indices = torch.tensor(sorted(set(slots)), - dtype=torch.long, - device=device) - committed_window = self.conv_state_shape[1] - section_offsets = [0] - for section_dim in self.conv_section_dims: - section_offsets.append(section_offsets[-1] + section_dim) - - replay_conv_buffers = (self.kda_conv_q, self.kda_conv_k, - self.kda_conv_v) - for layer_offset, conv_state in enumerate(self.all_conv_states): - selected_conv = conv_state.index_select(0, state_indices) - for section_idx, replay_buffer in enumerate(replay_conv_buffers): - assert replay_buffer is not None - replay_layer = replay_buffer[layer_offset] - replay_layer.index_fill_(0, state_indices, 0) - section = selected_conv[:, section_offsets[section_idx]: - section_offsets[section_idx + 1], :] - replay_layer[:, :, :committed_window].index_copy_( - 0, state_indices, section.to(replay_layer.dtype)) - - for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, - self.kda_beta_cache): - assert replay_buffer is not None - replay_buffer.index_fill_(1, state_indices, 0) - self.prev_num_accepted_tokens[state_indices] = 0 - - def mamba_layer_cache( - self, layer_idx: int - ) -> Union[PythonMambaCacheManager.State, - PythonMambaCacheManager.SpeculativeState, None]: - if not self.use_kda_replay_update: - return super().mamba_layer_cache(layer_idx) - - layer_offset = self.mamba_layer_offsets[layer_idx] - assert self.prev_num_accepted_tokens is not None - assert self.kda_conv_q is not None - assert self.kda_conv_k is not None - assert self.kda_conv_v is not None - assert self.kda_qkg_cache is not None - assert self.kda_v_cache is not None - assert self.kda_beta_cache is not None - spec_kwargs = { - "prev_num_accepted_tokens": self.prev_num_accepted_tokens, - "kda_conv_q": self.kda_conv_q[layer_offset], - "kda_conv_k": self.kda_conv_k[layer_offset], - "kda_conv_v": self.kda_conv_v[layer_offset], - "kda_qkg_cache": self.kda_qkg_cache[layer_offset], - "kda_v_cache": self.kda_v_cache[layer_offset], - "kda_beta_cache": self.kda_beta_cache[layer_offset], - } - if self.mamba_ssm_rand_seed is not None: - spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed - return PythonMambaCacheManager.SpeculativeState( - conv=self.get_conv_states(layer_idx), - temporal=self.get_ssm_states(layer_idx), - intermediate_conv_window=None, - **spec_kwargs, - ) - - def update_resources( - self, - scheduled_batch: ScheduledRequests, - attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float | None = None, - ) -> None: - super().update_resources(scheduled_batch, attn_metadata, - kv_cache_dtype_byte_size) - if (not self._record_kda_replay_in_update_resources - or self.local_num_mamba_layers == 0): - return - - generation_requests = scheduled_batch.generation_requests - drafted_requests = [ - request for request in generation_requests - if request.py_draft_tokens is not None - and len(request.py_draft_tokens) > 0 - ] - if not drafted_requests: - return - if len(drafted_requests) != len(generation_requests): - raise RuntimeError( - "Mixed drafted/undrafted generation batch is not supported " - "for KDA replay bookkeeping") - - active_requests = [ - request for request in generation_requests - if request.py_request_id in self._request_id_to_state_index - ] - if not active_requests: - return - assert self.prev_num_accepted_tokens is not None - device = self.prev_num_accepted_tokens.device - state_indices = torch.tensor([ - self._request_id_to_state_index[request.py_request_id] - for request in active_requests - ], - dtype=torch.int32, - device=device) - accepted_drafts = torch.tensor([ - request.py_num_accepted_draft_tokens for request in active_requests - ], - dtype=torch.int32, - device=device) - is_dummy_request = torch.tensor([ - self._request_id_to_is_dummy.get(request.py_request_id, False) - for request in active_requests - ], - dtype=torch.bool, - device=device) - self._record_kda_replay_acceptance(state_indices, accepted_drafts, - is_dummy_request) - - def get_state_indices(self, - request_ids: Optional[List[int]] = None, - is_padding: Optional[List[bool]] = None): - if self.local_num_mamba_layers == 0: - # Mamba metadata is still prepared on attention-only PP ranks, - # but no local kernel consumes these indices. Return harmless - # placeholders instead of consulting an intentionally empty map. - if request_ids is not None: - return [0] * len(request_ids) - return self.cuda_state_indices - if request_ids is not None: - indices = [ - self._request_id_to_state_index[rid] for rid in request_ids - ] - if is_padding is None: - is_padding = [False] * len(request_ids) - assert len(request_ids) == len(is_padding) - is_dummy = [ - self._request_id_to_is_dummy.get(rid, False) or padding - for rid, padding in zip(request_ids, is_padding) - ] - self._refresh_dummy_request_mask(is_dummy) - return indices - return self.cuda_state_indices - - def get_max_resource_count(self) -> int: - return self.max_batch_size - - def update_mamba_states( - self, - attn_metadata: "AttentionMetadata", - num_accepted_tokens: torch.Tensor, - state_indices: Optional[torch.Tensor] = None, - accepted_leaf_positions: Optional[torch.Tensor] = None): - if self.local_num_mamba_layers == 0: - return - batch_size = attn_metadata.num_seqs - num_contexts = attn_metadata.num_contexts - num_gens = batch_size - num_contexts - num_accepted_draft_tokens = ( - num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( - torch.int32) - 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) - if self.use_kda_replay_update: - assert self._dummy_request_mask is not None - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - self._record_kda_replay_acceptance( - state_indices_d, - num_accepted_draft_tokens, - is_dummy_request, - ) - return - - # Dynamic tree selects a tree node rather than a linear draft depth. - accepted_positions = (accepted_leaf_positions.to(torch.int32) - if accepted_leaf_positions is not None else - num_accepted_draft_tokens) - src_state_indices = self.intermediate_state_indices[:num_gens] - - if self._use_replay_state_update: - self._commit_gdn_cached_replay_history_layers( - attn_metadata, num_gens) - assert self._dummy_request_mask is not None - is_dummy_request = self._dummy_request_mask[ - num_contexts:num_contexts + num_gens] - replay_metadata = self.get_replay_state_update_metadata() - assert replay_metadata is not None - _advance_replay_state( - replay_metadata, - state_indices_d, - num_accepted_tokens[num_contexts:num_contexts + num_gens], - is_dummy_request, - ) - else: - for layer_offset, dst in enumerate(self.all_ssm_states): - _promote_mamba_state_triton( - dst.unsqueeze(0), - self.intermediate_ssm_states[layer_offset:layer_offset + 1], - src_state_indices, - accepted_positions, - state_indices_d, - ) - - for layer_offset, dst in enumerate(self.all_conv_states): - _promote_mamba_state_triton( - dst.unsqueeze(0), - self.intermediate_conv_states[layer_offset:layer_offset + 1], - src_state_indices, - accepted_positions, - state_indices_d, - ) - - def _mark_context_position_as_history(self, request: LlmRequest, - kv_cache) -> None: - """Advance history without making later recurrent state reusable.""" - history_length = request.context_current_position - if history_length <= kv_cache.history_length: - return - capacity = max(kv_cache.capacity, history_length) - if not kv_cache.resize(capacity, history_length=history_length): - raise ValueError( - "Failed to resize history length of V2 Mamba cache for " - f"request {request.py_request_id} to {history_length} tokens") - - def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: - should_block_reuse = (self.enable_block_reuse and not self.is_draft - and not request.is_dummy_request) - if not should_block_reuse: - return - - if kv_cache is None: - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is None: - return - - snapshot_points = request.expect_snapshot_points - commit_limit = (min(max(snapshot_points), request.prompt_len) - if snapshot_points else request.prompt_len) - commit_end = min(request.context_current_position, commit_limit) - if (request.context_current_position in request.expect_snapshot_points - and commit_end > kv_cache.num_committed_tokens): - tokens = self._augment_tokens_for_block_reuse( - request.get_tokens(DEFAULT_BEAM_INDEX), - request, - start=kv_cache.num_committed_tokens, - end=commit_end, - ) - kv_cache.commit(tokens) - if (self.local_num_mamba_layers > 0 - and self._branch_snapshot_points.get( - request.py_request_id) == commit_end): - self._branch_snapshots_taken_total += 1 - if request.context_current_position >= commit_limit: - self._mark_context_position_as_history(request, kv_cache) - if request.context_remaining_length == 0: - kv_cache.stop_committing() - - def update_context_resources(self, - scheduled_batch: ScheduledRequests) -> None: - for request in scheduled_batch.context_requests: - kv_cache = self.kv_cache_map.get(request.py_request_id) - if kv_cache is None or not kv_cache.is_active: - continue - - should_block_reuse = (self.enable_block_reuse and not self.is_draft - and not request.is_dummy_request) - is_all_reusable = ( - self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE) - is_snapshot_boundary = (request.context_current_position - in request.expect_snapshot_points) - has_pending_snapshot = any( - point > request.context_current_position - for point in request.expect_snapshot_points) - should_resize = (not should_block_reuse or - (not is_all_reusable and not has_pending_snapshot)) - should_commit = (is_all_reusable or is_snapshot_boundary - or request.context_remaining_length == 0) - - if should_resize and not kv_cache.resize( - None, request.context_current_position): - raise ValueError( - "Failed to resize history length of V2 Mamba cache for " - f"request {request.py_request_id} to " - f"{request.context_current_position} tokens at context " - "update") - if should_commit: - self.try_commit_blocks(request, kv_cache) - if request.context_remaining_length == 0: - if self.conversation_manager is not None: - self.conversation_manager.save_drop_plan(request, kv_cache) - kv_cache.enable_swa_scratch_reuse = False - - def shutdown(self): - self._gdn_cached_replay_state_descriptors = None - self._gdn_cached_replay_state_strides = None - self.all_ssm_states = [] - self.all_conv_states = [] - self._ple_ngram_contexts = {} - self._ple_conv_states = {} - self.intermediate_ssm_states = None - self.intermediate_conv_states = None - self.intermediate_state_indices = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None - self._dummy_request_mask_host = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - self._branch_snapshot_points.clear() - self.kda_conv_q = None - self.kda_conv_k = None - self.kda_conv_v = None - self.kda_qkg_cache = None - self.kda_v_cache = None - self.kda_beta_cache = None - super().shutdown() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py new file mode 100644 index 000000000000..9a1f85fd1e4a --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Stable, lazily loaded exports for Mamba cache managers. + +The former module is now a package. Lazy exports preserve its import surface +without making model-owned feature modules depend on manager import order. +""" + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .common import ( + BaseMambaCacheManager, + MambaHybridCacheManager, + MambaLayerCache, + MambaRole, + ReplayStateUpdateMetadata, + SpeculativeMambaLayerCache, + ) + from .legacy import ( + CppMambaHybridCacheManager, + MambaCacheManager, + MixedMambaHybridCacheManager, + PythonMambaCacheManager, + ) + from .mamba_cache_manager_v2 import MambaHybridCacheManagerV2 + +_EXPORTS = { + "get_tensor_size_bytes": ("common", "get_tensor_size_bytes"), + "_estimate_mamba_hybrid_cache_cost": ("common", "_estimate_mamba_hybrid_cache_cost"), + "BaseMambaCacheManager": ("common", "BaseMambaCacheManager"), + "CppMambaHybridCacheManager": ("legacy", "CppMambaHybridCacheManager"), + "MIN_REPLAY_HISTORY_SIZE": ("common", "MIN_REPLAY_HISTORY_SIZE"), + "MambaCacheManager": ("legacy", "MambaCacheManager"), + "MambaHybridCacheManager": ("common", "MambaHybridCacheManager"), + "MambaHybridCacheManagerV2": ( + "mamba_cache_manager_v2", + "MambaHybridCacheManagerV2", + ), + "MambaLayerCache": ("common", "MambaLayerCache"), + "MambaRole": ("common", "MambaRole"), + "MixedMambaHybridCacheManager": ("legacy", "MixedMambaHybridCacheManager"), + "PythonMambaCacheManager": ("legacy", "PythonMambaCacheManager"), + "ReplayStateUpdateMetadata": ("common", "ReplayStateUpdateMetadata"), + "SpeculativeMambaLayerCache": ("common", "SpeculativeMambaLayerCache"), + "_advance_replay_state": ("common", "_advance_replay_state"), + "_allocate_mamba_seed_buffer": ("common", "_allocate_mamba_seed_buffer"), + "_compute_deterministic_mamba_seed": ( + "common", + "_compute_deterministic_mamba_seed", + ), + "_get_local_mamba_cache_layout": ("common", "_get_local_mamba_cache_layout"), + "_get_mamba_hybrid_pool_size": ("common", "_get_mamba_hybrid_pool_size"), + "_get_num_cuda_graph_padding_dummy_slots": ( + "common", + "_get_num_cuda_graph_padding_dummy_slots", + ), + "_mamba_effective_tp_size": ("common", "_mamba_effective_tp_size"), + "_mamba_rank_offset": ("common", "_mamba_rank_offset"), + "_mamba_snapshot_rule_counts": ("common", "_mamba_snapshot_rule_counts"), + "_promote_mamba_state_triton": ("common", "_promote_mamba_state_triton"), + "use_py_mamba_cache_manager": ("common", "use_py_mamba_cache_manager"), +} + + +def __getattr__(name: str): + export = _EXPORTS.get(name) + if export is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module_name, attribute_name = export + value = getattr(import_module(f"{__name__}.{module_name}"), attribute_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted({*globals(), *_EXPORTS}) + + +__all__ = [ + "BaseMambaCacheManager", + "CppMambaHybridCacheManager", + "MIN_REPLAY_HISTORY_SIZE", + "MambaCacheManager", + "MambaHybridCacheManager", + "MambaHybridCacheManagerV2", + "MambaLayerCache", + "MambaRole", + "MixedMambaHybridCacheManager", + "PythonMambaCacheManager", + "ReplayStateUpdateMetadata", + "SpeculativeMambaLayerCache", + "use_py_mamba_cache_manager", +] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py new file mode 100644 index 000000000000..138e361445bd --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -0,0 +1,796 @@ +# 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"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +import os +import sys +from abc import ABC, abstractmethod +from dataclasses import dataclass, field, fields, replace +from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Protocol, Tuple + +import torch + +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override +import triton +import triton.language as tl + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.resource_manager import ( + BaseResourceManager, + KVCacheManager, + get_pp_layers, +) +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import DataRole + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + + +GB = 1 << 30 + + +class ReplayStateUpdateMetadata(NamedTuple): + """Shared tensors and fixed sizes for replay state updates.""" + + prev_num_accepted_tokens: torch.Tensor + cache_buf_idx: torch.Tensor + replay_step_width: int + replay_history_size: int + + +@dataclass(frozen=True, kw_only=True) +class MambaLayerCache: + """Persistent recurrent-state views for one local layer.""" + + conv: torch.Tensor + temporal: torch.Tensor + + def at_layer_idx(self, layer: int) -> "MambaLayerCache": + """Return a per-layer view while preserving shared slot metadata.""" + values = {} + for cache_field in fields(self): + value = getattr(self, cache_field.name) + values[cache_field.name] = ( + value + if value is None or cache_field.metadata.get("slot_shared", False) + else value[layer] + ) + return type(self)(**values) + + +@dataclass(frozen=True, kw_only=True) +class SpeculativeMambaLayerCache(MambaLayerCache): + """Generic intermediate-state view used during speculative decoding.""" + + intermediate_conv_window: torch.Tensor | None = None + intermediate_ssm: torch.Tensor | None = None + mamba_ssm_rand_seed: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + + +@dataclass(frozen=True, kw_only=True) +class MambaStateLayout: + """Resolved recurrent geometry; slot capacity is filled after pool allocation.""" + + layer_mask: tuple[bool, ...] + pp_layers: tuple[int, ...] + mamba_pp_layers: tuple[int, ...] + mapping: "Mapping" + conv_state_shape: tuple[int, ...] + ssm_state_shape: tuple[int, ...] + conv_state_dtype: torch.dtype + ssm_state_dtype: torch.dtype + n_groups_per_rank: int + max_batch_size: int + state_index_capacity: int + slot_capacity: int | None + spec_config: object | None + stochastic_rounding: bool = False + seed_rank_offset: int = 0 + conv_section_dims: tuple[int, ...] = field(default_factory=tuple) + conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c" + + def with_slot_capacity(self, slot_capacity: int) -> "MambaStateLayout": + """Return the post-allocation context for tensor-view binding.""" + return replace(self, slot_capacity=slot_capacity) + + +@dataclass(frozen=True, kw_only=True) +class MambaAcceptanceBatch: + """Accepted-token inputs for model-specific recurrent updates.""" + + attention_metadata: object + num_contexts: int + num_generations: int + num_accepted_tokens: torch.Tensor + accepted_positions: torch.Tensor + source_state_indices: torch.Tensor + destination_state_indices: torch.Tensor + is_dummy_request: torch.Tensor | None + + +class BaseMambaCacheManager(ABC): + """Abstract interface for accessing Mamba/recurrent state caches.""" + + @abstractmethod + def get_state_indices(self, *args, **kwargs) -> torch.Tensor: + """Return slot indices of each request with shape ``[max_batch_size]``.""" + + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + """Return replay metadata tensors and fixed replay sizes.""" + return None + + @abstractmethod + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + """Return conv states with shape ``[slot_size, conv_dim, d_conv - 1]``.""" + + @abstractmethod + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + """Return SSM states with shape ``[slot_size, heads, head_dim, d_state]``.""" + + @abstractmethod + def is_speculative(self) -> bool: + """Whether speculative state is allocated.""" + + @abstractmethod + def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache | None: + """Return the typed recurrent cache payload for one layer.""" + + def on_state_transfer_complete(self, request_ids: list[int]) -> None: + """Notify a model manager after a disaggregated transfer.""" + + +MIN_REPLAY_HISTORY_SIZE = 16 + + +def _advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: torch.Tensor | None = None, +) -> None: + """Compatibility wrapper for Mamba2-owned replay bookkeeping.""" + from tensorrt_llm._torch.modules.mamba.cache_manager import advance_replay_state + + advance_replay_state(replay_metadata, state_indices, accepted_tokens, is_dummy_request) + + +def _allocate_mamba_seed_buffer( + cache_size: int, rank_offset: int, device: torch.device +) -> torch.Tensor: + """Compatibility wrapper for Mamba2-owned seed allocation.""" + from tensorrt_llm._torch.modules.mamba.cache_manager import allocate_mamba_seed_buffer + + return allocate_mamba_seed_buffer(cache_size, rank_offset, device) + + +def _compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: + """Compatibility wrapper for Mamba2-owned deterministic seeding.""" + from tensorrt_llm._torch.modules.mamba.cache_manager import compute_deterministic_mamba_seed + + return compute_deterministic_mamba_seed(counter, slot, rank_offset) + + +def _mamba_rank_offset(mapping: Mapping) -> int: + """Stable identity for per-rank recurrent state initialization.""" + return mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + mapping.rank * 1_009 + + +class _PrefixReuseDiagnostics(Protocol): + def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: ... + + def _get_num_reusable_tokens_before_pruning(self) -> int: ... + + +def _get_num_cuda_graph_padding_dummy_slots( + spec_config: Optional["DecodingBaseConfig"], + max_batch_size: int, +) -> int: + """Return the number of persistent CUDA-graph padding dummy IDs. + + This is computed before ``ModelEngine`` exists and covers draft lengths + reachable at every batch size, including the zero-length acceptance-rate + fallback. ``ModelEngine._compute_dynamic_draft_len_mapping`` is created + later and covers only configured CUDA-graph batch sizes, so it cannot size + this persistent ID set. + """ + if spec_config is None: + return 1 + + draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + supports_dynamic_draft_len = ( + spec_dec_mode is not None + and hasattr(spec_dec_mode, "support_dynamic_draft_len") + and spec_dec_mode.support_dynamic_draft_len() + ) + if draft_len_schedule and supports_dynamic_draft_len: + runtime_draft_lengths = set() + first_uncovered_batch_size = 1 + for batch_size_threshold, draft_len in draft_len_schedule.items(): + if first_uncovered_batch_size > max_batch_size: + break + if batch_size_threshold >= first_uncovered_batch_size: + runtime_draft_lengths.add(draft_len) + first_uncovered_batch_size = batch_size_threshold + 1 + if first_uncovered_batch_size <= max_batch_size: + runtime_draft_lengths.add(0) + else: + max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 + max_total_draft_tokens = getattr(spec_config, "max_total_draft_tokens", 0) or 0 + is_linear_tree = getattr( + spec_config, + "is_linear_tree", + max_draft_len == max_total_draft_tokens, + ) + static_draft_len = max_draft_len if is_linear_tree else max_total_draft_tokens + runtime_draft_lengths = {static_draft_len or 0} + + if (getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and ( + getattr(spec_config, "acceptance_rate_threshold", 0) or 0 + ) > 0: + runtime_draft_lengths.add(0) + return len(runtime_draft_lengths) + + +class MambaRole: + """V2 buffer roles owned only by the hybrid Mamba manager.""" + + SSM_STATE = DataRole("ssm_state") + CONV_STATE = DataRole("conv_state") + + +def _mamba_effective_tp_size(mapping: Mapping) -> int: + """TP degree for sizing per-rank recurrent-state pools. + + Attention-DP replicates the state and takes precedence; helix + repurposes CP ranks as plain TP for recurrent-state layers. + """ + if mapping.enable_attention_dp: + return 1 + if mapping.has_cp_helix(): + return mapping.tp_size * mapping.cp_size + return mapping.tp_size + + +def get_tensor_size_bytes(tensor): + """Calculate tensor size in bytes.""" + if isinstance(tensor, torch.Tensor): + return tensor.element_size() * tensor.nelement() + elif isinstance(tensor, list): + return sum(get_tensor_size_bytes(t) for t in tensor) + return 0 + + +def use_py_mamba_cache_manager() -> bool: + """Check if PythonMambaCacheManager should be forced (agg mode override). + + Returns True if TRTLLM_USE_PY_MAMBA='1' is set, False otherwise. + + Agg-mode-only override: forces the V1-route MixedMambaHybridCacheManager + with PythonMambaCacheManager inside instead of the configured manager. + Disagg mode is unaffected — its compatibility routing selects Mixed or + Cpp based on the transceiver configuration. + """ + return os.environ.get("TRTLLM_USE_PY_MAMBA", "0") == "1" + + +class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): + """Common recurrent views and snapshot planning, independent of algorithms.""" + + _supports_additional_snapshot_offsets = False + + def prepare_expect_snapshot_points(self, requests: List[LlmRequest]) -> None: + """Set reusable Mamba snapshot boundaries before scheduling.""" + if not self.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + state_config = self.kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + for request in requests: + snapshot_points = set() + if interval is not None and interval > 0: + snapshot_points.update(range(interval, request.prompt_len + 1, interval)) + if self._supports_additional_snapshot_offsets: + for offset in state_config.additional_snapshot_offsets_from_start: + if offset <= request.prompt_len: + snapshot_points.add(offset) + for offset in state_config.additional_snapshot_offsets_from_end: + point = request.prompt_len - offset + if point > 0: + snapshot_points.add(point) + request.expect_snapshot_points = sorted(snapshot_points) + + @override + def is_speculative(self) -> bool: + return self.spec_config is not None + + @override + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] + + @override + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] + + def get_mamba_ssm_rand_seed(self) -> torch.Tensor | None: + return None + + +def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: + """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" + pool_size = max_batch_size + # One permanent slot is shared by every CUDA-graph padding sentinel. + pool_size += 1 + if mapping.enable_attention_dp: + # Attention-DP can insert a transient dummy request on an otherwise + # idle rank. Keep this headroom internal so scheduler-visible + # max_batch_size still limits real requests. + pool_size += 1 + return pool_size + + +@triton.jit +def _promote_mamba_state_kernel( + src_ptr, + dst_ptr, + src_idx_ptr, + accepted_position_source_ptr, + num_accepted_tokens_ptr, + blk_ptr, + replay_pnat_ptr, + replay_cache_buf_idx_ptr, + dummy_request_mask_ptr, + num_gens, + count, + src_s_layer, + src_s_row, + src_s_step, + dst_s_layer, + dst_s_block, + POSITION_SOURCE_IS_TOKEN_COUNT: tl.constexpr, + REPLAY_STEP_WIDTH: tl.constexpr, + REPLAY_HISTORY_SIZE: tl.constexpr, + UPDATE_REPLAY_BOOKKEEPING: tl.constexpr, + BLOCK: tl.constexpr, +): + # One program copies a BLOCK-sized tile of the contiguous inner state for a + # single (layer, gen) pair. grid = (num_layers * num_gens, ceil(count/BLOCK)). + pair = tl.program_id(0) + tile = tl.program_id(1) + layer = (pair // num_gens).to(tl.int64) + g = pair % num_gens + row = tl.load(src_idx_ptr + g).to(tl.int64) + accepted_position = tl.load(accepted_position_source_ptr + g) + if POSITION_SOURCE_IS_TOKEN_COUNT: + accepted_position -= 1 + acc = accepted_position.to(tl.int64) + blk = tl.load(blk_ptr + g).to(tl.int64) + + # The replay metadata update is batch-scoped, so layer 0's first copy tile + # owns it. This folds PNAT and active-buffer maintenance into the existing + # accepted-convolution-state promotion launch. + if UPDATE_REPLAY_BOOKKEEPING: + if (layer == 0) & (tile == 0): + num_accepted_tokens = tl.load(num_accepted_tokens_ptr + g) + is_dummy_request = tl.load(dummy_request_mask_ptr + g) + pnat = tl.load(replay_pnat_ptr + blk) + cache_buf_idx = tl.load(replay_cache_buf_idx_ptr + blk) + wrote_checkpoint = pnat + REPLAY_STEP_WIDTH > REPLAY_HISTORY_SIZE + next_pnat = tl.where(wrote_checkpoint, num_accepted_tokens, pnat + num_accepted_tokens) + next_cache_buf_idx = tl.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx) + tl.store(replay_pnat_ptr + blk, next_pnat, mask=~is_dummy_request) + tl.store(replay_cache_buf_idx_ptr + blk, next_cache_buf_idx, mask=~is_dummy_request) + + # int64 throughout: per-layer strides are O(1e8), so layer*stride overflows + # int32 and would corrupt addresses / fault. + src_base = layer * src_s_layer + row * src_s_row + acc * src_s_step + dst_base = layer * dst_s_layer + blk * dst_s_block + offs = tile * BLOCK + tl.arange(0, BLOCK) + mask = offs < count + v = tl.load(src_ptr + src_base + offs, mask=mask) + tl.store(dst_ptr + dst_base + offs, v, mask=mask) + + +def _promote_mamba_state_triton( + dst: torch.Tensor, + intermediate: torch.Tensor, + src_state_indices: torch.Tensor, + accepted_position_source: torch.Tensor, + dst_state_indices: torch.Tensor, + num_accepted_tokens: Optional[torch.Tensor] = None, + position_source_is_token_count: bool = False, + replay_pnat: Optional[torch.Tensor] = None, + replay_cache_buf_idx: Optional[torch.Tensor] = None, + dummy_request_mask: Optional[torch.Tensor] = None, + replay_step_width: int = 0, + replay_history_size: int = 0, + BLOCK: int = 2048, +) -> None: + """Scatter each generation request's accepted draft-step recurrent state + from the per-request intermediate buffer into the unified C++ pool view. + + Args: + dst: ``[num_layers, num_blocks, *state_shape]`` view of the C++ pool. + The block dim (dim 1) is strided (the pool interleaves + ``ssm_bytes | conv_bytes`` per block), so ``dst`` is non-contiguous. + The kernel writes through ``dst``'s real strides, touching only this + state's bytes -- the case that defeats torch.compile / aot_autograd + (dtype-view mutation) and Inductor codegen (``XBLOCK`` on the uint8 + pool). + intermediate: ``[num_layers, max_batch_size, T, *state_shape]`` dense. + src_state_indices, accepted_position_source, dst_state_indices: + ``[num_gens]`` int. + num_accepted_tokens: ``[num_gens]`` accepted-token counts. Required + when replay bookkeeping is enabled. + position_source_is_token_count: Whether ``accepted_position_source`` + contains accepted-token counts instead of zero-based intermediate + state positions. When true, the kernel subtracts one while + computing the source position. + replay_pnat, replay_cache_buf_idx, dummy_request_mask: Optional replay + metadata. When supplied together, the kernel advances PNAT and the + active history buffer while promoting the convolution state. + + For each gen ``g``:: + + accepted_position = ( + accepted_position_source[g] - 1 + if position_source_is_token_count + else accepted_position_source[g] + ) + dst[:, dst_state_indices[g]] = intermediate[:, src_state_indices[g], accepted_position] + + Pure gather->scatter copy; bandwidth-bound (~85% of HBM peak), one launch. + """ + num_layers = dst.shape[0] + num_gens = src_state_indices.shape[0] + if num_gens == 0: + return + update_replay_bookkeeping = replay_pnat is not None + if update_replay_bookkeeping: + assert num_accepted_tokens is not None + assert replay_cache_buf_idx is not None + assert dummy_request_mask is not None + assert replay_step_width > 0 + assert replay_history_size >= replay_step_width + assert replay_pnat.dtype == torch.int32 + assert replay_cache_buf_idx.dtype == torch.int32 + assert dummy_request_mask.dtype == torch.bool + assert dummy_request_mask.numel() >= num_gens + else: + # Compile-time-disabled pointer arguments still require tensors. + num_accepted_tokens = accepted_position_source + assert replay_cache_buf_idx is None + assert dummy_request_mask is None + replay_pnat = dst_state_indices + replay_cache_buf_idx = dst_state_indices + dummy_request_mask = dst_state_indices + count = 1 + for s in dst.shape[2:]: + count *= s + # Grid dim order matters: dim 0 -> CUDA grid.x (limit 2^31-1), dim 1 -> + # grid.y (limit 65535). Put the (layer, gen) pairs in dim 0 since that is + # what grows with batch size (num_layers*num_gens stays far below 2^31 for + # any real config); the tile count in dim 1 is small (~count/BLOCK). Do NOT + # swap them: pairs would hit the 65535 y-limit at num_layers*num_gens>65535. + grid = (num_layers * num_gens, triton.cdiv(count, BLOCK)) + _promote_mamba_state_kernel[grid]( + intermediate, + dst, + src_state_indices, + accepted_position_source, + num_accepted_tokens, + dst_state_indices, + replay_pnat, + replay_cache_buf_idx, + dummy_request_mask, + num_gens, + count, + intermediate.stride(0), + intermediate.stride(1), + intermediate.stride(2), + dst.stride(0), + dst.stride(1), + POSITION_SOURCE_IS_TOKEN_COUNT=position_source_is_token_count, + REPLAY_STEP_WIDTH=replay_step_width, + REPLAY_HISTORY_SIZE=replay_history_size, + UPDATE_REPLAY_BOOKKEEPING=update_replay_bookkeeping, + BLOCK=BLOCK, + ) + + +def _mamba_snapshot_rule_counts( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], + tokens_per_block: int, +) -> Tuple[int, int]: + """Return reachable fixed rules and their partial-block upper bound.""" + if not kv_cache_config.enable_block_reuse: + return 0, 0 + + num_rules = 0 + num_unaligned_rules = 0 + state_config = kv_cache_config.mamba_state_config + for offset in set(state_config.additional_snapshot_offsets_from_start): + if max_seq_len is not None and offset > max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += int(offset % tokens_per_block != 0) + for offset in set(state_config.additional_snapshot_offsets_from_end): + # A from-end offset is reachable iff some valid prompt is longer than + # the offset. Its absolute alignment depends on that prompt. + if max_seq_len is not None and offset >= max_seq_len: + continue + num_rules += 1 + num_unaligned_rules += 1 + return num_rules, num_unaligned_rules + + +def _mamba_regular_snapshot_interval( + kv_cache_config: KvCacheConfig, + max_seq_len: Optional[int], +) -> Optional[int]: + if not kv_cache_config.enable_block_reuse: + return None + interval = kv_cache_config.mamba_state_config.periodic_snapshot_interval + if interval is None or interval <= 0: + return None + if max_seq_len is not None and interval > max_seq_len: + return None + return interval + + +def _get_local_mamba_cache_layout( + model_config, + mapping: Mapping, + *, + spec_config=None, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, +): + """Return normalized params and local Mamba/attention layer counts. + + Cache construction and affine sizing must follow the model's PP layout: + partition base layers first, then place appended speculative layers on the + last PP rank. The normalized params retain target masks and the appended + draft-layer count so estimation selects the same combined or per-manager + layout as runtime. + """ + from tensorrt_llm._torch.pyexecutor.config_utils import extract_mamba_kv_cache_params + + params = extract_mamba_kv_cache_params( + model_config.pretrained_config, + spec_config=spec_config, + quant_config=model_config.quant_config, + ) + mamba_layer_mask, full_attention_layer_mask = params.get_layer_masks( + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + combined_layer_mask = [ + is_mamba or is_attention + for is_mamba, is_attention in zip(mamba_layer_mask, full_attention_layer_mask) + ] + local_layer_indices, _ = get_pp_layers( + sum(combined_layer_mask), + mapping, + spec_config=spec_config, + layer_mask=combined_layer_mask, + ) + local_mamba_layers = sum(mamba_layer_mask[layer_idx] for layer_idx in local_layer_indices) + local_attention_layers = sum( + full_attention_layer_mask[layer_idx] for layer_idx in local_layer_indices + ) + return params, local_mamba_layers, local_attention_layers + + +def _estimate_mamba_hybrid_cache_cost( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + tokens_per_block: int, + max_seq_len: Optional[int], + num_reserved_dummy_slots: int, + include_explicit_snapshots: bool, + cap_partial_attention_snapshots: bool, + extra_state_bytes_per_rank: int = 0, + is_draft: bool = False, + use_separate_draft_kv_cache: bool = False, + **kwargs, +) -> Tuple[int, int]: + spec_config = kwargs.get("spec_config") + params, local_mamba_layers, local_attention_layers = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=is_draft, + use_separate_draft_kv_cache=use_separate_draft_kv_cache, + ) + attention_slope = ( + KVCacheManager.get_cache_size_per_token( + model_config, + mapping, + num_layers=local_attention_layers, + **kwargs, + ) + if local_attention_layers > 0 + else 0 + ) + state_bytes_per_rank = local_mamba_layers * params.get_states_bytes_per_layer(mapping) + state_bytes_per_rank += extra_state_bytes_per_rank + max_resident_sequences = max_batch_size * mapping.pp_size + + if include_explicit_snapshots: + fixed_rules, unaligned_fixed_rules = _mamba_snapshot_rule_counts( + kv_cache_config, max_seq_len, tokens_per_block + ) + else: + fixed_rules = 0 + unaligned_fixed_rules = 0 + fixed_state_slots = ( + max_resident_sequences + num_reserved_dummy_slots + max_resident_sequences * fixed_rules + ) + attention_block_bytes = attention_slope * tokens_per_block + + interval = _mamba_regular_snapshot_interval(kv_cache_config, max_seq_len) + has_unaligned_periodic_snapshot = interval is not None and interval % tokens_per_block != 0 + if cap_partial_attention_snapshots: + # Snapshot alignment is unknown while estimating cache cost. Once a + # snapshot is possible, reserve one retained partial attention page + # per resident lineage. Dummy requests carry no attention capacity. + has_non_live_ssm_capacity = fixed_rules > 0 or interval is not None + partial_attention_slots = max_resident_sequences if has_non_live_ssm_capacity else 0 + else: + partial_attention_slots = max_resident_sequences * unaligned_fixed_rules + intercept = ( + fixed_state_slots * state_bytes_per_rank + partial_attention_slots * attention_block_bytes + ) + + if interval is None: + regular_slope = 0 + else: + regular_slope = math.ceil(state_bytes_per_rank / interval) + if has_unaligned_periodic_snapshot and not cap_partial_attention_snapshots: + regular_slope += math.ceil(attention_block_bytes / interval) + return attention_slope + regular_slope, intercept + + +# The split legacy modules import the original implementation helpers as a +# group. Keep private helpers in this internal export list while the package +# facade exposes only its compatibility surface. +__all__ = [ + "BaseMambaCacheManager", + "MambaHybridCacheManager", + "MambaLayerCache", + "SpeculativeMambaLayerCache", + "MambaStateLayout", + "MambaAcceptanceBatch", + "IntermediateState", + "ReplayStateUpdateMetadata", + "MambaRole", + "MIN_REPLAY_HISTORY_SIZE", + "use_py_mamba_cache_manager", + "get_tensor_size_bytes", +] + + +def _promote_intermediate_states( + destination_states: tuple[torch.Tensor, ...], + intermediate_states: torch.Tensor, + batch: MambaAcceptanceBatch, +) -> None: + """Copy accepted positions from [layer, batch, step, ...] into per-layer pools.""" + # Preserve the established kernel monkeypatch point without coupling the + # concrete intermediate and replay algorithms through inheritance. + from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager + + for layer_offset, destination in enumerate(destination_states): + mamba_cache_manager._promote_mamba_state_triton( + destination.unsqueeze(0), + intermediate_states[layer_offset : layer_offset + 1], + batch.source_state_indices, + batch.accepted_positions, + batch.destination_state_indices, + ) + + +class IntermediateState: + """Own conventional speculative intermediate states and promotion.""" + + def __init__(self) -> None: + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self.intermediate_ssm: torch.Tensor | None = None + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + """Borrow persistent per-layer views until shutdown; allocate owned scratch.""" + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + if not context.mamba_pp_layers: + return + + device = ssm_states[0].device + + if context.spec_config is not None: + tokens_per_step = context.spec_config.tokens_per_gen_step + common_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_ssm = torch.zeros( + common_shape + list(context.ssm_state_shape), + dtype=context.ssm_state_dtype, + device=device, + ) + self.intermediate_conv = torch.zeros( + common_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + fields: dict[str, torch.Tensor | None] = {} + if self.intermediate_conv is not None: + fields["intermediate_conv_window"] = self.intermediate_conv[layer_offset] + if self.intermediate_ssm is not None: + fields["intermediate_ssm"] = self.intermediate_ssm[layer_offset] + return fields + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + """Plain intermediate buffers have no request-persistent bookkeeping.""" + + def update(self, batch: MambaAcceptanceBatch) -> None: + if self.intermediate_indices is None: + return + _promote_intermediate_states(self._ssm_states, self.intermediate_ssm, batch) + _promote_intermediate_states(self._conv_states, self.intermediate_conv, batch) + + def shutdown(self) -> None: + self._ssm_states = () + self._conv_states = () + self.intermediate_ssm = None + self.intermediate_conv = None + self.intermediate_indices = None + + def make_layer_cache( + self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor + ) -> MambaLayerCache: + return SpeculativeMambaLayerCache( + conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset) + ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py new file mode 100644 index 000000000000..702ddb113049 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py @@ -0,0 +1,2320 @@ +# 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"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Union + +import torch + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + from tensorrt_llm.sampling_params import SamplingParams + +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, 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 .common import ( + MIN_REPLAY_HISTORY_SIZE, + BaseMambaCacheManager, + MambaHybridCacheManager, + ReplayStateUpdateMetadata, + _advance_replay_state, + _allocate_mamba_seed_buffer, + _compute_deterministic_mamba_seed, + _get_local_mamba_cache_layout, + _get_mamba_hybrid_pool_size, + _mamba_effective_tp_size, + _mamba_rank_offset, + _promote_mamba_state_triton, + get_tensor_size_bytes, +) +from .common import _estimate_mamba_hybrid_cache_cost as _estimate_common_mamba_cache_cost + +GB = 1 << 30 + + +def _estimate_mamba_hybrid_cache_cost(model_config, mapping, **kwargs): + """Keep V1's original model-specific estimate outside the V2 base.""" + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + _qwen4_exp_ple_state_bytes_per_rank, + ) + + extra_bytes = 0 + if not kwargs.get("is_draft", False): + params, _, _ = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=kwargs.get("spec_config"), + use_separate_draft_kv_cache=kwargs.get("use_separate_draft_kv_cache", False), + ) + extra_bytes = _qwen4_exp_ple_state_bytes_per_rank( + model_config, + params, + mapping, + spec_config=kwargs.get("spec_config"), + use_separate_draft_kv_cache=kwargs.get("use_separate_draft_kv_cache", False), + ) + return _estimate_common_mamba_cache_cost( + model_config, mapping, extra_state_bytes_per_rank=extra_bytes, **kwargs + ) + + +def warn_legacy_mamba_cache_manager( + implementation: Literal["python", "cpp"], manager_name: str +) -> None: + """Warn once when a V1 implementation is selected or constructed.""" + display_name = "Python" if implementation == "python" else "C++" + logger.warning_once( + f"{manager_name} uses the deprecated {display_name} Mamba cache " + "manager V1 path. Set kv_cache_config.use_kv_cache_manager_v2=True; " + "V1 remains available only for compatibility.", + key=f"deprecated_mamba_cache_manager_{implementation}_v1", + ) + + +class PythonMambaCacheManager(BaseResourceManager): + """Pure-Python mamba state manager with speculative decoding support. + + Manages only mamba states (conv + SSM) using PyTorch tensors on GPU. + Supports speculative decoding and disaggregated serving. + """ + + @dataclass(frozen=True, kw_only=True) + class State: + """Base state container for Mamba cache.""" + + conv: torch.Tensor + # Represents last state, or "two back" if prev_num_accepted_tokens is > 0 + temporal: torch.Tensor + + # Subclasses override to list fields shared across layers (not indexed by layer). + _SHARED_FIELDS = frozenset() + + def at_layer_idx(self, layer: int): + kwargs = {} + for k, v in vars(self).items(): + if v is None: + kwargs[k] = None + elif k in self._SHARED_FIELDS: + kwargs[k] = v + else: + kwargs[k] = v[layer] + return type(self)(**kwargs) + + @property + def has_kda_replay_caches(self) -> bool: + """Whether the Kimi K3 KDA fused-verify replay caches are + allocated. The base (non-speculative) state never has them; the + speculative state overrides this.""" + return False + + @dataclass(frozen=True, kw_only=True) + class SpeculativeState(State): + """Speculative state with intermediate states for draft tokens. + + Supports three SSM update paths (only one set of tensors is + allocated): + - Legacy: caches full intermediate SSM states (intermediate_ssm) + - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) + - KDA replay: per-slot draft-token caches consumed by the fused + ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted + drafts and commits states in place (kda_conv_*, kda_*_cache) + """ + + _SHARED_FIELDS = frozenset( + {"prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed"} + ) + + # Allocated for the legacy and Mamba2-replay paths; None for the + # KDA replay path (the kernel commits conv windows in place). + intermediate_conv_window: torch.Tensor | None = None + + # Legacy path: full intermediate SSM states at each step + intermediate_ssm: torch.Tensor | None = None + + # KDA replay path (fused multi-token verify, kimi_linear). + # Pool invariant under this path: `temporal` holds the state after + # the LAST GOLDEN token; the accepted drafts recorded in + # prev_num_accepted_tokens are pending in these caches and are + # replayed by the kernel at the start of the next verify round. + # Extended conv caches [layers, slots, dim, (W-1) + num_spec] fp32, + # dim-contiguous (stride(dim) == 1): columns [0, W-1) are the + # committed raw-input window, tail columns the pending drafts' raw + # inputs. + kda_conv_q: torch.Tensor | None = None + kda_conv_k: torch.Tensor | None = None + kda_conv_v: torch.Tensor | None = None + # Post-processed per-draft quantities for replay: + # [layers, slots, num_spec, 3, H*K] (q/k/gate), [.., num_spec, H*V], + # [.., num_spec, H] — all fp32. + kda_qkg_cache: torch.Tensor | None = None + kda_v_cache: torch.Tensor | None = None + kda_beta_cache: torch.Tensor | None = None + + # Replay path: compact double-buffered cache + # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). + # 0 means temporal saved state is actually the last state, not two back. + prev_num_accepted_tokens: torch.Tensor | None = None # (cache,) int — shared across layers + cache_buf_idx: torch.Tensor | None = None # (cache,) int32 — shared across layers + # Per-cache-slot Philox seeds. Replay bumps them in-place per launch so + # CUDA graph replay uses fresh SR draws without allocating RNG tensors. + # (cache,) int64 - shared across layers + mamba_ssm_rand_seed: torch.Tensor | None = None + old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) + old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) + # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. + old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + + @property + def has_kda_replay_caches(self) -> bool: + """True when the KDA fused-verify replay caches were allocated + (the fused ``trtllm::kda_mtp_decode`` verify path).""" + return self.kda_qkg_cache is not None + + def commit_conv_window(self, slot_indices: torch.Tensor, conv_pool: torch.Tensor) -> None: + """Seed the KDA replay caches from the live ``W - 1`` conv pool.""" + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import ( + copy_kda_replay_conv_window, + ) + + copy_kda_replay_conv_window( + conv_pool, + self.kda_conv_q, + self.kda_conv_k, + self.kda_conv_v, + slot_indices, + ) + + def __init__( + self, + d_state: int, + d_conv: int, + num_heads: int, + n_groups: int, + head_dim: int, + num_layers: int, + max_batch_size: int, + spec_state_size: int, + mapping: Mapping, + dtype: torch.dtype, + ssm_cache_dtype: torch.dtype, + layer_mask: Optional[List[bool]] = None, + speculative_num_draft_tokens: Optional[int] = None, + model_type: str = "nemotron_hybrid", + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + ) -> None: + warn_legacy_mamba_cache_manager("python", type(self).__name__) + self.mamba_ssm_cache_dtype = ssm_cache_dtype + self.speculative_num_draft_tokens = speculative_num_draft_tokens + self.spec_state_size = spec_state_size + self._use_replay_state_update = use_replay_state_update + # KDA replay path (kimi_linear fused multi-token verify). Mutually + # exclusive with use_replay_state_update; requires speculative mode. + self._kda_replay_num_spec = kda_replay_num_spec + self._use_kda_replay_update = kda_replay_num_spec is not None + if self._use_kda_replay_update: + assert not use_replay_state_update, ( + "kda_replay_num_spec and use_replay_state_update are mutually exclusive" + ) + assert speculative_num_draft_tokens is not None, ( + "KDA replay caches require speculative decoding" + ) + assert kda_replay_num_spec == speculative_num_draft_tokens, ( + f"KDA replay cache width ({kda_replay_num_spec}) must match " + f"the draft length ({speculative_num_draft_tokens}): the " + "fused verify kernel is compiled with a static NUM_SPEC" + ) + self.replay_history_size: Optional[int] = None + self.replay_step_width: Optional[int] = None + # When True, allocate the per-slot Philox seed buffer even outside + # the replay path so the non-replay flashinfer SR kernel reads a + # persistent deterministic seed instead of a per-call torch.randint. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh cache-slot assignment. Combined + # with slot index and rank offset to produce reproducible per-slot + # seed values. Starts at 0 so the post-init "reset" stream is + # disjoint from the counter=0 stream used at allocation time. + self._seed_request_counter = 0 + + # get tp size + tp_size = _mamba_effective_tp_size(mapping) + + # derive mamba parameters for conv and ssm states + d_inner = head_dim * num_heads + conv_dim = d_inner + 2 * n_groups * d_state + nheads = num_heads + + # check that can be partitioned + assert nheads % tp_size == 0, "nheads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + + # partition conv_dim and nheads + d_inner_local = d_inner // tp_size + ng_ds_local = n_groups * d_state // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + d_inner = d_inner // tp_size + + # Per-section dims for conv_state. + # Qwen3-Next: [Q | K | V] = [ng*ds, ng*ds, d_inner] + # Nemotron_hybrid: [x | B | C] = [d_inner, ng*ds, ng*ds] + 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}") + + # conv and ssm states device + device = torch.device("cuda") + + pp_layers, num_layers = get_pp_layers( + num_layers, + mapping, + layer_mask=layer_mask, + ) + num_local_layers = len(pp_layers) + self.mamba_layer_offsets = {idx: offset for offset, idx in enumerate(pp_layers)} + + conv_state_shape = (conv_dim, d_conv - 1) + ssm_state_shape = (nheads, head_dim, d_state) + + # create mamba conv and ssm states + conv_states = torch.zeros( + size=(num_local_layers, max_batch_size) + conv_state_shape, + dtype=dtype, + device=device, + ) + + ssm_states = torch.zeros( + size=(num_local_layers, max_batch_size) + ssm_state_shape, + dtype=self.mamba_ssm_cache_dtype, + device=device, + ) + + # Per-slot Philox seeds. Allocated whenever stochastic rounding can + # fire, even outside the replay path and even when MTP/spec is off, + # so all consumers (replay kernel, MTP non-replay flashinfer, + # non-MTP flashinfer) read a persistent deterministic seed from the + # cache manager instead of calling torch.randint per forward. + self._mamba_ssm_rand_seed: Optional[torch.Tensor] = None + if self._use_replay_state_update or self._mamba_ssm_stochastic_rounding: + self._mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + max_batch_size, self._seed_rank_offset, device + ) + + # create state container + if speculative_num_draft_tokens is not None: + T = speculative_num_draft_tokens + 1 + self.replay_step_width = T + + # Conv intermediate cache — legacy and Mamba2-replay paths only. + # The KDA replay kernel commits conv windows in place, so the + # per-step window scratch is not needed. + intermediate_conv_window_cache = None + if not self._use_kda_replay_update: + intermediate_conv_window_cache = torch.zeros( + size=(num_local_layers, self.spec_state_size, T) + conv_state_shape, + dtype=dtype, + device=device, + ) + + # SSM speculative cache — path-specific tensors + spec_kwargs = {} + # Share the manager-level seed buffer through SpeculativeState + # so the MTP path can still read it via layer_cache. + if self._mamba_ssm_rand_seed is not None: + spec_kwargs["mamba_ssm_rand_seed"] = self._mamba_ssm_rand_seed + if self._use_kda_replay_update: + # KDA replay caches for the fused multi-token verify kernel + # (trtllm::kda_mtp_decode). Per cache slot (persistent + # across rounds — the next round replays from them), unlike + # the batch-row-indexed intermediate buffers. + M = self._kda_replay_num_spec + # conv_dim covers the [q | k | v] short-conv sections; the + # kernel consumes them as three dim-contiguous caches. + assert conv_dim % 3 == 0, ( + "KDA replay caches expect the [q | k | v] conv-state " + "sectioning (3 equal sections)" + ) + section_dim = conv_dim // 3 + # d_conv is KDA's convolution width; the live pool stores + # its W - 1 committed input columns. + w_kernel = d_conv + extended_s = w_kernel - 1 + M + + def _dim_contiguous_conv_cache(): + return torch.zeros( + num_local_layers, + max_batch_size, + extended_s, + section_dim, + dtype=torch.float32, + device=device, + ).transpose(-1, -2) + + spec_kwargs["prev_num_accepted_tokens"] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device + ) + spec_kwargs["kda_conv_q"] = _dim_contiguous_conv_cache() + spec_kwargs["kda_conv_k"] = _dim_contiguous_conv_cache() + spec_kwargs["kda_conv_v"] = _dim_contiguous_conv_cache() + spec_kwargs["kda_qkg_cache"] = torch.zeros( + num_local_layers, + max_batch_size, + M, + 3, + section_dim, + dtype=torch.float32, + device=device, + ) + spec_kwargs["kda_v_cache"] = torch.zeros( + num_local_layers, + max_batch_size, + M, + section_dim, + dtype=torch.float32, + device=device, + ) + spec_kwargs["kda_beta_cache"] = torch.zeros( + num_local_layers, max_batch_size, M, nheads, dtype=torch.float32, device=device + ) + ssm_spec_cache = [ + spec_kwargs["kda_conv_q"], + spec_kwargs["kda_conv_k"], + spec_kwargs["kda_conv_v"], + spec_kwargs["kda_qkg_cache"], + spec_kwargs["kda_v_cache"], + spec_kwargs["kda_beta_cache"], + ] + spec_path_label = "kda-replay" + elif self._use_replay_state_update: + assert n_groups % tp_size == 0, ( + "replay state update requires n_groups divisible by tp_size" + ) + n_groups_per_rank = n_groups // tp_size + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) + + # Compact replay cache. + spec_kwargs["prev_num_accepted_tokens"] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device + ) + spec_kwargs["cache_buf_idx"] = torch.zeros( + max_batch_size, dtype=torch.int32, device=device + ) + spec_kwargs["old_x"] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + nheads, + head_dim, + dtype=dtype, + device=device, + ) + spec_kwargs["old_B"] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + n_groups_per_rank, + d_state, + dtype=dtype, + device=device, + ) + spec_kwargs["old_dt"] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device, + ) + spec_kwargs["old_dA_cumsum"] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device, + ) + ssm_spec_cache = [ + spec_kwargs["old_x"], + spec_kwargs["old_B"], + spec_kwargs["old_dt"], + spec_kwargs["old_dA_cumsum"], + ] + spec_path_label = "replay" + else: + # Legacy: full intermediate SSM states at each step + spec_kwargs["intermediate_ssm"] = torch.zeros( + size=(num_local_layers, self.spec_state_size, T) + ssm_state_shape, + dtype=self.mamba_ssm_cache_dtype, + device=device, + ) + ssm_spec_cache = [spec_kwargs["intermediate_ssm"]] + spec_path_label = "legacy" + + self.mamba_cache = self.SpeculativeState( + conv=conv_states, + temporal=ssm_states, + intermediate_conv_window=intermediate_conv_window_cache, + **spec_kwargs, + ) + + logger.info( + f"Mamba Cache ({spec_path_label}) is allocated. " + f"max_mamba_cache_size: {max_batch_size}, " + f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " + f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB, " + f"ssm_spec_cache size: {get_tensor_size_bytes(ssm_spec_cache) / GB:.2f}GB, " + "intermediate_conv_window_cache size: " + f"{get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" + ) + else: + self.mamba_cache = self.State( + conv=conv_states, + temporal=ssm_states, + ) + + logger.info( + f"Mamba Cache is allocated. " + f"max_mamba_cache_size: {max_batch_size}, " + f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " + f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB" + ) + + # mamba cache available blocks + self.mamba_cache_free_blocks = [i for i in range(max_batch_size)] + + # mamba cache index, maps request_id -> state indices + self.mamba_cache_index: Dict[int, int] = {} + self._dummy_request_ids: set[int] = set() + # Batch-order mask aligned with state_indices; duplicate dummy request + # IDs mark every batch row even when they share one cache slot. + self._dummy_request_mask = torch.zeros(max_batch_size, dtype=torch.bool, device=device) + self._dummy_request_mask_host = torch.zeros( + max_batch_size, dtype=torch.bool, pin_memory=prefer_pinned() + ) + + # Permanent slot shared by every CUDA-graph padding sentinel id + # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per + # draft length). Pool sizing must include +1 headroom for this; + # see MixedMambaHybridCacheManager. + self._padding_slot: int = self.mamba_cache_free_blocks.pop() + + # Reserved slot for the attention-DP padding dummy. Keep it out of the + # free pool so dummy insertion never consumes real-request capacity. + self._attention_dp_dummy_slot: Optional[int] = ( + self.mamba_cache_free_blocks.pop() if mapping.enable_attention_dp else None + ) + + # save intermediate state indices for requests + self.intermediate_state_indices = torch.arange( + max_batch_size, dtype=torch.int32, device=device + ) + + # Physical tensor rows include reserved dummy slots. Resource capacity + # is the number of real request slots remaining after those + # reservations. + self._max_batch_size = max_batch_size + self._max_resource_count = len(self.mamba_cache_free_blocks) + + def get_max_resource_count(self) -> int: + """Return the maximum number of real requests that can be cached.""" + return self._max_resource_count + + def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: + """Return the prefix of *context_requests* that fits in the + available Mamba cache blocks. Requests that already have a + cached block do not consume a free block.""" + free = len(self.mamba_cache_free_blocks) + result = [] + for r in context_requests: + if r.py_request_id in self.mamba_cache_index: + result.append(r) + elif free > 0: + result.append(r) + free -= 1 + else: + break + return result + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + """For Mamba cache manager, we always need one slot per request.""" + return 1 + + @torch.inference_mode() + def _prepare_mamba_cache_blocks(self, request_ids: List[int]): + for r in request_ids: + if r in self.mamba_cache_index: + continue + if len(self.mamba_cache_free_blocks) == 0: + raise RuntimeError("run out of mamba cache blocks") + block = self.mamba_cache_free_blocks.pop() + self.mamba_cache_index[r] = block + if ( + isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update + ): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif ( + isinstance(self.mamba_cache, self.SpeculativeState) and self._use_kda_replay_update + ): + # Fresh request: no drafts pending in the replay caches. + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + if self._mamba_ssm_rand_seed is not None: + # Deterministic per-slot rotation on fresh assignment. + # `block` is pulled from mamba_cache_free_blocks, which + # excludes _padding_slot by construction (see __init__), + # so padding sentinels never reach this branch. + self._seed_request_counter += 1 + self._mamba_ssm_rand_seed[block] = _compute_deterministic_mamba_seed( + self._seed_request_counter, block, self._seed_rank_offset + ) + + @torch.inference_mode() + def seed_kda_replay_caches_for_disagg_gen(self, request_ids: List[int]) -> None: + """Seed the fused-verify KDA replay conv caches from the conv pool. + + On a disaggregated generation server the ctx->gen transfer populates + only the base ``conv`` / ``temporal`` pools (see + ``disaggregation/resource/page.py``); the per-slot ``kda_conv_*`` + replay caches consumed by ``trtllm::kda_mtp_decode`` are normally + seeded by a *local* prefill/decode via + ``_sync_kda_replay_conv_window`` and would otherwise hold zeros (or a + previous occupant's window) for a transferred request — corrupting + the recurrent state from the first verify step onward. Call this + after the state transfer completes, before the first generation + forward. Mirrors ``_sync_kda_replay_conv_window``: the conv pool row + stores the committed ``W - 1`` raw inputs directly. Pending-draft + scratch is cleared (no drafts are pending for a freshly transferred + request). + """ + if not ( + self._use_kda_replay_update and isinstance(self.mamba_cache, self.SpeculativeState) + ): + return + blocks = [ + self.mamba_cache_index[rid] for rid in request_ids if rid in self.mamba_cache_index + ] + if not blocks: + return + conv = self.mamba_cache.conv # [L, slots, 3D, W - 1] + idx = torch.tensor(sorted(set(blocks)), dtype=torch.long, device=conv.device) + d = conv.shape[2] // 3 + committed = conv.shape[3] + cs = conv.index_select(1, idx) + for cache, section in ( + (self.mamba_cache.kda_conv_q, cs[:, :, :d]), + (self.mamba_cache.kda_conv_k, cs[:, :, d : 2 * d]), + (self.mamba_cache.kda_conv_v, cs[:, :, 2 * d :]), + ): + # cache: [L, slots, D, committed + num_spec]; zero the draft + # tail columns and seed the committed window in one copy. + seeded = torch.zeros( + (cache.shape[0], idx.numel()) + cache.shape[2:], + dtype=cache.dtype, + device=cache.device, + ) + seeded[:, :, :, :committed] = section.to(cache.dtype) + cache.index_copy_(1, idx, seeded) + for buf in ( + self.mamba_cache.kda_qkg_cache, + self.mamba_cache.kda_v_cache, + self.mamba_cache.kda_beta_cache, + ): + buf.index_fill_(1, idx, 0) + self.mamba_cache.prev_num_accepted_tokens[idx] = 0 + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + requests = scheduled_batch.context_requests + scheduled_batch.generation_requests + context_ids = [i.py_request_id for i in scheduled_batch.context_requests] + generation_ids = [i.py_request_id for i in scheduled_batch.generation_requests] + request_ids = context_ids + generation_ids + self._prepare_mamba_cache_blocks(request_ids) + self._refresh_dummy_request_mask([req.is_dummy for req in requests]) + + def _is_padding_sentinel(self, request_id: int) -> bool: + # cuda_graph_runner caches one dummy per runtime_draft_len value + # (see _get_padded_batch), so any id in the range of dummy request IDs + # may be live concurrently. + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID + + max_dl = self.speculative_num_draft_tokens or 0 + return CUDA_GRAPH_DUMMY_REQUEST_ID - max_dl <= request_id <= CUDA_GRAPH_DUMMY_REQUEST_ID + + @torch.inference_mode() + def add_dummy_requests(self, request_ids: List[int], **kwargs): + # Sentinels alias to the shared _padding_slot; non-sentinel + # dummies (warmup, attention-DP idle padding) get their own + # slot and are freed individually. + if not request_ids: + return + self._dummy_request_ids.update(request_ids) + for r in request_ids: + if r in self.mamba_cache_index: + block = self.mamba_cache_index[r] + if ( + isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update + ): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif ( + isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_kda_replay_update + ): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + continue + if self._is_padding_sentinel(r): + block = self._padding_slot + elif r == ATTENTION_DP_DUMMY_REQUEST_ID and self._attention_dp_dummy_slot is not None: + block = self._attention_dp_dummy_slot + else: + if len(self.mamba_cache_free_blocks) == 0: + raise RuntimeError("run out of mamba cache blocks") + block = self.mamba_cache_free_blocks.pop() + self.mamba_cache_index[r] = block + if ( + isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update + ): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 + elif ( + isinstance(self.mamba_cache, self.SpeculativeState) and self._use_kda_replay_update + ): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + + def free_resources(self, request: LlmRequest): + request_id = request.py_request_id + if request_id not in self.mamba_cache_index: + return + self._dummy_request_ids.discard(request_id) + block = self.mamba_cache_index.pop(request_id) + # Reserved slots must not re-enter the real-request free pool. + if block != self._padding_slot and block != self._attention_dp_dummy_slot: + self.mamba_cache_free_blocks.append(block) + + def get_state_indices(self, request_ids: List[int], is_padding: List[bool]) -> List[int]: + assert len(request_ids) == len(is_padding) + indices = [self.mamba_cache_index[rid] for rid in request_ids] + is_dummy = [ + rid in self._dummy_request_ids or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_(torch.as_tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, non_blocking=True) + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).conv + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).temporal + + def get_intermediate_ssm_states(self, layer_idx: int) -> Optional[torch.Tensor]: + if not isinstance(self.mamba_cache, self.SpeculativeState): + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).intermediate_ssm + + def get_intermediate_conv_states(self, layer_idx: int) -> Optional[torch.Tensor]: + """Get intermediate conv states for speculative decoding.""" + if not isinstance(self.mamba_cache, self.SpeculativeState): + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset).intermediate_conv_window + + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_x[layer_offset] + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_B[layer_offset] + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dt[layer_offset] + + def get_replay_old_dA_cumsum(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_dA_cumsum[layer_offset] + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.cache_buf_idx + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + return self.mamba_cache.prev_num_accepted_tokens + + def is_speculative(self) -> bool: + return isinstance(self.mamba_cache, self.SpeculativeState) + + def mamba_layer_cache(self, layer_idx: int) -> Union[State, SpeculativeState]: + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.at_layer_idx(layer_offset) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.mamba_ssm_cache_dtype + + 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. + + Used by mamba2_mixer non-MTP paths that don't hold a SpeculativeState. + Callers must bump in-place (`.add_(1)` or slice-and-add) and pass a + view that matches the consuming kernel's expected shape. + """ + return self._mamba_ssm_rand_seed + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + @property + def use_kda_replay_update(self) -> bool: + return self._use_kda_replay_update + + def get_replay_state_update_metadata(self) -> Optional[ReplayStateUpdateMetadata]: + if ( + not self._use_replay_state_update + or not isinstance(self.mamba_cache, self.SpeculativeState) + or self.mamba_cache.prev_num_accepted_tokens is None + or self.mamba_cache.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=(self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=self.mamba_cache.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size, + ) + + def shutdown(self): + """Release tensor memory.""" + # Clear mamba cache states + empty = torch.tensor([]) + + def _drop(tensor): + return empty if tensor is not None else None + + if isinstance(self.mamba_cache, self.SpeculativeState): + self.mamba_cache = self.SpeculativeState( + conv=empty, + temporal=empty, + intermediate_conv_window=_drop(self.mamba_cache.intermediate_conv_window), + intermediate_ssm=_drop(self.mamba_cache.intermediate_ssm), + prev_num_accepted_tokens=_drop(self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=_drop(self.mamba_cache.cache_buf_idx), + mamba_ssm_rand_seed=_drop(self.mamba_cache.mamba_ssm_rand_seed), + old_x=_drop(self.mamba_cache.old_x), + old_B=_drop(self.mamba_cache.old_B), + old_dt=_drop(self.mamba_cache.old_dt), + old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), + kda_conv_q=_drop(self.mamba_cache.kda_conv_q), + kda_conv_k=_drop(self.mamba_cache.kda_conv_k), + kda_conv_v=_drop(self.mamba_cache.kda_conv_v), + kda_qkg_cache=_drop(self.mamba_cache.kda_qkg_cache), + kda_v_cache=_drop(self.mamba_cache.kda_v_cache), + kda_beta_cache=_drop(self.mamba_cache.kda_beta_cache), + ) + else: + self.mamba_cache = self.State(conv=empty, temporal=empty) + + torch.cuda.empty_cache() + + @torch.compile(options={"max-autotune": True}) + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None, + ): + 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 + # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. + accepted_positions = ( + accepted_leaf_positions + if accepted_leaf_positions is not None + else num_accepted_draft_tokens + ) + state_indices_d = state_indices[num_contexts : num_contexts + num_gens] + src_state_indices = self.intermediate_state_indices[:num_gens] + + if self._use_kda_replay_update: + # KDA replay: the fused verify kernel already committed the SSM + # state and conv windows (in place, after the golden token) and + # cached this round's drafts. All that remains is recording how + # many of those drafts the sampler accepted, so the next round's + # kernel launch replays exactly that prefix. + is_dummy_request = self._dummy_request_mask[num_contexts : num_contexts + num_gens] + prev = self.mamba_cache.prev_num_accepted_tokens + current = prev[state_indices_d] + accepted = num_accepted_draft_tokens.to(torch.int32).clamp(min=0) + prev[state_indices_d] = torch.where(is_dummy_request, current, accepted) + return + + if self._use_replay_state_update: + is_dummy_request = self._dummy_request_mask[num_contexts : num_contexts + num_gens] + replay_metadata = self.get_replay_state_update_metadata() + assert replay_metadata is not None + _advance_replay_state( + replay_metadata, + state_indices_d, + num_accepted_tokens[num_contexts : num_contexts + num_gens], + is_dummy_request, + ) + else: + # Legacy: copy accepted SSM state from intermediate cache. + ssm_states = self.mamba_cache.temporal + intermediate_ssm_cache = self.mamba_cache.intermediate_ssm + accepted_ssm_state = intermediate_ssm_cache[:, src_state_indices, accepted_positions] + ssm_states[:, state_indices_d, :] = accepted_ssm_state + + # Conv: both paths save all intermediate conv windows, carry over the accepted one. + conv_states = self.mamba_cache.conv + intermediate_conv_window_cache = self.mamba_cache.intermediate_conv_window + accepted_conv_state = intermediate_conv_window_cache[ + :, src_state_indices, accepted_positions + ] + conv_states[:, state_indices_d, :] = accepted_conv_state + + def on_state_transfer_complete(self, request_ids: list[int]) -> None: + self.seed_kda_replay_caches_for_disagg_gen(request_ids) + + +class MambaCacheManager(BaseResourceManager, BaseMambaCacheManager): + """Facade for standalone mamba state management (no KV cache). + + Delegates to PythonMambaCacheManager. + """ + + def __init__( + self, + d_state: int, + d_conv: int, + num_heads: int, + n_groups: int, + head_dim: int, + num_layers: int, + max_batch_size: int, + spec_state_size: int, + mapping: Mapping, + dtype: torch.dtype, + ssm_cache_dtype: torch.dtype, + layer_mask: Optional[List[bool]] = None, + stream: Optional[torch.cuda.Stream] = None, + speculative_num_draft_tokens: Optional[int] = None, + model_type: str = "nemotron_hybrid", + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + ) -> None: + warn_legacy_mamba_cache_manager("python", type(self).__name__) + max_num_sequences = max_batch_size * mapping.pp_size + + self._impl = PythonMambaCacheManager( + d_state=d_state, + d_conv=d_conv, + num_heads=num_heads, + n_groups=n_groups, + head_dim=head_dim, + num_layers=num_layers, + max_batch_size=max_num_sequences, + spec_state_size=spec_state_size, + mapping=mapping, + dtype=dtype, + ssm_cache_dtype=ssm_cache_dtype, + layer_mask=layer_mask, + speculative_num_draft_tokens=speculative_num_draft_tokens, + model_type=model_type, + use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, + ) + + def get_max_resource_count(self) -> int: + return self._impl.get_max_resource_count() + + def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: + return self._impl.filter_ctx_requests_by_capacity(context_requests) + + def get_needed_resource_to_completion(self, request: LlmRequest) -> int: + return self._impl.get_needed_resource_to_completion(request) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + self._impl.prepare_resources(scheduled_batch) + + def free_resources(self, request: LlmRequest): + self._impl.free_resources(request) + + def add_dummy_requests(self, request_ids: List[int], **kwargs): + self._impl.add_dummy_requests(request_ids, **kwargs) + + def get_state_indices( + self, request_ids: Optional[List[int]] = None, is_padding: Optional[List[bool]] = None + ) -> Union[torch.Tensor, List[int]]: + return self._impl.get_state_indices(request_ids, is_padding) + + @property + def mamba_cache_free_blocks(self) -> List[int]: + return self._impl.mamba_cache_free_blocks + + @property + def mamba_cache_index(self) -> Dict[int, int]: + return self._impl.mamba_cache_index + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self._impl.get_conv_states(layer_idx) + + def seed_kda_replay_caches_for_disagg_gen(self, request_ids: List[int]) -> None: + self._impl.seed_kda_replay_caches_for_disagg_gen(request_ids) + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self._impl.get_ssm_states(layer_idx) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self._impl.get_mamba_ssm_cache_dtype() + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + """Delegate to the underlying Python manager.""" + return self._impl.get_mamba_ssm_rand_seed() + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + @property + def use_kda_replay_update(self) -> bool: + return self._impl.use_kda_replay_update + + def get_replay_state_update_metadata(self) -> Optional[ReplayStateUpdateMetadata]: + return self._impl.get_replay_state_update_metadata() + + def get_intermediate_ssm_states(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_intermediate_ssm_states(layer_idx) + + def get_intermediate_conv_states(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_intermediate_conv_states(layer_idx) + + def get_replay_old_x(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_x(layer_idx) + + def get_replay_old_B(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_B(layer_idx) + + def get_replay_old_dt(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_dt(layer_idx) + + def get_replay_old_dA_cumsum(self, layer_idx: int) -> Optional[torch.Tensor]: + return self._impl.get_replay_old_dA_cumsum(layer_idx) + + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: + return self._impl.get_replay_cache_buf_idx() + + def get_replay_prev_num_accepted_tokens(self) -> Optional[torch.Tensor]: + return self._impl.get_replay_prev_num_accepted_tokens() + + def is_speculative(self) -> bool: + return self._impl.is_speculative() + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, PythonMambaCacheManager.SpeculativeState, None]: + return self._impl.mamba_layer_cache(layer_idx) + + def shutdown(self): + self._impl.shutdown() + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: torch.Tensor, + accepted_leaf_positions: Optional[torch.Tensor] = None, + ): + # Non-speculative configs don't allocate intermediate state; the + # promotion is a clean no-op. + if not self._impl.is_speculative(): + return + self._impl.update_mamba_states( + attn_metadata, num_accepted_tokens, state_indices, accepted_leaf_positions + ) + + def on_state_transfer_complete(self, request_ids: list[int]) -> None: + self.seed_kda_replay_caches_for_disagg_gen(request_ids) + + +class _LegacyMambaHybridCacheManager(MambaHybridCacheManager): + """Shared interface and state plumbing for hybrid Mamba managers. + + Concrete storage, state-index, and resource lifecycles remain owned by the + Cpp and V2 implementations. + """ + + _supports_additional_snapshot_offsets = False + + def _setup_mtp_intermediate_states(self, spec_config, max_batch_size: int) -> None: + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + if spec_config is None or self.local_num_mamba_layers == 0: + return + + tokens_per_gen_step = spec_config.tokens_per_gen_step + if not self._use_replay_state_update: + self.intermediate_ssm_states = torch.zeros( + size=[self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step] + + self.ssm_state_shape, + dtype=self.ssm_state_dtype, + device="cuda", + ) + self.intermediate_conv_states = torch.zeros( + size=[self.local_num_mamba_layers, max_batch_size, tokens_per_gen_step] + + self.conv_state_shape, + dtype=self.conv_state_dtype, + device="cuda", + ) + self.intermediate_state_indices = torch.arange( + max_batch_size, dtype=torch.int32, device="cuda" + ) + + def _allocate_pool_replay_buffers( + self, + spec_config, + cache_size: int, + device: Optional[torch.device], + ) -> bool: + """Allocate replay tensors shared by the Cpp and V2 state pools.""" + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + if self.local_num_mamba_layers == 0 or ( + not self._use_replay_state_update and not self._mamba_ssm_stochastic_rounding + ): + return False + + assert device is not None + self.mamba_ssm_rand_seed = _allocate_mamba_seed_buffer( + cache_size, self._seed_rank_offset, device + ) + if spec_config is None or not self._use_replay_state_update: + return False + + history_size = self.replay_history_size + assert history_size is not None + nheads, head_dim, d_state = self.ssm_state_shape + common_shape = [self.local_num_mamba_layers, cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.old_x = torch.zeros( + common_shape + [history_size, nheads, head_dim], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_shape + [history_size, self._n_groups_per_rank, d_state], + dtype=self.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros( + common_shape + [nheads, history_size], + dtype=torch.float32, + device=device, + ) + return True + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + if self._dummy_request_mask is None: + return + + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_(torch.tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, non_blocking=True) + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + if num_contexts == 0: + return + + context_slots = self.cuda_state_indices[:num_contexts].long() + if ( + self._use_replay_state_update + and self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None + ): + self.prev_num_accepted_tokens[context_slots] = 0 + self.cache_buf_idx[context_slots] = 0 + if self.old_x is not None: + self.old_x[:, context_slots] = 0 + if self.old_B is not None: + self.old_B[:, context_slots] = 0 + if self.old_dt is not None: + self.old_dt[:, context_slots] = 0 + if self.old_dA_cumsum is not None: + self.old_dA_cumsum[:, context_slots] = 0 + + if self.mamba_ssm_rand_seed is None: + return + self._seed_request_counter += 1 + counter = self._seed_request_counter + rank_offset = self._seed_rank_offset + host_slots = self._host_state_indices[:num_contexts].tolist() + new_seeds = [ + _compute_deterministic_mamba_seed(counter, slot, rank_offset) for slot in host_slots + ] + self.mamba_ssm_rand_seed[context_slots] = torch.tensor( + new_seeds, + dtype=torch.int64, + device=self.mamba_ssm_rand_seed.device, + ) + + def prepare_expect_snapshot_points(self, requests: List[LlmRequest]) -> None: + """Set reusable Mamba snapshot boundaries before scheduling.""" + if not self.enable_block_reuse: + for request in requests: + request.expect_snapshot_points = [] + return + + state_config = self.kv_cache_config.mamba_state_config + interval = state_config.periodic_snapshot_interval + for request in requests: + snapshot_points = set() + if interval is not None and interval > 0: + snapshot_points.update(range(interval, request.prompt_len + 1, interval)) + if self._supports_additional_snapshot_offsets: + for offset in state_config.additional_snapshot_offsets_from_start: + if offset <= request.prompt_len: + snapshot_points.add(offset) + for offset in state_config.additional_snapshot_offsets_from_end: + point = request.prompt_len - offset + if point > 0: + snapshot_points.add(point) + request.expect_snapshot_points = sorted(snapshot_points) + + def is_speculative(self) -> bool: + return self.spec_config is not None + + def get_ssm_states(self, layer_idx: int) -> torch.Tensor: + return self.all_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_conv_states(self, layer_idx: int) -> torch.Tensor: + return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_ssm_states(self, layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_ssm_states is None: + return None + return self.intermediate_ssm_states[self.mamba_layer_offsets[layer_idx]] + + def get_intermediate_conv_states(self, layer_idx: int) -> Optional[torch.Tensor]: + if self.intermediate_conv_states is None: + return None + return self.intermediate_conv_states[self.mamba_layer_offsets[layer_idx]] + + def mamba_layer_cache( + self, layer_idx: int + ) -> Union[PythonMambaCacheManager.State, PythonMambaCacheManager.SpeculativeState, None]: + conv = self.get_conv_states(layer_idx) + ssm = self.get_ssm_states(layer_idx) + if self.spec_config is not None: + layer_offset = self.mamba_layer_offsets[layer_idx] + spec_kwargs = {} + if self.mamba_ssm_rand_seed is not None: + spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed + if self._use_replay_state_update: + spec_kwargs["old_x"] = self.old_x[layer_offset] + spec_kwargs["old_B"] = self.old_B[layer_offset] + spec_kwargs["old_dt"] = self.old_dt[layer_offset] + spec_kwargs["old_dA_cumsum"] = self.old_dA_cumsum[layer_offset] + spec_kwargs["cache_buf_idx"] = self.cache_buf_idx + spec_kwargs["prev_num_accepted_tokens"] = self.prev_num_accepted_tokens + else: + spec_kwargs["intermediate_ssm"] = self.intermediate_ssm_states[layer_offset] + return PythonMambaCacheManager.SpeculativeState( + conv=conv, + temporal=ssm, + intermediate_conv_window=self.intermediate_conv_states[layer_offset], + **spec_kwargs, + ) + return PythonMambaCacheManager.State(conv=conv, temporal=ssm) + + @property + def use_replay_state_update(self) -> bool: + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata(self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, "prev_num_accepted_tokens", None) + cache_buf_idx = getattr(self, "cache_buf_idx", None) + if ( + not self._use_replay_state_update + or prev_num_accepted_tokens is None + or cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None + ): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size, + ) + + def get_mamba_ssm_cache_dtype(self) -> torch.dtype: + return self.ssm_state_dtype + + def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: + return getattr(self, "mamba_ssm_rand_seed", None) + + @property + def use_kda_replay_update(self) -> bool: + return getattr(self, "_use_kda_replay_update", False) + + +class MixedMambaHybridCacheManager(KVCacheManager, MambaCacheManager, MambaHybridCacheManager): + """Hybrid cache manager combining separate KVCacheManager and MambaCacheManager. + + Manages KV cache and mamba states in independent pools, with support of + speculative decoding and disaggregated serving. + Does not support block reuse / prefix caching for mamba states. + """ + + 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 parameters + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + layer_mask: List[bool], + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + max_seq_len: int, + max_batch_size: int, + mapping: Mapping, + max_num_tokens: int = 8192, + dtype: DataType = DataType.HALF, + spec_config: Optional["DecodingBaseConfig"] = None, + is_estimating_kv_cache: bool = False, + execution_stream: Optional[torch.cuda.Stream] = None, + model_type: str = "nemotron_hybrid", + is_draft: bool = False, + use_replay_state_update: bool = False, + mamba_ssm_stochastic_rounding: bool = False, + kda_replay_num_spec: Optional[int] = None, + # Per-pool configurations forwarded to the C++ KVCacheManager ctor. + # Lets a single manager host pools with mixed shapes (e.g. Gemma4 + # hybrid attention). See KVCacheManager.__init__. + pool_configurations: Optional[List[PoolConfiguration]] = None, + ) -> None: + # mamba hybrid cache requires block reuse to be disabled in KV cache config + warn_legacy_mamba_cache_manager("python", type(self).__name__) + assert not kv_cache_config.enable_block_reuse, ( + "mamba hybrid cache requires block reuse to be disabled in KV cache config" + ) + + # Host-drafter spec modes (NGram) have no spec worker to call + # update_mamba_states; this manager promotes accepted states itself + # in update_resources. One-model modes (MTP/Eagle/DFlash) and the + # suffix-automaton worker promote from their spec workers and must + # NOT be promoted twice. + self._promote_states_in_update_resources = ( + spec_config is not None and getattr(spec_config, "decoding_type", None) == "NGram" + ) + + pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) + + MambaCacheManager.__init__( + self, + mamba_d_state, + mamba_d_conv, + mamba_num_heads, + mamba_n_groups, + mamba_head_dim, + mamba_num_layers, + pool_size, + max_batch_size, + mapping, + mamba_cache_dtype, + mamba_ssm_cache_dtype, + mamba_layer_mask, + execution_stream, + speculative_num_draft_tokens=( + spec_config.tokens_per_gen_step - 1 if spec_config is not None else None + ), + model_type=model_type, + use_replay_state_update=use_replay_state_update, + mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, + kda_replay_num_spec=kda_replay_num_spec, + ) + + # initialize kv cache manager + KVCacheManager.__init__( + self, + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_num_tokens=max_num_tokens, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + execution_stream=execution_stream, + is_draft=is_draft, + pool_configurations=pool_configurations, + ) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + MambaCacheManager.prepare_resources(self, scheduled_batch) + KVCacheManager.prepare_resources(self, scheduled_batch) + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + MambaCacheManager.free_resources(self, request) + KVCacheManager.free_resources(self, request, pin_on_release) + + def add_dummy_requests(self, request_ids: List[int], **kwargs): + MambaCacheManager.add_dummy_requests(self, request_ids) + return KVCacheManager.add_dummy_requests(self, request_ids, **kwargs) + + def shutdown(self): + MambaCacheManager.shutdown(self) + KVCacheManager.shutdown(self) + + def update_resources( + self, + scheduled_batch: ScheduledRequests, + attn_metadata: "AttentionMetadata" = None, + kv_cache_dtype_byte_size: float = None, + ): + KVCacheManager.update_resources( + self, scheduled_batch, attn_metadata, kv_cache_dtype_byte_size + ) + self._maybe_promote_drafter_states(scheduled_batch, attn_metadata) + + def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): + """Commit accepted verification states for host-drafter spec modes. + + One-model spec workers (MTP/Eagle/DFlash) and the suffix-automaton + worker call update_mamba_states themselves right after on-device + acceptance. Host-drafter modes (NGram) have no spec worker: + acceptance lands on the requests as + ``py_num_accepted_draft_tokens`` during sampler update, and this + hook — running right after, alongside the KV rewind — promotes the + accepted step's intermediate state into the live pools (or, on the + KDA replay path, records the accepted-draft count for the next + round's replay). Without it, the pools would keep the + pre-verification state and the next forward would resume from a + stale prefix. + """ + if not self._promote_states_in_update_resources: + return + if not self.is_speculative() or attn_metadata is None: + return + gen_requests = scheduled_batch.generation_requests + if not gen_requests: + return + drafted = [ + r for r in gen_requests if r.py_draft_tokens is not None and len(r.py_draft_tokens) > 0 + ] + if not drafted: + # Drafter skipped this step (e.g. speculation gated off): the + # forward ran the plain in-place decode path; nothing to promote. + return + assert len(drafted) == len(gen_requests), ( + "mixed drafted/undrafted generation batch is not supported for " + "hybrid state promotion (drafts are padded to the static max)" + ) + device = self._impl.mamba_cache.temporal.device + num_contexts = len(scheduled_batch.context_requests) + num_accepted = torch.tensor( + [0] * num_contexts + [r.py_num_accepted_draft_tokens + 1 for r in gen_requests], + dtype=torch.int32, + device=device, + ) + # Batch-ordered slots (contexts then gens), matching the ordering + # the forward used for the intermediate scratch buffers. Requests + # that finished this step were already freed by response handling + # (which runs before update_resources) and are gone from the index; + # their rows must stay in place for alignment, so redirect them to + # the reserved padding slot (a harmless scratch write). + slot_index = self.mamba_cache_index + padding_slot = self._impl._padding_slot + state_indices = torch.tensor( + [ + slot_index.get(r.py_request_id, padding_slot) + for r in scheduled_batch.context_requests + gen_requests + ], + dtype=torch.int32, + device=device, + ) + self.update_mamba_states(attn_metadata, num_accepted, state_indices) + + +class CppMambaHybridCacheManager(KVCacheManager, _LegacyMambaHybridCacheManager): + """Hybrid cache manager storing mamba states inside the KVCacheManager pool. + + Both KV cache blocks and recurrent state blocks are managed by the unified + C++ KVCacheManager, enabling block reuse / prefix caching across attention + and mamba layers. This compatibility manager remains available through the + manager preference override and legacy disaggregated routing. + + """ + + 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, + # Note that max_seq_len is not necessarily equal to kv_cache_config.num_tokens. + # It's derived from the model's BuildConfig for consistency with the C++ backend. + 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, # this is the full attention layer mask + 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: + # 3 kinds of layers: + # 1) Mamba layers (mamba_layer_mask is True) + # 2) Full attention layers (full_attention_layer_mask is True) + # 3) Not managed layers (both masks are False) + warn_legacy_mamba_cache_manager("cpp", type(self).__name__) + 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) + layer_mask = [ + mamba_layer_mask[i] or full_attention_layer_mask[i] for i in range(total_layers) + ] + # PP sharding is done across all layers. + # This is called again in the super().__init__, but we want it to run first + # to set up mtp states before the C++ backend is initialized. + self.pp_layers, _ = get_pp_layers( + mamba_num_layers + num_layers, + mapping, + spec_config=spec_config, + layer_mask=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) + self.requests = [] + # Seed externally visible mamba fields before any early return so that + # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work + # on ranks with no local mamba layers. + self._use_replay_state_update = use_replay_state_update + self._use_gdn_cached_replay_all_layer_commit = ( + use_replay_state_update + and model_type == "qwen3_next" + and self.local_num_mamba_layers > 0 + ) + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update + else None + ) + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None + else None + ) + # Same allocation gate as PythonMambaCacheManager: the rand_seed + # buffer must exist whenever SR can fire, not only on the replay path. + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + # Host-side counter bumped per fresh context-slot assignment; combined + # with the slot index and rank offset to produce reproducible per-slot + # seed values without any torch.randint. + self._seed_request_counter = 0 + self.ssm_state_dtype = mamba_ssm_cache_dtype + # Keep the shared Mamba interface valid on PP ranks that do not own a + # local Mamba layer. + self.spec_config = spec_config + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + + if self.local_num_mamba_layers == 0: + logger.info("No local mamba layers for this rank, skipping mamba cache initialization") + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=num_layers, + num_kv_heads=num_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=full_attention_layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, + **kwargs, + ) + # PP ranks replay the same scheduling decisions, so a rank without + # local Mamba layers must still publish the configured boundaries. + self.kv_cache_config = kv_cache_config + self.linear_attention_metadata = LinearAttentionMetadata() + self.linear_attention_metadata.states_snapshot_interval = ( + kv_cache_config.mamba_state_config.periodic_snapshot_interval + if kv_cache_config.enable_block_reuse + else 0 + ) + return + + # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) + tp_size = _mamba_effective_tp_size(mapping) + 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 + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + self.conv_state_dtype = mamba_cache_dtype + + # Store GLOBAL (pre-TP-division) mamba params for disagg RnnModelConfig. + # d_inner is computed at the top of __init__ before TP division. + d_inner_global = d_inner # = mamba_head_dim * mamba_num_heads (GLOBAL) + conv_dim_global = d_inner_global + 2 * mamba_n_groups * mamba_d_state + self._rnn_d_state = mamba_d_state + self._rnn_d_conv = mamba_d_conv + self._rnn_num_heads = mamba_num_heads # GLOBAL + self._rnn_n_groups = mamba_n_groups + self._rnn_head_dim = mamba_head_dim + self._rnn_hidden_size = d_inner_global # GLOBAL = head_dim * num_heads + self._rnn_conv_dim_size = conv_dim_global # GLOBAL conv dim + self._rnn_num_layers = mamba_num_layers + + # Conv section layout for section-aware split/concat in TP mismatch. + # Section dims are derived from mHiddenSize and mNGroups*mDState in C++; + # we just need to tell C++ which ordering to use. + self._rnn_conv_section_layout = model_type # "nemotron_hybrid" or "qwen3_next" + self.ssm_bytes = math.prod(self.ssm_state_shape) * self.ssm_state_dtype.itemsize + self.conv_bytes = math.prod(self.conv_state_shape) * self.conv_state_dtype.itemsize + + total_bytes = self.ssm_bytes + self.conv_bytes + if total_bytes % self.ssm_state_dtype.itemsize != 0: + raise RuntimeError( + f"Total state bytes ({total_bytes}) not divisible by " + f"ssm_state_dtype size ({self.ssm_state_dtype.itemsize})" + ) + if total_bytes % self.conv_state_dtype.itemsize != 0: + raise RuntimeError( + f"Total state bytes ({total_bytes}) not divisible by " + f"conv_state_dtype size ({self.conv_state_dtype.itemsize})" + ) + if self.ssm_bytes % self.conv_state_dtype.itemsize != 0: + raise RuntimeError( + f"SSM state bytes ({self.ssm_bytes}) not divisible by " + f"conv_state_dtype size ({self.conv_state_dtype.itemsize})" + ) + self.linear_attention_metadata = LinearAttentionMetadata() + self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value + self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes + self.linear_attention_metadata.states_snapshot_interval = ( + kv_cache_config.mamba_state_config.periodic_snapshot_interval + if kv_cache_config.enable_block_reuse + else 0 + ) + # RNN model params for disagg TP-mismatch split/concat. + conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} + self.linear_attention_metadata.rnn_num_heads = self._rnn_num_heads + self.linear_attention_metadata.rnn_head_dim = self._rnn_head_dim + self.linear_attention_metadata.rnn_d_state = self._rnn_d_state + self.linear_attention_metadata.rnn_d_conv = self._rnn_d_conv + self.linear_attention_metadata.rnn_n_groups = self._rnn_n_groups + self.linear_attention_metadata.rnn_conv_section_layout = conv_section_map.get( + self._rnn_conv_section_layout, 0 + ) + self.linear_attention_metadata.rnn_ssm_bytes = self.ssm_bytes + self.linear_attention_metadata.rnn_ssm_dtype_size = self.ssm_state_dtype.itemsize + self.linear_attention_metadata.rnn_conv_dtype_size = self.conv_state_dtype.itemsize + kv_cache_config = kv_cache_config.model_copy(deep=True) + if kv_cache_config.enable_partial_reuse: + logger.warning( + "Partial reuse is not supported for mamba hybrid models, disabling partial reuse" + ) + kv_cache_config.enable_partial_reuse = False + + # Keep the vector in physical global-layer order. Disabled entries are + # placeholders that are never projected into this manager. + kv_cache_config.max_attention_window = [ + LinearCacheType.RECURRENT_STATES.value if mamba_layer_mask[layer_idx] else max_seq_len + for layer_idx in range(len(layer_mask)) + ] + + recurrent_states_window = LinearCacheType.RECURRENT_STATES.value + local_windows = { + recurrent_states_window if mamba_layer_mask[layer_idx] else max_seq_len + for layer_idx in self.pp_layers + } + kwargs["pool_configurations"] = [ + PoolConfiguration( + window_size=window_size, + head_dim=head_dim, + dtype=torch_dtype_to_binding(self.ssm_state_dtype) + if window_size == recurrent_states_window + else dtype, + ) + for window_size in sorted(local_windows) + ] + + # Normalize num_kv_heads to a per-layer list and zero out mamba + # layer positions: those layers carry SSM/conv state instead of KV + # heads, so the parent KV cache should not allocate KV head storage + # for them. + 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) + + # pass remaining arguments to super class + 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=layer_mask, + is_estimating_kv_cache=is_estimating_kv_cache, + is_draft=is_draft, + linear_attention_metadata=self.linear_attention_metadata, + **kwargs, + ) + + assert self.local_num_mamba_layers > 0, "At least one mamba layer is required" + self.mamba_layer_offsets = {} + for idx, layer_id in enumerate(self.mamba_pp_layers): + self.mamba_layer_offsets[layer_id] = idx + + self.host_block_offsets = torch.zeros( + [self.impl.num_pools, self.max_batch_size, 2, self.max_blocks_per_seq], + dtype=torch.int32, + device="cpu", + ) + self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ + self.layer_offsets[self.mamba_pp_layers[0]] + ][0] + + 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() + ) + self._row_indices = torch.arange(self.max_batch_size, dtype=torch.long, device="cpu") + self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} + self.kv_cache_config = kv_cache_config + self.is_estimating_kv_cache = is_estimating_kv_cache + + self._setup_states() + self._setup_replay_buffers(spec_config) + if use_replay_state_update and model_type == "qwen3_next": + logger.info_once( + "Configured GDN cached replay commit mode: small-batch fused, " + "large-batch all-layer", + key="gdn_cached_replay_commit_mode_fused", + ) + + @staticmethod + def get_cache_size_per_token( + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs, + ): + """Affine memory model for the unified hybrid KV pool. + + Returns ``(slope_bytes_per_token, intercept_bytes)``: + + * ``slope`` = attention KV bytes per token plus conservative periodic + Mamba-state and partial-attention-snapshot costs. + * ``intercept`` = live and CUDA-graph dummy Mamba state. + + Memory budget -> max tokens then becomes + ``T = (budget - intercept) // slope`` instead of plain + ``T = budget // bytes_per_token``. + """ + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=1, + include_explicit_snapshots=False, + cap_partial_attention_snapshots=False, + **kwargs, + ) + + @property + def use_gdn_cached_replay_all_layer_commit(self) -> bool: + return self._use_gdn_cached_replay_all_layer_commit + + def _commit_gdn_cached_replay_history_layers( + self, + attn_metadata: "AttentionMetadata", + num_decodes: int, + ) -> None: + """Synchronously advance every local GDN checkpoint in one launch.""" + from tensorrt_llm._torch.modules.fla.cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + commit_gdn_cached_replay_history_layers, + ) + + if ( + not self._use_gdn_cached_replay_all_layer_commit + or num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + ): + return + if ( + self.all_ssm_states is None + or self.old_x is None + or self.old_B is None + or self.old_dt is None + or self.replay_history_size is None + ): + raise RuntimeError("GDN cached replay all-layer commit requires replay state buffers.") + + mamba_metadata = attn_metadata.mamba_metadata + if mamba_metadata.replay_num_decodes != num_decodes: + raise RuntimeError( + "GDN replay metadata contains " + f"{mamba_metadata.replay_num_decodes} decode requests, " + f"but state update received {num_decodes}." + ) + commit_gdn_cached_replay_history_layers( + ssm_states=self.all_ssm_states, + old_u=self.old_x, + old_k=self.old_B, + old_G=self.old_dt, + replay_work_items=mamba_metadata.replay_work_items[:num_decodes], + n_writes=mamba_metadata.replay_n_writes, + history_size=self.replay_history_size, + ) + + def shutdown(self): + # Release tensor views into the pool before the pool memory is freed, + # so their deleters don't see stale pointers. + self.all_ssm_states = None + self.all_conv_states = None + self.intermediate_ssm_states = None + self.intermediate_conv_states = None + self.intermediate_state_indices = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + super().shutdown() + + def add_dummy_requests( + self, + request_ids: List[int], + # Note that token_nums should be past_kv_len + input_len (without + # spec decoding). The draft tokens will be added in this function, + # so we don't need to take care of it in the caller. When preparing + # token_nums, we should not take the draft tokens into account, so + # don't use the kv_cache_manager.max_seq_len, which includes both + # extra tokens and draft tokens. + 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, + draft_kv_cache_manager: Optional[KVCacheManager] = None, + capture_sampling_params: Optional["SamplingParams"] = 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, + draft_kv_cache_manager=draft_kv_cache_manager, + capture_sampling_params=capture_sampling_params, + ) + if requests: + self.requests.extend(requests) + # Only process the newly added requests, not all of self.requests. + # self.requests may contain stale entries from the previous + # _prepare_resources call (e.g., disagg transfer-pending requests) + # that would exceed max_batch_size and cause out-of-bounds access. + self._setup_state_indices(requests) + return requests + + 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) + + @nvtx_range("hybrid_prepare_resources") + def _prepare_resources(self, scheduled_batch: ScheduledRequests): + self.requests = scheduled_batch.context_requests + scheduled_batch.generation_requests + # Issue all async block onboards; defer the syncTransfers() (refresh_blocks) + # until just before forward so the CPU-side prep (attn metadata, input + # tensor assembly, draft model prep, etc.) can overlap with the in-flight + # cudaMemcpyAsync calls instead of being serialized behind them. + # copy_linear_attention_block returns True iff it actually issued a copy; + # we skip refresh_blocks entirely when nothing was scheduled. + self._pending_state_transfers = self.impl.copy_linear_attention_block_batch(self.requests) + self._setup_state_indices() + self._reset_context_mamba_slots(len(scheduled_batch.context_requests)) + + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + self._prepare_resources(scheduled_batch) + + @nvtx_range("hybrid_flush_state_transfers") + def flush_state_transfers(self) -> None: + """Complete any deferred state-block onboards scheduled by + prepare_resources(). Must be called before forward() reads recurrent + state blocks. Cheap no-op when nothing was scheduled. + """ + if getattr(self, "_pending_state_transfers", False): + self.impl.refresh_blocks() + self._pending_state_transfers = False + + @nvtx_range("hybrid_update_mamba_states") + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None, + ): + if self.local_num_mamba_layers == 0: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + accepted_tokens = num_accepted_tokens[num_contexts : num_contexts + num_gens] + # Dynamic tree supplies explicit tree-node leaf positions. Linear MTP + # passes accepted-token counts and lets the promotion kernel subtract + # one, avoiding a separate GPU subtraction/cast. + if accepted_leaf_positions is None: + accepted_position_source = accepted_tokens + position_source_is_token_count = True + else: + accepted_position_source = accepted_leaf_positions.to(torch.int32) + position_source_is_token_count = False + # Match the API of MambaCacheManager.update_mamba_states: callers + # may pass per-request state slot indices explicitly (e.g. MTP via + # attn_metadata.mamba_metadata.state_indices). Fall back to this + # manager's own slot mapping when not provided. + 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] + + # The accepted SSM/conv promotion is a bandwidth-bound gather->scatter + # into the dtype-reinterpreted, strided C++ pool view. torch.compile + # can't handle the dtype-view mutation (and Inductor chokes on the + # uint8 pool with "XBLOCK too large"), so a dedicated Triton kernel + # writes through the view's real strides (~85% of HBM peak, one launch + # per state). + if self._use_replay_state_update: + # Every GDN layer has finished reading the old checkpoint and + # writing its candidate history. Advance all local checkpoints + # in one launch before PNAT and the active history buffer change. + self._commit_gdn_cached_replay_history_layers(attn_metadata, num_gens) + assert self._dummy_request_mask is not None + else: + # Legacy: copy the accepted SSM state from the intermediate buffer. + _promote_mamba_state_triton( + self.all_ssm_states, + self.intermediate_ssm_states, + src_state_indices, + accepted_position_source, + state_indices_d, + position_source_is_token_count=position_source_is_token_count, + ) + + # Conv: both paths save all intermediate conv windows, carry over the + # accepted one. The replay path also advances PNAT and the active cache + # buffer in this same launch. + replay_dummy_request_mask = None + if self._use_replay_state_update: + replay_dummy_request_mask = self._dummy_request_mask[ + num_contexts : num_contexts + num_gens + ] + _promote_mamba_state_triton( + self.all_conv_states, + self.intermediate_conv_states, + src_state_indices, + accepted_position_source, + state_indices_d, + num_accepted_tokens=accepted_tokens, + position_source_is_token_count=position_source_is_token_count, + replay_pnat=(self.prev_num_accepted_tokens if self._use_replay_state_update else None), + replay_cache_buf_idx=(self.cache_buf_idx if self._use_replay_state_update else None), + dummy_request_mask=replay_dummy_request_mask, + replay_step_width=(self.replay_step_width if self._use_replay_state_update else 0), + replay_history_size=(self.replay_history_size if self._use_replay_state_update else 0), + ) + + def get_num_available_tokens( + self, token_num_upper_bound: int, max_num_draft_tokens: int = 0, **kwargs + ) -> int: + # Base bound from attention KV cache pool (the parent's behaviour). + result = super().get_num_available_tokens( + token_num_upper_bound, max_num_draft_tokens, **kwargs + ) + # Also bound by the recurrent-state pool capacity: each request needs + # roughly ceil(N / states_snapshot_interval) recurrent-state blocks + # (+1 for the corner case where N is a multiple of tokens_per_block). + # When block reuse is disabled, only one snapshot is needed per + # request, so no additional capping is required here. + interval = ( + self.linear_attention_metadata.states_snapshot_interval + if self.linear_attention_metadata is not None + else 0 + ) + # Attention-only PP ranks keep the interval so every rank publishes + # identical scheduling boundaries, but they have no recurrent-state + # pool whose capacity should constrain their attention KV cache. + if self.local_num_mamba_layers > 0 and interval and interval > 0: + stats = self.impl.get_kv_cache_stats() + rs_free = stats.num_free_blocks_per_window_size.get( + LinearCacheType.RECURRENT_STATES.value, 0 + ) + # Reserve 1 block for the always-allocated last block (corner case + # / final live state) so we don't promise more tokens than the + # pool can actually back at allocation time. + usable_rs_blocks = max(0, rs_free - 1) + rs_token_cap = usable_rs_blocks * interval + result = min(result, rs_token_cap) + return max(result, 0) + + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + if request in self.requests: + self.requests.remove(request) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) + + def _setup_state_indices(self, requests=None) -> None: + if self.local_num_mamba_layers == 0: + return + if requests is None: + requests = self.requests + block_indices = [] + for req in requests: + if req.is_context_finished: + next_step = self.get_num_tokens(req) - 1 + elif self.kv_cache_config.enable_block_reuse: + next_step = req.context_current_position - 1 + req.context_chunk_size + else: + next_step = req.prompt_len - 1 + block_indices.append(next_step // self.tokens_per_block) + self.impl.copy_batch_block_offsets( + self.host_block_offsets, [req.py_request_id for req in requests], 1, 0 + ) + max_blocks = self.blocks_per_window[LinearCacheType.RECURRENT_STATES.value][0] + n = len(requests) + self._host_state_indices.zero_() + if n > 0: + # Vectorized gather: replace per-element Python loop with a single + # tensor index op. block_indices is a small Python list so + # torch.tensor() conversion here is O(n) at C level, much cheaper + # than n round-trips through Python indexing. + bi = torch.tensor(block_indices, dtype=torch.long) + rows = self._row_indices[:n] + # host_block_offsets: [num_pools, max_batch_size, 2, max_blocks_per_seq] + values = self.host_block_offsets[self.recurrent_states_pool_index, rows, 0, bi] + invalid_mask = (values < 0) | (values >= max_blocks) + # The C++ recurrent-state manager uses null page-table entries for + # logical blocks that are not snapshot boundaries. Usually a + # context chunk ends exactly at an allocated snapshot (or at the + # final live-state block), but the scheduler may shorten a chunk + # further when KV capacity is tight. In that case the Mamba + # kernel must keep accumulating into the next allocated snapshot + # or final block, matching copyLinearAttentionBlock(), which also + # skips placeholders when it advances the live state. + for bad_i in invalid_mask.nonzero(as_tuple=False).flatten().tolist(): + req = requests[bad_i] + if req.is_context_finished: + continue + last_prompt_block = (req.prompt_len - 1) // self.tokens_per_block + row = self.host_block_offsets[self.recurrent_states_pool_index, rows[bad_i], 0] + candidates = row[block_indices[bad_i] : last_prompt_block + 1] + valid_candidates = ((candidates >= 0) & (candidates < max_blocks)).nonzero( + as_tuple=False + ) + if valid_candidates.numel() > 0: + values[bad_i] = candidates[valid_candidates[0, 0]] + + invalid_mask = (values < 0) | (values >= max_blocks) + if invalid_mask.any(): + bad_i = int(invalid_mask.nonzero(as_tuple=False)[0, 0]) + req = requests[bad_i] + value = int(values[bad_i]) + raise RuntimeError( + f"Invalid recurrent state block index {value} " + f"(expected 0 <= index < {max_blocks}) for request {bad_i}, " + f"prompt_len={req.prompt_len}, " + f"is_context_finished={req.is_context_finished}, " + f"context_current_position={req.context_current_position}, " + f"prepopulated_token_num={req.prepopulated_prompt_len}, " + f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " + f"block_index for next step is {block_indices[bad_i]}, " + "\nblock_ids=" + f"{self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" + ) + self._host_state_indices[:n] = values + + self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask(is_dummy) + + # Build request_id → pool block offset mapping so that + # get_state_indices can return indices in arbitrary request order. + # Bulk tolist avoids a per-request tensor-index + .item() round-trip. + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy + + def get_state_indices( + self, request_ids: Optional[List[int]] = None, is_padding: Optional[List[bool]] = None + ) -> list: + if self.local_num_mamba_layers == 0: + # Mamba metadata is prepared on every PP rank even when this rank + # owns only attention layers. No local kernel consumes these + # indices, so avoid consulting state that is intentionally absent. + return [0] * len(request_ids) if request_ids is not None else [] + if request_ids is not None: + # Return indices in the order of the caller's request_ids, + # not the internal self.requests order. This is critical when + # the batch is reordered after prepare_resources (e.g. disagg + # serving sorts generation_requests by py_batch_idx). + indices = [self._request_id_to_state_index[rid] for rid in request_ids] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + return self.cuda_state_indices + + def _setup_states(self) -> None: + # Pool layout: {numLocalLayers, numBlocks, ssm_bytes + conv_bytes} (as uint8) + pool: torch.Tensor = ( + self.impl.get_recurrent_states_pool() + .view(torch.uint8) + .reshape(self.local_num_mamba_layers, -1, self.ssm_bytes + self.conv_bytes) + ) + num_blocks_in_pool = pool.shape[1] + self.all_ssm_states = ( + pool[:, :, : self.ssm_bytes] + .view(self.ssm_state_dtype) + .view([self.local_num_mamba_layers, num_blocks_in_pool] + self.ssm_state_shape) + ) + self.all_conv_states = ( + pool[:, :, self.ssm_bytes : self.ssm_bytes + self.conv_bytes] + .view(self.conv_state_dtype) + .view([self.local_num_mamba_layers, num_blocks_in_pool] + self.conv_state_shape) + ) + self.all_ssm_states.zero_() + self.all_conv_states.zero_() + + def _setup_replay_buffers(self, spec_config) -> None: + cache_size = self.all_ssm_states.shape[1] + device = self.all_ssm_states.device + if not self._allocate_pool_replay_buffers(spec_config, cache_size, device): + return + + self._dummy_request_mask = torch.zeros(self.max_batch_size, dtype=torch.bool, device=device) + self._dummy_request_mask_host = torch.zeros( + self.max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned(), + ) + + +__all__ = [ + "PythonMambaCacheManager", + "MambaCacheManager", + "MixedMambaHybridCacheManager", + "CppMambaHybridCacheManager", + "warn_legacy_mamba_cache_manager", +] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py new file mode 100644 index 000000000000..99260329abb5 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -0,0 +1,1279 @@ +# 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"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +import math +import sys +from dataclasses import replace +from typing import ( + TYPE_CHECKING, + Dict, + Iterable, + List, + Literal, + Optional, + Sequence, + Tuple, + Union, + cast, +) + +import torch + +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + BlockReusePolicy, + KVCacheManagerV2, + Role, +) +from tensorrt_llm._torch.pyexecutor.kv_cache_stats import KVCacheV2IterationStatsReport +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.resource_manager import ( + BaseResourceManager, + CacheTypeCpp, + DataType, + get_pp_layers, +) +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, prefer_pinned +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + DEFAULT_BEAM_INDEX, + BatchDesc, + BufferConfig, + DataRole, + KVCacheDesc, + LayerId, + PageIndexMode, + SsmLayerConfig, + TokenIdExt, + _KVCache, +) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy + +from .common import ( + MambaAcceptanceBatch, + MambaHybridCacheManager, + MambaLayerCache, + MambaRole, + MambaStateLayout, + _estimate_mamba_hybrid_cache_cost, + _get_num_cuda_graph_padding_dummy_slots, + _mamba_effective_tp_size, + _mamba_rank_offset, + _mamba_regular_snapshot_interval, + _mamba_snapshot_rule_counts, + _PrefixReuseDiagnostics, +) + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata + from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig + from tensorrt_llm.sampling_params import SamplingParams + +GB = 1 << 30 + + +class MambaHybridCacheManagerV2(KVCacheManagerV2, MambaHybridCacheManager): + """Hybrid Mamba cache manager backed by KVCacheManagerV2. + + Attention KV pages and Mamba recurrent-state pages are both owned by the + Python V2 cache manager. Mamba layers are represented as V2 SSM layers, + while this wrapper exposes the state tensors and slot indices expected by + the PyTorch Mamba kernels. + """ + + _supports_additional_snapshot_offsets = True + + # Construction and cache configuration + + @override + 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, + mamba_ssm_stochastic_rounding: bool = False, + conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", + **kwargs, + ) -> None: + if conv_state_layout not in ("x_b_c", "q_k_v"): + raise ValueError(f"Unsupported convolution state layout: {conv_state_layout!r}") + if "model_type" in kwargs: + # The V1 managers select the conv-state layout by model_type; this + # class takes it explicitly. Silently absorbing model_type here + # means a caller's layout request would be dropped on the floor. + raise TypeError( + "MambaHybridCacheManagerV2 does not accept 'model_type' " + f"(got {kwargs['model_type']!r}); pass " + "conv_state_layout='x_b_c' or 'q_k_v' instead" + ) + 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.spec_config = spec_config + self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding + self._seed_rank_offset = _mamba_rank_offset(mapping) + self._recurrent_evicted_blocks_total = 0 + self._recurrent_onboarded_blocks_total = 0 + self._recurrent_dropped_blocks_total = 0 + self._recurrent_status_logged = False + # Branch points keyed by request id. prepare_expect_snapshot_points() + # overwrites the per-request list on every scheduler pass, so the point + # has to be re-merged from here rather than stored only on the request. + self._branch_snapshot_points: Dict[int, int] = {} + self._snapshot_pruned_tokens_total = 0 + self._page_pruned_tokens_total = 0 + self._branch_snapshots_taken_total = 0 + self._branch_snapshots_skipped_total: Dict[str, int] = {} + num_cuda_graph_padding_dummy_slots = _get_num_cuda_graph_padding_dummy_slots( + spec_config, max_batch_size + ) + self._num_reserved_dummy_slots = num_cuda_graph_padding_dummy_slots + int( + mapping.enable_attention_dp + ) + self.ssm_state_dtype = ( + mamba_ssm_cache_dtype if mamba_ssm_cache_dtype is not None else mamba_cache_dtype + ) + self.conv_state_dtype = mamba_cache_dtype + self._global_n_groups = mamba_n_groups + + 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 = _mamba_effective_tp_size(mapping) + d_inner = mamba_head_dim * mamba_num_heads + grouped_state_dim = mamba_n_groups * mamba_d_state + conv_dim = d_inner + 2 * grouped_state_dim + nheads = mamba_num_heads + assert nheads % tp_size == 0, "mamba_num_heads must be divisible by tp_size" + assert conv_dim % tp_size == 0, "conv_dim must be divisible by tp_size" + if kwargs.get("is_disagg", False) and grouped_state_dim % tp_size != 0: + raise ValueError( + "Disaggregated Mamba transfer requires each convolution " + "state section to be divisible by tp_size" + ) + self._n_groups_per_rank = mamba_n_groups // tp_size + d_inner_local = d_inner // tp_size + grouped_state_dim_local = grouped_state_dim // tp_size + conv_dim = conv_dim // tp_size + nheads = nheads // tp_size + self.conv_state_shape = [conv_dim, mamba_d_conv - 1] + self.ssm_state_shape = [nheads, mamba_head_dim, mamba_d_state] + # TP-mismatch disaggregated transfers split the flat convolution + # state at the semantic boundaries selected by the layout. + if conv_state_layout == "x_b_c": + self.conv_section_dims = [ + d_inner_local, + grouped_state_dim_local, + grouped_state_dim_local, + ] + else: + self.conv_section_dims = [ + grouped_state_dim_local, + grouped_state_dim_local, + d_inner_local, + ] + self.ssm_bytes = math.prod(self.ssm_state_shape) * self.ssm_state_dtype.itemsize + self.conv_bytes = math.prod(self.conv_state_shape) * self.conv_state_dtype.itemsize + else: + logger.info("No local mamba layers for this rank, skipping mamba state views") + self._n_groups_per_rank = 0 + self.conv_state_shape = [] + self.ssm_state_shape = [] + self.conv_section_dims = [] + self.ssm_bytes = 0 + self.conv_bytes = 0 + + self._state_layout = MambaStateLayout( + layer_mask=tuple(self._mamba_layer_mask), + pp_layers=tuple(self.pp_layers), + mamba_pp_layers=tuple(self.mamba_pp_layers), + mapping=mapping, + conv_state_shape=tuple(self.conv_state_shape), + ssm_state_shape=tuple(self.ssm_state_shape), + conv_state_dtype=self.conv_state_dtype, + ssm_state_dtype=self.ssm_state_dtype, + n_groups_per_rank=self._n_groups_per_rank, + max_batch_size=max_batch_size, + state_index_capacity=(max_batch_size + self._num_reserved_dummy_slots), + slot_capacity=None, + spec_config=spec_config, + stochastic_rounding=mamba_ssm_stochastic_rounding, + seed_rank_offset=self._seed_rank_offset, + conv_section_dims=tuple(self.conv_section_dims), + conv_state_layout=conv_state_layout, + ) + self._initialize_model_state() + + 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 + + kv_cache_config = kv_cache_config.model_copy(deep=True) + if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse: + block_reuse_config = kv_cache_config.block_reuse_config + block_reuse_policy = BlockReusePolicy(block_reuse_config.policy) + if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE: + # SSM reuse is valid only at explicit snapshot boundaries. + kv_cache_config.block_reuse_config = block_reuse_config.model_copy( + update={"policy": BlockReusePolicy.PER_REQUEST.value} + ) + self.kv_cache_config = kv_cache_config + + super().__init__( + kv_cache_config, + kv_cache_type, + num_layers=mamba_num_layers + num_layers, + num_kv_heads=per_layer_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=dtype, + spec_config=spec_config, + layer_mask=combined_layer_mask, + is_draft=is_draft, + is_estimating_kv_cache=is_estimating_kv_cache, + num_reserved_index_slots=self._num_reserved_dummy_slots, + **kwargs, + ) + + self._setup_state_pool() + + def _initialize_model_state(self) -> None: + """Initialize and validate model-owned state before constructing the pool.""" + if self.spec_config is not None: + raise NotImplementedError( + "Speculative decoding requires a model-specific Mamba cache manager" + ) + + @classmethod + @override + def get_cache_size_per_token( + cls, + model_config, + mapping: Mapping, + *, + max_batch_size: int, + kv_cache_config: KvCacheConfig, + num_layers: Optional[int] = None, + tokens_per_block: int = 32, + max_seq_len: Optional[int] = None, + **kwargs, + ): + spec_config = kwargs.get("spec_config") + num_reserved_dummy_slots = _get_num_cuda_graph_padding_dummy_slots( + spec_config, max_batch_size + ) + int(mapping.enable_attention_dp) + return _estimate_mamba_hybrid_cache_cost( + model_config, + mapping, + max_batch_size=max_batch_size, + kv_cache_config=kv_cache_config, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + num_reserved_dummy_slots=num_reserved_dummy_slots, + extra_state_bytes_per_rank=cls._extra_state_bytes_per_rank( + model_config, mapping, **kwargs + ), + include_explicit_snapshots=True, + cap_partial_attention_snapshots=True, + **kwargs, + ) + + @classmethod + def _extra_state_bytes_per_rank(cls, model_config, mapping: Mapping, **kwargs) -> int: + return 0 + + def _extra_scratch_bytes_per_slot(self) -> int: + return 0 + + def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConfig]: + """Declare all buffers sharing one recurrent layer's slot lifecycle.""" + return [ + BufferConfig(role=MambaRole.SSM_STATE, size=self.ssm_bytes), + BufferConfig(role=MambaRole.CONV_STATE, size=self.conv_bytes), + ] + + def _is_local_mamba_layer(self, local_layer_idx: int) -> bool: + return self._mamba_layer_mask[self.pp_layers[local_layer_idx]] + + @override + def _get_pool_roles(self, pool_id: int) -> Tuple[DataRole, Optional[DataRole]]: + layer_id = int(self.impl.layer_grouping[pool_id][0]) + if self._is_local_mamba_layer(layer_id): + return MambaRole.SSM_STATE, None + return super()._get_pool_roles(pool_id) + + def _max_resident_sequences(self) -> int: + return self.max_batch_size * self.mapping.pp_size + + def _mamba_state_bytes_per_slot(self) -> int: + return ( + sum( + buffer.size + for layer_id in self.mamba_pp_layers + for buffer in self._get_recurrent_buffer_configs(layer_id) + ) + + self._extra_scratch_bytes_per_slot() + ) + + def _num_ssm_snapshots_for_capacity( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + if capacity <= 0 or not kv_cache_config.enable_block_reuse: + return 0 + + fixed_rules, _ = _mamba_snapshot_rule_counts( + kv_cache_config, self.max_seq_len, self.tokens_per_block + ) + interval = _mamba_regular_snapshot_interval(kv_cache_config, self.max_seq_len) + regular_snapshots = capacity // interval if interval is not None else 0 + return self._max_resident_sequences() * fixed_rules + regular_snapshots + + def _num_ssm_states_per_typical_request( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> int: + fixed_rules, _ = _mamba_snapshot_rule_counts( + kv_cache_config, + capacity, + self.tokens_per_block, + ) + # Additional snapshots are stable boundaries that must remain alive. + # Periodic snapshots are evictable cache entries and therefore do not + # increase the guaranteed state count represented by BatchDesc. + return 1 + fixed_rules + + def _typical_request_descs( + self, + capacity: int, + kv_cache_config: KvCacheConfig, + ) -> List[KVCacheDesc]: + """Model one request with one descriptor per live SSM state.""" + num_states = self._num_ssm_states_per_typical_request(capacity, kv_cache_config) + capacity_per_state, capacity_remainder = divmod(capacity, num_states) + capacities = [capacity_per_state + int(i < capacity_remainder) for i in range(num_states)] + return [ + KVCacheDesc( + capacity=state_capacity, + history_length=max(0, state_capacity - 1), + ) + for state_capacity in capacities + ] + + def _get_typical_request_capacity( + self, + kv_cache_config: KvCacheConfig, + ) -> int: + if kv_cache_config.avg_seq_len is not None: + return kv_cache_config.avg_seq_len + + fallback_capacity = max(1, self.max_seq_len // 2) + logger.warning( + "'kv_cache_config.avg_seq_len' is not set for a hybrid Mamba " + "model using KV cache manager V2. Falling back to " + f"max_seq_len / 2={fallback_capacity} for cache-pool sizing. Set " + "'kv_cache_config.avg_seq_len' in the YAML configuration to the " + "workload's average total sequence length for an accurate KV/SSM " + "pool ratio." + ) + return fallback_capacity + + @override + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + attention_quota = super()._get_quota_from_max_tokens(max_tokens) + num_request_lineages = self._max_resident_sequences() + snapshot_slots = self._num_ssm_snapshots_for_capacity(max_tokens, self.kv_cache_config) + state_slots = num_request_lineages + self._num_reserved_dummy_slots + snapshot_slots + state_quota = state_slots * self._mamba_state_bytes_per_slot() + # Once the plan contains any non-live SSM capacity, reserve one partial + # attention page per request lineage. This remains conservative when + # the plan contains fewer than one non-live slot per lineage. + extra_attention_quota = ( + num_request_lineages * self._attention_cache_bytes_per_token() * self.tokens_per_block + if snapshot_slots > 0 + else 0 + ) + return attention_quota + state_quota + extra_attention_quota + + @override + def _get_max_tokens_from_quota(self, quota: int) -> float: + if self._get_quota_from_max_tokens(0) > quota: + return 0 + + low = 0 + high = 1 + while self._get_quota_from_max_tokens(high) <= quota: + low = high + high *= 2 + if high >= 1 << 62: + return float("inf") + + while low + 1 < high: + mid = (low + high) // 2 + if self._get_quota_from_max_tokens(mid) <= quota: + low = mid + else: + high = mid + return low + + def _minimum_live_gpu_quota(self) -> int: + """Return the minimum quota for live states and one attention page.""" + attention_block_quota = self._attention_cache_bytes_per_token() * self.tokens_per_block + num_state_slots = self._max_resident_sequences() + self._num_reserved_dummy_slots + state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + return max( + self._get_quota_from_max_tokens(0), + state_quota + attention_block_quota, + ) + + @override + def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: + kv_cache_config = self.kv_cache_config + cache_tiers = config.cache_tiers + gpu_quota = cache_tiers[0].quota + minimum_live_quota = self._minimum_live_gpu_quota() + if minimum_live_quota > gpu_quota: + raise ValueError( + "The V2 Mamba GPU cache quota is too small for live recurrent " + f"states and attention pages: got {gpu_quota} bytes, need at " + f"least {minimum_live_quota} bytes." + ) + # _build_base_config already constructed every attention layer, + # including dtype-specific scale and subclass-provided side buffers. + # Preserve those configs and replace only the local Mamba layers. + layers = list(config.layers) + # Register pool-owned roles here, then bind their typed views in + # _get_state_buffer. Model scratch must not participate in warmup cleanup. + self._recurrent_buffers: dict[tuple[int, DataRole], torch.Tensor | None] = {} + for local_layer_idx, global_layer_idx in enumerate(self.pp_layers): + if self._mamba_layer_mask[global_layer_idx]: + layer_id = LayerId(local_layer_idx) + buffers = self._get_recurrent_buffer_configs(global_layer_idx) + for buffer in buffers: + self._recurrent_buffers[local_layer_idx, buffer.role] = None + layers[local_layer_idx] = SsmLayerConfig( + layer_id=layer_id, + buffers=buffers, + ) + + dummy_requests = [ + KVCacheDesc(capacity=0, history_length=0) for _ in range(self._num_reserved_dummy_slots) + ] + constraints = [ + replace( + batch, + kv_caches=[*batch.kv_caches, *dummy_requests], + ) + for batch in config.constraints + ] + + typical_step = config.typical_step + if config.initial_pool_ratio is None: + typical_capacity = self._get_typical_request_capacity(kv_cache_config) + request_descs = self._typical_request_descs(typical_capacity, kv_cache_config) + typical_step = BatchDesc( + request_descs * self._max_resident_sequences() + dummy_requests + ) + # The recurrent (SSM) state pool must hold one slot per resident + # sequence plus every reserved dummy slot. Unlike attention pages, a + # Mamba state is fixed-size per sequence, so this floor is independent + # of sequence length. The base config only emits constraints when + # ``avg_seq_len`` is set, and speculative decoding inflates the reserved + # dummy slots (CUDA-graph padding), so without an explicit floor the SSM + # pool can be undersized (see the live/dummy-slot check in _setup_states + # / __init__). Add a min-slots constraint of zero-capacity requests: + # these cost no attention pages but reserve one SSM slot each. + if any(isinstance(layer, SsmLayerConfig) for layer in layers): + ssm_floor_slots = self._max_resident_sequences() + self._num_reserved_dummy_slots + constraints = [ + *constraints, + BatchDesc( + [KVCacheDesc(capacity=0, history_length=0) for _ in range(ssm_floor_slots)] + ), + ] + return replace( + config, + layers=layers, + typical_step=typical_step, + constraints=constraints, + # SSM lifecycles require minimum-snapshot commit semantics. The + # flag is harmless when reuse is disabled because no commits are + # attempted, while the runtime config still needs the invariant. + commit_min_snapshot=True, + ) + + def _setup_state_pool(self) -> None: + """Resolve logical slots, bind base views, then initialize model state.""" + self.mamba_layer_offsets = { + layer_id: offset for offset, layer_id in enumerate(self.mamba_pp_layers) + } + self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} + + state_index_capacity = self.max_batch_size + self._num_reserved_dummy_slots + self.cuda_state_indices = torch.zeros( + [state_index_capacity], dtype=torch.int32, device="cuda" + ) + self._host_state_indices = torch.zeros( + [state_index_capacity], dtype=torch.int32, pin_memory=prefer_pinned() + ) + + self._dummy_request_mask = torch.zeros( + state_index_capacity, dtype=torch.bool, device="cuda" + ) + self._dummy_request_mask_host = torch.zeros( + state_index_capacity, dtype=torch.bool, pin_memory=prefer_pinned() + ) + self._generation_state_indices = torch.arange( + self.max_batch_size, dtype=torch.int32, device="cuda" + ) + + if self.local_num_mamba_layers > 0: + first_mamba_local_layer = self.layer_offsets[self.mamba_pp_layers[0]] + self.ssm_layer_group_id = self.impl.get_layer_group_id(LayerId(first_mamba_local_layer)) + self._ssm_page_index_scale = self.impl.get_page_index_scale( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE + ) + num_ssm_pages = self.impl.get_page_index_upper_bound( + LayerId(first_mamba_local_layer), MambaRole.SSM_STATE + ) + num_ssm_slots = ( + num_ssm_pages + self._ssm_page_index_scale - 1 + ) // self._ssm_page_index_scale + required_live_slots = self._max_resident_sequences() + self._num_reserved_dummy_slots + if num_ssm_slots < required_live_slots: + KVCacheManagerV2.shutdown(self) + raise ValueError( + "The V2 Mamba state pool has only " + f"{num_ssm_slots} slots but needs at least " + f"{required_live_slots} live/dummy slots. Increase the " + "KV cache budget or allocate a larger Mamba pool_ratio." + ) + self._state_layout = self._state_layout.with_slot_capacity(num_ssm_slots) + self._setup_states() + else: + self.ssm_layer_group_id = None + self._ssm_page_index_scale = 1 + self.all_ssm_states = [] + self.all_conv_states = [] + self._state_layout = self._state_layout.with_slot_capacity(0) + + self._setup_model_state() + + def _get_state_buffer( + self, local_layer_idx: int, role, dtype: torch.dtype, state_shape: List[int] + ) -> torch.Tensor: + key = (local_layer_idx, role) + if key not in self._recurrent_buffers: + raise ValueError( + f"Recurrent buffer was not registered: layer={local_layer_idx}, role={role}" + ) + addr = self.impl.get_mem_pool_base_address( + LayerId(local_layer_idx), role, PageIndexMode.SHARED + ) + num_pages = self.impl.get_page_index_upper_bound(LayerId(local_layer_idx), role) + raw = convert_to_torch_tensor(TensorWrapper(addr, dtype, [num_pages] + state_shape)) + page_index_scale = self.impl.get_page_index_scale(LayerId(local_layer_idx), role) + num_slots = (num_pages + page_index_scale - 1) // page_index_scale + # V2 coalesces same-size per-layer buffers inside each slot. Kernels + # index Mamba states by logical slot id, so expose only this layer's + # sub-page from each coalesced slot instead of the raw page-index view. + view = raw.as_strided( + [num_slots] + state_shape, + [raw.stride(0) * page_index_scale] + list(raw.stride()[1:]), + ) + self._recurrent_buffers[key] = view + return view + + def _setup_states(self) -> None: + local_layer_ids = [self.layer_offsets[layer_id] for layer_id in self.mamba_pp_layers] + self.all_ssm_states = [ + self._get_state_buffer( + local_layer_idx, MambaRole.SSM_STATE, self.ssm_state_dtype, self.ssm_state_shape + ) + for local_layer_idx in local_layer_ids + ] + self.all_conv_states = [ + self._get_state_buffer( + local_layer_idx, MambaRole.CONV_STATE, self.conv_state_dtype, self.conv_state_shape + ) + for local_layer_idx in local_layer_ids + ] + + def _setup_model_state(self) -> None: + """Bind views only after layout, pool and slot capacity are final.""" + + @override + def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: + return MambaLayerCache( + conv=self.get_conv_states(layer_idx), temporal=self.get_ssm_states(layer_idx) + ) + + def _attention_cache_bytes_per_token(self) -> int: + # Mamba layers have zero KV heads, so the generic calculation naturally + # returns only bytes owned by local attention layers. + return super().get_cache_bytes_per_token() + + @override + def get_cache_bytes_per_token(self) -> int: + cache_bytes = self._attention_cache_bytes_per_token() + + interval = self.kv_cache_config.mamba_state_config.periodic_snapshot_interval + if self.kv_cache_config.enable_block_reuse and interval is not None and interval > 0: + cache_bytes += self._mamba_state_bytes_per_slot() // interval + if cache_bytes == 0 and self.local_num_mamba_layers > 0: + cache_bytes = self._mamba_state_bytes_per_slot() + return max(1, cache_bytes) + + @override + def get_num_free_blocks(self) -> int: + assert len(self.kv_cache_map) == 0, ( + "get_num_free_blocks is only used when the kv cache manager is empty" + ) + attention_pages = [] + ssm_pages = [] + for local_layer_idx in range(self.num_local_layers): + layer_id = LayerId(local_layer_idx) + if self._is_local_mamba_layer(local_layer_idx): + ssm_pages.append( + self.impl.get_page_index_upper_bound(layer_id, MambaRole.SSM_STATE) + // self._ssm_page_index_scale + ) + else: + attention_pages.append( + self.impl.get_page_index_upper_bound(layer_id, Role.KEY) // self.kv_factor + ) + if attention_pages: + return max(attention_pages) + return max(ssm_pages) if ssm_pages else 0 + + @property + @override + 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 + + @override + 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) + + @override + def _create_kv_cache( + self, + request_id: int, + lora_task_id: Optional[int], + input_tokens: Optional[Sequence[TokenIdExt]], + *, + cache_salt: Optional[str] = None, + is_dummy: bool = False, + enable_request_stats: bool = False, + expected_prompt_length: Optional[int] = None, + ) -> Optional[_KVCache]: + kv_cache = super()._create_kv_cache( + request_id, + lora_task_id, + input_tokens, + cache_salt=cache_salt, + is_dummy=is_dummy, + enable_request_stats=enable_request_stats, + expected_prompt_length=expected_prompt_length, + ) + if ( + self.mapping.rank == 0 + and kv_cache is not None + and input_tokens is not None + and not is_dummy + and self.local_num_mamba_layers > 0 + ): + # Prefix lookup excludes the final prompt token because it must be + # recomputed by prefill. Restore it for the request-length metric. + request_total_tokens = len(input_tokens) + 1 + prefix_reuse_diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) + logger.debug( + f"[MambaHybridCacheManagerV2] prefix reuse rank={self.mapping.rank} " + f"request_id={request_id} " + f"request_total_tokens={request_total_tokens} " + f"longest_attention_match_tokens=" + f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_hybrid_pruning()} " + f"content_divergence_tokens=" + f"{prefix_reuse_diagnostics._get_num_reusable_tokens_before_pruning()} " + f"latest_recurrent_snapshot_tokens=" + f"{kv_cache.num_committed_tokens}" + ) + return kv_cache + + @override + 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, + draft_kv_cache_manager: Optional[BaseResourceManager] = None, + capture_sampling_params: Optional["SamplingParams"] = 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, + draft_kv_cache_manager=draft_kv_cache_manager, + capture_sampling_params=capture_sampling_params, + ) + if requests and prepare_resource: + self._setup_state_indices( + requests, + num_contexts=0 if is_gen else len(requests), + ) + if not is_gen: + self._reset_context_mamba_slots(len(requests)) + return requests + + @override + def prepare_resources(self, scheduled_batch: ScheduledRequests): + super().prepare_resources(scheduled_batch) + if self.local_num_mamba_layers == 0: + return + requests = scheduled_batch.context_requests + scheduled_batch.generation_requests + num_contexts = len(scheduled_batch.context_requests) + self._setup_state_indices(requests, num_contexts=num_contexts) + self._reset_context_mamba_slots(num_contexts) + + def _setup_state_indices( + self, + requests: List[LlmRequest], + num_contexts: int = 0, + ) -> None: + if self.local_num_mamba_layers == 0: + return + generation_requests = requests[num_contexts:] + old_state_values = [ + self._request_id_to_state_index.get(request.py_request_id, -1) + for request in generation_requests + ] + n = len(requests) + assert n <= self._host_state_indices.shape[0], ( + f"State-index batch size {n} exceeds max_batch_size {self._host_state_indices.shape[0]}" + ) + self._host_state_indices.zero_() + if n > 0: + for i, req in enumerate(requests): + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + raise RuntimeError(f"Missing V2 KV cache for request {req.py_request_id}") + base_index = kv_cache.get_ssm_block_base_index(self.ssm_layer_group_id) + if base_index < 0: + raise RuntimeError( + f"Invalid SSM state block index {base_index} for " + f"request {req.py_request_id}" + ) + self._host_state_indices[i] = base_index + + self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) + is_dummy = [req.is_dummy for req in requests] + self._refresh_dummy_request_mask(is_dummy) + state_values = self._host_state_indices[:n].tolist() + for req, value, dummy in zip(requests, state_values, is_dummy): + self._request_id_to_state_index[req.py_request_id] = value + self._request_id_to_is_dummy[req.py_request_id] = dummy + + new_state_values = [ + self._request_id_to_state_index[request.py_request_id] + for request in generation_requests + ] + self._on_state_slots_relocated(old_state_values, new_state_values) + + def _on_state_slots_relocated(self, old_slots: list[int], new_slots: list[int]) -> None: + """Allow model-owned scratch state to follow the updated slot mapping.""" + + def _reset_context_mamba_slots(self, num_contexts: int) -> None: + if num_contexts: + self._reset_model_slots( + self.cuda_state_indices[:num_contexts].long(), + self._host_state_indices[:num_contexts].tolist(), + ) + + def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + """Reset any model-owned request state for the supplied slots.""" + + def _refresh_dummy_request_mask(self, is_dummy: list[bool]) -> None: + count = len(is_dummy) + if count > self._dummy_request_mask_host.shape[0]: + raise ValueError("Dummy-request batch exceeds state-index capacity") + self._dummy_request_mask_host.zero_() + if count: + self._dummy_request_mask_host[:count].copy_(torch.tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, non_blocking=True) + + @override + def get_state_indices( + self, request_ids: Optional[List[int]] = None, is_padding: Optional[List[bool]] = None + ): + if self.local_num_mamba_layers == 0: + # Mamba metadata is still prepared on attention-only PP ranks, + # but no local kernel consumes these indices. Return harmless + # placeholders instead of consulting an intentionally empty map. + if request_ids is not None: + return [0] * len(request_ids) + return self.cuda_state_indices + if request_ids is not None: + indices = [self._request_id_to_state_index[rid] for rid in request_ids] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + return self.cuda_state_indices + + @override + def get_max_resource_count(self) -> int: + return self.max_batch_size + + @override + def free_resources(self, request: LlmRequest, pin_on_release: bool = False): + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is not None and kv_cache.is_active: + self.try_commit_blocks(request, kv_cache) + self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) + self._branch_snapshot_points.pop(request.py_request_id, None) + super().free_resources(request, pin_on_release) + + def update_mamba_states( + self, + attn_metadata: "AttentionMetadata", + num_accepted_tokens: torch.Tensor, + state_indices: Optional[torch.Tensor] = None, + accepted_leaf_positions: Optional[torch.Tensor] = None, + ): + if self.local_num_mamba_layers == 0: + return + batch_size = attn_metadata.num_seqs + num_contexts = attn_metadata.num_contexts + num_gens = batch_size - num_contexts + num_accepted_draft_tokens = ( + num_accepted_tokens[num_contexts : num_contexts + num_gens] - 1 + ).to(torch.int32) + # Dynamic tree selects a tree node rather than a linear draft depth. + accepted_positions = ( + accepted_leaf_positions.to(torch.int32) + if accepted_leaf_positions is not None + else num_accepted_draft_tokens + ) + if state_indices is None: + state_indices = self.get_state_indices() + state_indices_d = state_indices[num_contexts : num_contexts + num_gens].to(torch.int32) + src_state_indices = self._generation_state_indices[:num_gens] + is_dummy_request = self._dummy_request_mask[num_contexts : num_contexts + num_gens] + self._update_speculative_state( + MambaAcceptanceBatch( + attention_metadata=attn_metadata, + num_contexts=num_contexts, + num_generations=num_gens, + num_accepted_tokens=num_accepted_tokens[num_contexts : num_contexts + num_gens], + accepted_positions=accepted_positions, + source_state_indices=src_state_indices, + destination_state_indices=state_indices_d, + is_dummy_request=is_dummy_request, + ), + ) + + def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: + raise NotImplementedError("The model manager must implement speculative state updates") + + @override + def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: + should_block_reuse = ( + self.enable_block_reuse and not self.is_draft and not request.is_dummy_request + ) + if not should_block_reuse: + return + + if kv_cache is None: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None: + return + + snapshot_points = request.expect_snapshot_points + commit_limit = ( + min(max(snapshot_points), request.prompt_len) if snapshot_points else request.prompt_len + ) + commit_end = min(request.context_current_position, commit_limit) + if ( + request.context_current_position in request.expect_snapshot_points + and commit_end > kv_cache.num_committed_tokens + ): + tokens = self._augment_tokens_for_block_reuse( + request.get_tokens(DEFAULT_BEAM_INDEX), + request, + start=kv_cache.num_committed_tokens, + end=commit_end, + ) + kv_cache.commit(tokens) + if ( + self.local_num_mamba_layers > 0 + and self._branch_snapshot_points.get(request.py_request_id) == commit_end + ): + self._branch_snapshots_taken_total += 1 + if request.context_current_position >= commit_limit: + self._mark_context_position_as_history(request, kv_cache) + if request.context_remaining_length == 0: + kv_cache.stop_committing() + + def _mark_context_position_as_history(self, request: LlmRequest, kv_cache) -> None: + """Advance history without making later recurrent state reusable.""" + history_length = request.context_current_position + if history_length <= kv_cache.history_length: + return + capacity = max(kv_cache.capacity, history_length) + if not kv_cache.resize(capacity, history_length=history_length): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to {history_length} tokens" + ) + + @override + def update_context_resources(self, scheduled_batch: ScheduledRequests) -> None: + for request in scheduled_batch.context_requests: + kv_cache = self.kv_cache_map.get(request.py_request_id) + if kv_cache is None or not kv_cache.is_active: + continue + + should_block_reuse = ( + self.enable_block_reuse and not self.is_draft and not request.is_dummy_request + ) + is_all_reusable = self.block_reuse_policy == BlockReusePolicy.ALL_REUSABLE + is_snapshot_boundary = ( + request.context_current_position in request.expect_snapshot_points + ) + has_pending_snapshot = any( + point > request.context_current_position for point in request.expect_snapshot_points + ) + should_resize = not should_block_reuse or ( + not is_all_reusable and not has_pending_snapshot + ) + should_commit = ( + is_all_reusable or is_snapshot_boundary or request.context_remaining_length == 0 + ) + + if should_resize and not kv_cache.resize(None, request.context_current_position): + raise ValueError( + "Failed to resize history length of V2 Mamba cache for " + f"request {request.py_request_id} to " + f"{request.context_current_position} tokens at context " + "update" + ) + if should_commit: + self.try_commit_blocks(request, kv_cache) + if request.context_remaining_length == 0: + if self.conversation_manager is not None: + self.conversation_manager.save_drop_plan(request, kv_cache) + kv_cache.enable_swa_scratch_reuse = False + + @override + def prepare_expect_snapshot_points(self, requests: List[LlmRequest]) -> None: + super().prepare_expect_snapshot_points(requests) + if ( + not self.enable_block_reuse + or not self.kv_cache_config.mamba_state_config.enable_branch_snapshot + ): + return + for request in requests: + self._apply_branch_snapshot_point(request) + + def _skip_branch_snapshot(self, reason: str) -> None: + self._branch_snapshots_skipped_total[reason] = ( + self._branch_snapshots_skipped_total.get(reason, 0) + 1 + ) + + @override + def _record_branch_snapshot_point( + self, req: LlmRequest, kv_cache: _KVCache, num_lookup_tokens: Optional[int] + ) -> None: + """Snapshot where this request leaves the reuse tree, for its siblings.""" + if not self.kv_cache_config.mamba_state_config.enable_branch_snapshot: + return + if req.is_dummy_request: + return + if num_lookup_tokens is None: + self._skip_branch_snapshot("no_fresh_match") + # A fresh lookup can record a point before resume() fails. On the + # next attempt, reapply that point against the cache's authoritative + # reuse depth: the request cursor is not rewound until preparation + # succeeds and may still describe a deeper, abandoned attempt. + if req.py_request_id in self._branch_snapshot_points: + self._apply_branch_snapshot_point( + req, context_current_position=kv_cache.num_committed_tokens + ) + return + + diagnostics = cast(_PrefixReuseDiagnostics, kv_cache) + divergence = diagnostics._get_num_reusable_tokens_before_pruning() + reused = kv_cache.num_committed_tokens + if self.local_num_mamba_layers > 0: + hybrid_depth = diagnostics._get_num_reusable_tokens_before_hybrid_pruning() + # Loss attributable to a missing recurrent snapshot, versus loss + # attributable to attention pages having been evicted. These call + # for different fixes, so they are counted apart. + self._snapshot_pruned_tokens_total += max(0, hybrid_depth - reused) + self._page_pruned_tokens_total += max(0, divergence - hybrid_depth) + + # The whole lookup range matched, so there is no fork here. + if divergence >= num_lookup_tokens: + self._skip_branch_snapshot("no_divergence") + return + # Align down: tokens past a block boundary come from a partial-match + # child rather than a confirmed shared prefix, and an unaligned commit + # takes the partial-block snapshot path. + point = (divergence // self.tokens_per_block) * self.tokens_per_block + if point == 0: + self._skip_branch_snapshot("aligned_to_zero") + return + # Reuse already reached the fork, so a snapshot there exists. + if point <= reused: + self._skip_branch_snapshot("already_reused") + return + # A branch point at or beyond the prompt end cannot improve sibling + # reuse. + if point >= req.prompt_len: + self._skip_branch_snapshot("at_prompt_end") + return + + self._branch_snapshot_points[req.py_request_id] = point + self._apply_branch_snapshot_point(req, context_current_position=reused) + + def _apply_branch_snapshot_point( + self, + request: LlmRequest, + context_current_position: Optional[int] = None, + ) -> None: + points = set(request.expect_snapshot_points) + # Keep the flag usable without any configured snapshot placement. If a + # placement is configured, preserve it: adding prompt_len can replace a + # nearby recurrent checkpoint in the same radix-tree block because a + # block stores one recurrent page per life cycle. + if not points: + points.add(request.prompt_len) + point = self._branch_snapshot_points.get(request.py_request_id) + current_position = ( + request.context_current_position + if context_current_position is None + else context_current_position + ) + if point is not None and point > current_position: + points.add(point) + request.expect_snapshot_points = sorted(points) + + def _format_branch_snapshot_counters(self) -> str: + """Attribute reuse loss between missing snapshots and evicted pages. + + A flat cache hit rate means something different depending on whether + branch points were never found or were found and not reused, so the + skip reasons are reported alongside the totals. + """ + skipped = self._branch_snapshots_skipped_total + return ( + f"snapshot_pruned_tokens={self._snapshot_pruned_tokens_total} " + f"page_pruned_tokens={self._page_pruned_tokens_total} " + f"branch_snapshots_taken={self._branch_snapshots_taken_total} " + f"branch_snapshots_skipped={sum(skipped.values())} " + f"branch_snapshots_skipped_by_reason=" + f"{dict(sorted(skipped.items()))}" + ) + + @override + def get_iteration_stats(self) -> Optional[KVCacheV2IterationStatsReport]: + """Log recurrent-cache movement across V2 storage tiers.""" + report = super().get_iteration_stats() + if report is None: + return None + + pool_group_ids = sorted( + { + pool_group_id + for pool_group_id, _, kind in self._stats_life_cycle_metadata().values() + if kind == "ssm" + } + ) + if not pool_group_ids: + return report + pool_group_reports = [ + report.by_pool_group[pool_group_id] for pool_group_id in pool_group_ids + ] + + stats = [pool_group_report.stats for pool_group_report in pool_group_reports] + evicted_blocks = sum(stat.iter_offload_blocks for stat in stats) + evicted_bytes = sum(stat.iter_offload_bytes for stat in stats) + onboarded_blocks = sum(stat.iter_onboard_blocks for stat in stats) + onboarded_bytes = sum(stat.iter_onboard_bytes for stat in stats) + dropped_blocks = sum(stat.iter_host_dropped_blocks for stat in stats) + dropped_bytes = sum(stat.iter_host_dropped_bytes for stat in stats) + has_movement = bool(evicted_blocks or onboarded_blocks or dropped_blocks) + if has_movement: + self._recurrent_evicted_blocks_total += evicted_blocks + self._recurrent_onboarded_blocks_total += onboarded_blocks + self._recurrent_dropped_blocks_total += dropped_blocks + if self.mapping.rank == 0 and (has_movement or not self._recurrent_status_logged): + logger.debug( + f"[MambaHybridCacheManagerV2] recurrent cache status " + f"rank={self.mapping.rank} pool_group_ids={pool_group_ids} " + f"evicted_recurrent_blocks={evicted_blocks} " + f"evicted_recurrent_bytes={evicted_bytes} " + f"onboarded_recurrent_blocks={onboarded_blocks} " + f"onboarded_recurrent_bytes={onboarded_bytes} " + f"dropped_recurrent_blocks={dropped_blocks} " + f"dropped_recurrent_bytes={dropped_bytes} " + f"total_evicted_recurrent_blocks=" + f"{self._recurrent_evicted_blocks_total} " + f"total_onboarded_recurrent_blocks=" + f"{self._recurrent_onboarded_blocks_total} " + f"total_dropped_recurrent_blocks=" + f"{self._recurrent_dropped_blocks_total} " + f"gpu_used_recurrent_blocks=" + f"{sum(stat.primary_used_num_blocks for stat in stats)} " + f"gpu_free_recurrent_blocks=" + f"{sum(stat.primary_free_num_blocks for stat in stats)} " + f"gpu_evictable_recurrent_blocks=" + f"{sum(stat.primary_evictable_num_blocks for stat in stats)} " + f"host_used_recurrent_blocks=" + f"{sum(stat.secondary_used_num_blocks for stat in stats)} " + f"host_free_recurrent_blocks=" + f"{sum(stat.secondary_free_num_blocks for stat in stats)} " + f"{self._format_branch_snapshot_counters()}" + ) + self._recurrent_status_logged = True + return report + + @override + def _iter_cache_buffers_for_invalid_check(self) -> Iterable[torch.Tensor]: + for global_layer_id, local_layer_id in self.layer_offsets.items(): + if self._is_local_mamba_layer(local_layer_id): + continue + # A layer group is a lifecycle, not a physical memory pool. + # Differently sized attention buffers can share one lifecycle, + # so scan every attention layer in this diagnostic path. + yield KVCacheManagerV2.get_buffers(self, global_layer_id) + + for (layer_id, role), buffer in self._recurrent_buffers.items(): + if buffer is None: + raise RuntimeError( + f"Recurrent buffer view was not bound: layer={layer_id}, role={role}" + ) + yield buffer + + @override + def shutdown(self): + self._shutdown_model_state() + self._recurrent_buffers.clear() + self.all_ssm_states = [] + self.all_conv_states = [] + self._dummy_request_mask = None + self._dummy_request_mask_host = None + self._generation_state_indices = None + self._branch_snapshot_points.clear() + super().shutdown() + + def _shutdown_model_state(self) -> None: + """Release model-owned state before the common pool is destroyed.""" + + +__all__ = ["MambaHybridCacheManagerV2"] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index dcb09edd5814..3f79ee122374 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -7541,21 +7541,11 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): ResourceManagerType.SEQ_SLOT_MANAGER].prepare_resources( requests) self._setup_sampler_step(requests) - # KDA fused-verify replay caches: the ctx->gen state transfer - # populates only the base mamba conv/ssm pools; with one-model - # spec decoding (e.g. SA) the first forward for these requests - # takes the fused verify path, which reads the per-slot - # kda_conv_* replay caches instead of the conv pool. Seed them - # from the transferred conv states before the first step - # (otherwise the recurrent state is permanently contaminated by - # a zero/stale conv window; K3 SA-in-disagg GSM8K -0.6pp). if self.model_engine.enable_spec_decode: kv_mgr = self.resource_manager.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER) - seed = getattr(kv_mgr, 'seed_kda_replay_caches_for_disagg_gen', - None) - if seed is not None: - seed([ + if isinstance(kv_mgr, BaseMambaCacheManager): + kv_mgr.on_state_transfer_complete([ req.py_request_id for req in cache_trans_complete_requests ]) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 49fef0c8d50e..087dc13b39fa 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -12,6 +12,16 @@ from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager from tensorrt_llm._torch.disaggregation.resource.page import AttentionLayerGroup, MambaLayerGroup from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +from tensorrt_llm._torch.modules.fla.cache_manager import GDNReplayState, Qwen35HybridCacheManagerV2 +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAReplayState, + KimiK3HybridCacheManagerV2, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + IntermediateState, + NemotronHybridCacheManagerV2, + ReplayHistory, +) from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( KvCacheCreator, @@ -268,7 +278,7 @@ def _capture_kimi_v2_manager_ctor( V2 manager and capture the constructor arguments.""" captured: dict[str, object] = {} - class RecordingV2Manager(MambaHybridCacheManagerV2): + class RecordingV2Manager(KimiK3HybridCacheManagerV2): def __init__(self, *args: object, **kwargs: object) -> None: captured["args"] = args captured["kwargs"] = kwargs @@ -278,7 +288,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - assert get_kv_cache_manager_cls(model_config, kv_cache_config) is MambaHybridCacheManagerV2 + assert get_kv_cache_manager_cls(model_config, kv_cache_config) is KimiK3HybridCacheManagerV2 _create_kv_cache_manager( model_engine=None, @@ -323,7 +333,7 @@ def test_kimi_explicit_v2_manager_geometry(monkeypatch: pytest.MonkeyPatch) -> N assert "kda_replay_num_spec" not in kwargs -def test_kimi_explicit_v2_manager_enables_kda_replay( +def test_kimi_explicit_v2_manager_enables_owned_kda_replay_policy( monkeypatch: pytest.MonkeyPatch, ) -> None: spec_config = MTPDecodingConfig(max_draft_len=3) @@ -334,8 +344,17 @@ def test_kimi_explicit_v2_manager_enables_kda_replay( _, kwargs = _capture_kimi_v2_manager_ctor(monkeypatch, spec_config) - assert kwargs["kda_replay_num_spec"] == spec_config.tokens_per_gen_step - 1 + assert "speculative_state" not in kwargs + mgr = object.__new__(KimiK3HybridCacheManagerV2) + mgr.spec_config = spec_config + mgr._requested_num_spec = None + mgr._state_layout = SimpleNamespace(spec_config=spec_config, mamba_pp_layers=()) + mgr._initialize_model_state() + policy = mgr._kda_replay + assert isinstance(policy, KDAReplayState) + assert policy.num_speculative_tokens == spec_config.tokens_per_gen_step - 1 assert kwargs["conv_state_layout"] == "q_k_v" + assert "kda_replay_num_spec" not in kwargs def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout( @@ -422,8 +441,8 @@ def test_v2_manager_rejects_model_type_kwarg() -> None: @pytest.mark.parametrize( ("use_v2", "enable_block_reuse", "expected"), [ - (True, False, MambaHybridCacheManagerV2), - (True, True, MambaHybridCacheManagerV2), + (True, False, NemotronHybridCacheManagerV2), + (True, True, NemotronHybridCacheManagerV2), (False, True, CppMambaHybridCacheManager), ("auto", True, CppMambaHybridCacheManager), ], @@ -486,7 +505,7 @@ def __init__(self, *args, **kwargs): mamba_ssm_cache_dtype=torch.bfloat16, ) monkeypatch.setenv("TRTLLM_USE_GDN_REPLAY", "1") - monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 90) + monkeypatch.setattr("tensorrt_llm._torch.modules.fla.cache_manager.get_sm_version", lambda: 90) monkeypatch.setattr( "tensorrt_llm._torch.pyexecutor._util.extract_mamba_kv_cache_params", lambda *args, **kwargs: mamba_params, @@ -535,7 +554,18 @@ def __init__(self, *args, **kwargs): assert captured_mixed["use_replay_state_update"] is False assert captured_mixed["model_type"] == "qwen3_next" assert captured_mixed["max_num_tokens"] == 256 - assert captured_v2["use_replay_state_update"] is True + assert captured_v2["use_replay_state_update"] is None + assert "speculative_state" not in captured_v2 + mgr = object.__new__(Qwen35HybridCacheManagerV2) + mgr.spec_config = common_kwargs["spec_config"] + mgr.ssm_state_dtype = torch.bfloat16 + mgr._requested_replay = None + mgr.local_num_mamba_layers = 0 + mgr._state_layout = SimpleNamespace( + spec_config=mgr.spec_config, mamba_pp_layers=(), n_groups_per_rank=1 + ) + mgr._initialize_model_state() + assert isinstance(mgr._speculative_state, GDNReplayState) assert captured_v2["max_num_tokens"] == 256 assert "model_type" not in captured_v2 assert captured_v2["conv_state_layout"] == "q_k_v" @@ -636,7 +666,7 @@ def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, bac backend=backend, transceiver_runtime="PYTHON" ), ) - is MambaHybridCacheManagerV2 + is NemotronHybridCacheManagerV2 ) @@ -824,7 +854,7 @@ def test_qwen3_gdn_replay_uses_v2_preference( ) assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected_v2 - expected_manager = MambaHybridCacheManagerV2 if expected_v2 else CppMambaHybridCacheManager + expected_manager = NemotronHybridCacheManagerV2 if expected_v2 else CppMambaHybridCacheManager assert ( get_kv_cache_manager_cls( _hybrid_model_config(), @@ -1261,16 +1291,13 @@ def test_cpp_hybrid_replay_bookkeeping_is_fused_into_conv_promotion( promote_calls = [] monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.legacy._promote_mamba_state_triton", lambda *args, **kwargs: promote_calls.append((args, kwargs)), ) - monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._advance_replay_state", - lambda *args, **kwargs: pytest.fail( - "Cpp replay bookkeeping must not advance before conv promotion" - ), - ) + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask = torch.zeros(64, dtype=torch.bool) mgr.update_mamba_states( SimpleNamespace(num_seqs=2, num_contexts=0), torch.tensor([3, 2], dtype=torch.int32), @@ -1646,7 +1673,7 @@ def test_v2_hybrid_estimator_counts_dummy_states_without_attention_capacity( lambda *args, **kwargs: 11, ) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._get_local_mamba_cache_layout", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common._get_local_mamba_cache_layout", lambda *args, **kwargs: ( SimpleNamespace(get_states_bytes_per_layer=lambda mapping: 64), 1, @@ -1702,6 +1729,13 @@ def _base_attention_layer_configs(num_layers): ] +def _install_default_v2_state(mgr): + """Populate local layer IDs for tests of V2 config building alone.""" + mgr.mamba_pp_layers = [ + layer_id for layer_id in mgr.pp_layers if mgr._mamba_layer_mask[layer_id] + ] + + def test_v2_hybrid_typical_batch_splits_capacity_across_ssm_states_and_dummies(): mgr = object.__new__(MambaHybridCacheManagerV2) mgr.kv_cache_type = CacheTypeCpp.SELF @@ -1718,6 +1752,7 @@ def test_v2_hybrid_typical_batch_splits_capacity_across_ssm_states_and_dummies() mgr.tokens_per_block = 32 mgr.num_local_layers = 2 mgr.local_num_mamba_layers = 1 + _install_default_v2_state(mgr) mgr._num_reserved_dummy_slots = 1 mgr.dtype = DataType.HALF mgr.enable_swa_scratch_reuse = False @@ -1782,7 +1817,7 @@ def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): mgr.max_seq_len = 4096 warnings_seen = [] monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.logger.warning", + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2.logger.warning", lambda message: warnings_seen.append(message), ) @@ -1801,12 +1836,14 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) mgr.local_num_mamba_layers = 1 + mgr._mamba_layer_mask = [True, False] mgr.ssm_bytes = 64 mgr.conv_bytes = 32 mgr._num_reserved_dummy_slots = 1 mgr.tokens_per_block = 32 mgr.num_local_layers = 2 mgr.pp_layers = [0, 1] + _install_default_v2_state(mgr) mgr.max_attention_window_vec = [128, 128] mgr.max_num_tokens = 128 mgr.enable_swa_scratch_reuse = False @@ -1831,12 +1868,14 @@ def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) mgr.local_num_mamba_layers = 1 + mgr._mamba_layer_mask = [True] mgr.ssm_bytes = 64 mgr.conv_bytes = 32 mgr._num_reserved_dummy_slots = 1 mgr.tokens_per_block = 32 mgr.num_local_layers = 1 mgr.pp_layers = [0] + _install_default_v2_state(mgr) mgr.max_attention_window_vec = [128] mgr.max_num_tokens = 128 mgr.enable_swa_scratch_reuse = False @@ -2069,6 +2108,7 @@ def allocated_memory(pool_ratio): mgr.tokens_per_block = 32 mgr.num_local_layers = 2 mgr.local_num_mamba_layers = 1 + _install_default_v2_state(mgr) mgr._num_reserved_dummy_slots = 1 mgr.dtype = DataType.HALF mgr.enable_swa_scratch_reuse = False @@ -2196,10 +2236,10 @@ def _build_v2_hybrid_with_mamba_layer( enable_swa_scratch_reuse=False, dtype=DataType.HALF, conv_state_layout="x_b_c", + speculative_state=None, mamba_d_conv=4, mamba_n_groups=1, mamba_ssm_cache_dtype=torch.float16, - kda_replay_num_spec=None, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2228,7 +2268,29 @@ def _build_v2_hybrid_with_mamba_layer( ), dtype="nvfp4" if dtype == DataType.NVFP4 else "auto", ) - return MambaHybridCacheManagerV2( + model_cls = ( + KimiK3HybridCacheManagerV2 + if isinstance(speculative_state, KDAReplayState) + else Qwen35HybridCacheManagerV2 + if conv_state_layout == "q_k_v" + else NemotronHybridCacheManagerV2 + ) + + class ModelStateManager(model_cls): + def _initialize_model_state(self): + if speculative_state is None: + super()._initialize_model_state() + elif isinstance(speculative_state, KDAReplayState): + self._kda_replay = speculative_state + self._kda_replay.validate(self._state_layout) + else: + from tensorrt_llm._torch.modules.fla.cache_manager import validate_gdn_layout + + self._speculative_state = speculative_state + validate_gdn_layout(self, self._speculative_state) + + manager_cls = ModelStateManager + return manager_cls( mamba_d_state=8, mamba_d_conv=mamba_d_conv, mamba_num_heads=4, @@ -2253,7 +2315,6 @@ def _build_v2_hybrid_with_mamba_layer( use_replay_state_update=use_replay_state_update, dtype=dtype, conv_state_layout=conv_state_layout, - kda_replay_num_spec=kda_replay_num_spec, ) @@ -2575,7 +2636,8 @@ def test_v2_hybrid_debug_logs_prefix_reuse_only_on_rank_zero( log_debug = MagicMock() monkeypatch.setattr(KVCacheManagerV2, "_create_kv_cache", create_kv_cache) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.logger.debug", log_debug + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2.logger.debug", + log_debug, ) mgr = object.__new__(MambaHybridCacheManagerV2) @@ -2643,7 +2705,8 @@ def test_v2_hybrid_logs_aggregated_recurrent_cache_status_only_on_rank_zero( monkeypatch.setattr(KVCacheManagerV2, "get_iteration_stats", MagicMock(return_value=report)) log_debug = MagicMock() monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.logger.debug", log_debug + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2.logger.debug", + log_debug, ) mgr = object.__new__(MambaHybridCacheManagerV2) @@ -3424,8 +3487,11 @@ def test_v2_hybrid_intermediate_states_size_by_tokens_per_gen_step(): @skip_no_cuda -def test_v2_kda_replay_allocates_logical_slot_caches(): +def test_v2_kda_replay_policy_allocates_logical_slot_caches(): spec_config = MTPDecodingConfig(max_draft_len=2) + policy = KDAReplayState( + num_speculative_tokens=2, + ) mgr = _build_v2_hybrid_with_mamba_layer( max_batch_size=4, num_mamba_layers=2, @@ -3434,15 +3500,16 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): mamba_d_conv=5, mamba_n_groups=4, mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=2, + speculative_state=policy, ) try: - assert mgr.intermediate_ssm_states is None - assert mgr.intermediate_conv_states is None - assert mgr.intermediate_state_indices.shape == (4,) - assert not mgr.use_replay_state_update + assert mgr._kda_replay is policy + assert policy.intermediate_indices.shape == (4,) + assert mgr.get_replay_state_update_metadata() is None layer_cache = mgr.mamba_layer_cache(0) + assert layer_cache.intermediate_ssm is None + assert layer_cache.intermediate_conv_window is None cache_size = layer_cache.temporal.shape[0] assert layer_cache.kda_conv_q.shape == (cache_size, 32, 6) assert layer_cache.kda_conv_k.shape == (cache_size, 32, 6) @@ -3454,7 +3521,7 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): assert layer_cache.kda_beta_cache.shape == (cache_size, 2, 4) assert layer_cache.kda_qkg_cache.dtype is torch.float32 assert layer_cache.prev_num_accepted_tokens.data_ptr() == ( - mgr.prev_num_accepted_tokens.data_ptr() + policy.prev_num_accepted_tokens.data_ptr() ) assert layer_cache.intermediate_ssm is None assert layer_cache.intermediate_conv_window is None @@ -3463,92 +3530,125 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): @pytest.mark.parametrize( - ("spec_config", "use_replay_state_update", "conv_state_layout", "num_spec", "match"), + ("spec_config", "conv_state_layout", "num_spec", "match"), [ - (None, False, "q_k_v", 2, "require speculative decoding"), - ( - MTPDecodingConfig(max_draft_len=2), - True, - "q_k_v", - 2, - "mutually exclusive", - ), + (None, "q_k_v", 2, "require speculative decoding"), ( MTPDecodingConfig(max_draft_len=2), - False, "q_k_v", 1, "must match the draft length", ), ( MTPDecodingConfig(max_draft_len=2), - False, "x_b_c", 2, "requires conv_state_layout='q_k_v'", ), + ( + MTPDecodingConfig(max_draft_len=2), + "q_k_v", + 2, + "require equal", + ), ], ) -def test_v2_kda_replay_validates_configuration( +def test_v2_kda_replay_policy_validates_configuration( spec_config, - use_replay_state_update, conv_state_layout, num_spec, match, ): + policy = KDAReplayState( + num_speculative_tokens=num_spec, + ) with pytest.raises(ValueError, match=match): _build_v2_hybrid_with_mamba_layer( spec_config=spec_config, - use_replay_state_update=use_replay_state_update, conv_state_layout=conv_state_layout, - kda_replay_num_spec=num_spec, + speculative_state=policy, ) -def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: - mgr = object.__new__(MambaCacheManager) - mgr._impl = SimpleNamespace(use_kda_replay_update=True) - - assert mgr.use_kda_replay_update +@skip_no_cuda +def test_v2_kda_replay_policy_supports_attention_only_pp_rank(): + policy = KDAReplayState( + num_speculative_tokens=2, + ) + mapping = Mapping(world_size=2, rank=1, tp_size=1, pp_size=2) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=1, + mapping=mapping, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_n_groups=4, + speculative_state=policy, + ) + try: + assert mgr.local_num_mamba_layers == 0 + assert mgr._kda_replay is policy + assert policy.prev_num_accepted_tokens is None + finally: + mgr.shutdown() -def test_v2_kda_replay_records_acceptance_and_skips_dummy_rows(): - mgr = object.__new__(MambaHybridCacheManagerV2) +def test_v2_kda_replay_policy_records_acceptance_and_skips_dummy_rows(): + mgr = object.__new__(KimiK3HybridCacheManagerV2) mgr.local_num_mamba_layers = 1 - mgr._use_kda_replay_update = True - mgr._use_replay_state_update = False - mgr.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) + policy = KDAReplayState( + num_speculative_tokens=2, + ) + policy.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) + policy.intermediate_indices = torch.arange(3, dtype=torch.int32) mgr._dummy_request_mask = torch.tensor([False, False, True]) + mgr._kda_replay = policy + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask = torch.zeros(64, dtype=torch.bool) mgr.update_mamba_states( SimpleNamespace(num_seqs=3, num_contexts=1), torch.tensor([1, 3, 2], dtype=torch.int32), state_indices=torch.tensor([0, 1, 2], dtype=torch.int32), ) - assert mgr.prev_num_accepted_tokens.tolist() == [7, 2, 9] + assert policy.prev_num_accepted_tokens.tolist() == [7, 2, 9] -def test_v2_kda_replay_resets_context_slots(): - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr._use_replay_state_update = False +def test_legacy_mamba_cache_manager_delegates_kda_replay_capability() -> None: + mgr = object.__new__(MambaCacheManager) + mgr._impl = SimpleNamespace(use_kda_replay_update=True) + + assert mgr.use_kda_replay_update + + +def test_v2_kda_replay_policy_resets_context_slots(): + mgr = object.__new__(KimiK3HybridCacheManagerV2) + policy = KDAReplayState( + num_speculative_tokens=2, + ) + policy.prev_num_accepted_tokens = torch.tensor([5, 6, 7], dtype=torch.int32) + mgr._kda_replay = policy mgr.cuda_state_indices = torch.tensor([2, 1], dtype=torch.int32) - mgr.prev_num_accepted_tokens = torch.tensor([5, 6, 7], dtype=torch.int32) - mgr.mamba_ssm_rand_seed = None + mgr._host_state_indices = torch.tensor([2, 1], dtype=torch.int32) + mgr._host_state_indices = mgr.cuda_state_indices.clone() mgr._reset_context_mamba_slots(num_contexts=1) - assert mgr.prev_num_accepted_tokens.tolist() == [5, 6, 0] + assert policy.prev_num_accepted_tokens.tolist() == [5, 6, 0] def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): - mgr = object.__new__(MambaHybridCacheManagerV2) + mgr = object.__new__(KimiK3HybridCacheManagerV2) mgr.local_num_mamba_layers = 1 - mgr._record_kda_replay_in_update_resources = True + mgr.spec_config = SimpleNamespace(decoding_type="NGram") mgr._request_id_to_state_index = {101: 1, 303: 2} mgr._request_id_to_is_dummy = {101: False, 303: True} - mgr.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) + policy = KDAReplayState( + num_speculative_tokens=2, + ) + policy.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) + mgr._kda_replay = policy base_update = MagicMock() monkeypatch.setattr(KVCacheManagerV2, "update_resources", base_update) active = SimpleNamespace( @@ -3571,42 +3671,36 @@ def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): mgr.update_resources(scheduled_batch) base_update.assert_called_once_with(scheduled_batch, None, None) - assert mgr.prev_num_accepted_tokens.tolist() == [7, 2, 9] + assert policy.prev_num_accepted_tokens.tolist() == [7, 2, 9] @skip_no_cuda -def test_v2_kda_replay_relocates_live_slot_history(): +def test_v2_kda_replay_policy_relocates_live_slot_history(): + policy = KDAReplayState( + num_speculative_tokens=2, + ) mgr = _build_v2_hybrid_with_mamba_layer( spec_config=MTPDecodingConfig(max_draft_len=2), conv_state_layout="q_k_v", mamba_n_groups=4, mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=2, + speculative_state=policy, ) try: - assert mgr.prev_num_accepted_tokens is not None - mgr.prev_num_accepted_tokens[:3].copy_( + policy.prev_num_accepted_tokens[:3].copy_( torch.tensor([10, 20, 30], dtype=torch.int32, device="cuda") ) - replay_buffers = ( - mgr.kda_conv_q, - mgr.kda_conv_k, - mgr.kda_conv_v, - mgr.kda_qkg_cache, - mgr.kda_v_cache, - mgr.kda_beta_cache, - ) + replay_buffers = policy._layer_replay_buffers() source_rows = [] for offset, replay_buffer in enumerate(replay_buffers): - assert replay_buffer is not None replay_buffer.zero_() replay_buffer[:, 0].fill_(offset * 10) replay_buffer[:, 1].fill_(offset * 10 + 1) source_rows.append((replay_buffer[:, 0].clone(), replay_buffer[:, 1].clone())) - mgr._relocate_kda_replay_slots([0, 1], [1, 2]) + policy.relocate_slots([0, 1], [1, 2]) - assert mgr.prev_num_accepted_tokens[:3].tolist() == [10, 10, 20] + assert policy.prev_num_accepted_tokens[:3].tolist() == [10, 10, 20] for replay_buffer, (source_zero, source_one) in zip(replay_buffers, source_rows): torch.testing.assert_close(replay_buffer[:, 1], source_zero) torch.testing.assert_close(replay_buffer[:, 2], source_one) @@ -3615,23 +3709,27 @@ def test_v2_kda_replay_relocates_live_slot_history(): @skip_no_cuda -def test_v2_kda_state_index_setup_relocates_generation_history(): +def test_v2_state_index_setup_relocates_kda_generation_history(): class StateCache: def __init__(self, slot: int) -> None: self.slot = slot - def get_ssm_block_base_index(self, _layer_group_id) -> int: + def get_ssm_block_base_index(self, layer_group_id) -> int: + del layer_group_id return self.slot def close(self) -> None: pass + policy = KDAReplayState( + num_speculative_tokens=2, + ) mgr = _build_v2_hybrid_with_mamba_layer( spec_config=MTPDecodingConfig(max_draft_len=2), conv_state_layout="q_k_v", mamba_n_groups=4, mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=2, + speculative_state=policy, ) context = SimpleNamespace(py_request_id=101, is_dummy=False) relocated = SimpleNamespace(py_request_id=202, is_dummy=False) @@ -3639,79 +3737,77 @@ def close(self) -> None: mgr.kv_cache_map.update({101: StateCache(0), 202: StateCache(1)}) try: mgr._setup_state_indices([context, relocated], num_contexts=1) - assert mgr.prev_num_accepted_tokens is not None - mgr.prev_num_accepted_tokens[1] = 7 - mgr.prev_num_accepted_tokens[3] = 9 - replay_buffers = ( - mgr.kda_conv_q, - mgr.kda_conv_k, - mgr.kda_conv_v, - mgr.kda_qkg_cache, - mgr.kda_v_cache, - mgr.kda_beta_cache, - ) + policy.prev_num_accepted_tokens[1] = 7 + policy.prev_num_accepted_tokens[3] = 9 + replay_buffers = policy._layer_replay_buffers() relocated_rows = [] for offset, replay_buffer in enumerate(replay_buffers): - assert replay_buffer is not None replay_buffer[:, 1].fill_(offset + 1) relocated_rows.append(replay_buffer[:, 1].clone()) mgr.kv_cache_map[202].slot = 2 mgr.kv_cache_map[303] = StateCache(3) - mgr._setup_state_indices([context, relocated, first_generation], num_contexts=1) + mgr._setup_state_indices( + [context, relocated, first_generation], + num_contexts=1, + ) assert mgr._request_id_to_state_index == {101: 0, 202: 2, 303: 3} - assert mgr.prev_num_accepted_tokens[2] == 7 - assert mgr.prev_num_accepted_tokens[3] == 0 + assert policy.prev_num_accepted_tokens[2] == 7 + assert policy.prev_num_accepted_tokens[3] == 0 for replay_buffer, relocated_row in zip(replay_buffers, relocated_rows): torch.testing.assert_close(replay_buffer[:, 2], relocated_row) finally: mgr.shutdown() -def test_v2_kda_replay_seeds_disaggregated_generation_slots(): - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr.local_num_mamba_layers = 2 - mgr.conv_state_shape = [6, 4] - mgr.conv_section_dims = [2, 2, 2] - mgr._request_id_to_state_index = {101: 1, 202: 3} - mgr.prev_num_accepted_tokens = torch.full((4,), 5, dtype=torch.int32) - mgr.all_conv_states = [ - torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4), - torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4) + 1000, - ] +def test_v2_kda_replay_policy_seeds_disaggregated_generation_slots(): + policy = KDAReplayState( + num_speculative_tokens=2, + ) + policy._conv_section_dims = (2, 2, 2) + policy.prev_num_accepted_tokens = torch.full((4,), 5, dtype=torch.int32) def dim_contiguous_conv_cache() -> torch.Tensor: return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) - mgr.kda_conv_q = dim_contiguous_conv_cache() - mgr.kda_conv_k = dim_contiguous_conv_cache() - mgr.kda_conv_v = dim_contiguous_conv_cache() - mgr.kda_qkg_cache = torch.full((2, 4, 2, 3, 2), 7.0) - mgr.kda_v_cache = torch.full((2, 4, 2, 2), 7.0) - mgr.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0) + policy.kda_conv_q = dim_contiguous_conv_cache() + policy.kda_conv_k = dim_contiguous_conv_cache() + policy.kda_conv_v = dim_contiguous_conv_cache() + policy.kda_qkg_cache = torch.full((2, 4, 2, 3, 2), 7.0) + policy.kda_v_cache = torch.full((2, 4, 2, 2), 7.0) + policy.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0) + manager = SimpleNamespace( + conv_state_shape=[6, 4], + _request_id_to_state_index={101: 1, 202: 3}, + all_conv_states=[ + torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4), + torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4) + 1000, + ], + ) - mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) + policy._conv_states = tuple(manager.all_conv_states) + policy._committed_window = manager.conv_state_shape[1] + policy.seed_transferred_slots(torch.tensor([1, 3])) - for layer_offset, conv_state in enumerate(mgr.all_conv_states): + for layer_offset, conv_state in enumerate(manager.all_conv_states): for slot in (1, 3): for replay_buffer, start in ( - (mgr.kda_conv_q, 0), - (mgr.kda_conv_k, 2), - (mgr.kda_conv_v, 4), + (policy.kda_conv_q, 0), + (policy.kda_conv_k, 2), + (policy.kda_conv_v, 4), ): torch.testing.assert_close( replay_buffer[layer_offset, slot, :, :4], conv_state[slot, start : start + 2, :], ) assert torch.count_nonzero(replay_buffer[layer_offset, slot, :, 4:]) == 0 - assert torch.count_nonzero(mgr.kda_qkg_cache[layer_offset, slot]) == 0 - assert torch.count_nonzero(mgr.kda_v_cache[layer_offset, slot]) == 0 - assert torch.count_nonzero(mgr.kda_beta_cache[layer_offset, slot]) == 0 - assert mgr.prev_num_accepted_tokens[slot] == 0 - assert mgr.prev_num_accepted_tokens[0] == 5 - assert mgr.prev_num_accepted_tokens[2] == 5 + assert torch.count_nonzero(policy.kda_qkg_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(policy.kda_v_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(policy.kda_beta_cache[layer_offset, slot]) == 0 + assert policy.prev_num_accepted_tokens[slot] == 0 + assert policy.prev_num_accepted_tokens[0] == 5 + assert policy.prev_num_accepted_tokens[2] == 5 @skip_no_cuda @@ -3731,7 +3827,7 @@ def test_v2_hybrid_static_dynamic_tree_capacity(): assert mgr.intermediate_conv_states.shape[2] == 32 assert mgr._kv_reserve_draft_tokens == 31 assert mgr._num_reserved_dummy_slots == 1 - assert not mgr.use_replay_state_update + assert mgr.get_replay_state_update_metadata() is None finally: mgr.shutdown() @@ -3751,7 +3847,6 @@ def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): ) 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( @@ -3767,17 +3862,18 @@ def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): @skip_no_cuda def test_v2_gdn_replay_builds_affine_state_layout(): + policy = GDNReplayState(tokens_per_gen_step=4) mgr = _build_v2_hybrid_with_mamba_layer( max_batch_size=16, num_mamba_layers=3, spec_config=MTPDecodingConfig(max_draft_len=3), - use_replay_state_update=True, conv_state_layout="q_k_v", + speculative_state=policy, ) try: assert mgr.use_gdn_cached_replay_all_layer_commit - assert mgr._gdn_cached_replay_state_descriptors is None - state_strides = mgr._gdn_cached_replay_state_strides + assert policy._state_descriptors is None + state_strides = policy._state_strides assert state_strides is not None states = mgr.all_ssm_states first_state = states[0] @@ -3800,27 +3896,28 @@ def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): ) batch_size = 16 + policy = GDNReplayState(tokens_per_gen_step=4) mgr = _build_v2_hybrid_with_mamba_layer( max_batch_size=batch_size, num_mamba_layers=3, spec_config=MTPDecodingConfig(max_draft_len=3), - use_replay_state_update=True, conv_state_layout="q_k_v", + speculative_state=policy, ) try: torch.manual_seed(1234) for state in mgr.all_ssm_states: state.normal_() - mgr.old_x.normal_() - mgr.old_B.normal_() - mgr.old_dt.uniform_(-0.2, -0.01) + policy.old_x.normal_() + policy.old_B.normal_() + policy.old_dt.uniform_(-0.2, -0.01) positions = torch.arange(batch_size, device="cuda", dtype=torch.int32) replay_work_items = torch.stack( ( positions, positions, - torch.full_like(positions, mgr.replay_history_size), + torch.full_like(positions, policy.replay_history_size), torch.zeros_like(positions), ), dim=1, @@ -3829,15 +3926,15 @@ def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): expected = torch.stack(mgr.all_ssm_states) commit_gdn_cached_replay_history_layers( ssm_states=expected, - old_u=mgr.old_x, - old_k=mgr.old_B, - old_G=mgr.old_dt, + old_u=policy.old_x, + old_k=policy.old_B, + old_G=policy.old_dt, replay_work_items=replay_work_items, n_writes=n_writes, - history_size=mgr.replay_history_size, + history_size=policy.replay_history_size, ) - mgr._commit_gdn_cached_replay_history_layers( + policy.commit_all_layers( SimpleNamespace( mamba_metadata=SimpleNamespace( replay_num_decodes=batch_size, @@ -3854,22 +3951,30 @@ def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): - mgr = object.__new__(MambaHybridCacheManagerV2) + mgr = object.__new__(Qwen35HybridCacheManagerV2) batch_size = 16 mgr.local_num_mamba_layers = 1 - mgr._use_replay_state_update = True - mgr._use_gdn_cached_replay_all_layer_commit = True - mgr.replay_step_width = 4 - mgr.replay_history_size = MIN_REPLAY_HISTORY_SIZE - mgr.prev_num_accepted_tokens = torch.zeros(batch_size, dtype=torch.int32) - mgr.cache_buf_idx = torch.zeros(batch_size, dtype=torch.int32) - mgr.intermediate_state_indices = torch.arange(batch_size, dtype=torch.int32) + policy = GDNReplayState(tokens_per_gen_step=4) + policy.prev_num_accepted_tokens = torch.zeros(batch_size, dtype=torch.int32) + policy.cache_buf_idx = torch.zeros(batch_size, dtype=torch.int32) + policy.intermediate_indices = torch.arange(batch_size, dtype=torch.int32) mgr._dummy_request_mask = torch.zeros(batch_size, dtype=torch.bool) + policy.intermediate_conv = torch.empty((1, batch_size, 1)) + policy.old_x = torch.empty(0) + policy.old_B = torch.empty(0) + policy.old_dt = torch.empty(0) + policy._state_strides = (1, 0, 0) + mgr._speculative_state = policy + mgr.all_ssm_states = [torch.empty(0)] mgr.all_conv_states = [torch.empty(0)] - mgr.intermediate_conv_states = torch.empty(0) + policy._ssm_states = tuple(mgr.all_ssm_states) + policy._conv_states = tuple(mgr.all_conv_states) events = [] - mgr._commit_gdn_cached_replay_history_layers = lambda *_args, **_kwargs: events.append("commit") + monkeypatch.setattr( + "tensorrt_llm._torch.modules.fla.cache_manager.commit_gdn_cached_replay_history_layers", + lambda **_kwargs: events.append("commit"), + ) monkeypatch.setattr( "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._advance_replay_state", lambda *_args, **_kwargs: events.append("advance"), @@ -3879,8 +3984,19 @@ def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): lambda *_args, **_kwargs: None, ) + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask = torch.zeros(64, dtype=torch.bool) mgr.update_mamba_states( - SimpleNamespace(num_seqs=batch_size, num_contexts=0), + SimpleNamespace( + num_seqs=batch_size, + num_contexts=0, + mamba_metadata=SimpleNamespace( + replay_num_decodes=batch_size, + replay_work_items=torch.empty((batch_size, 4), dtype=torch.int32), + replay_n_writes=torch.empty(1, dtype=torch.int32), + ), + ), torch.full((batch_size,), 3, dtype=torch.int32), state_indices=torch.arange(batch_size, dtype=torch.int32), ) @@ -3889,7 +4005,7 @@ def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): - mgr = object.__new__(MambaHybridCacheManagerV2) + mgr = object.__new__(NemotronHybridCacheManagerV2) mgr.local_num_mamba_layers = 1 mgr._request_id_to_state_index = { 100: 0, @@ -3903,17 +4019,18 @@ def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): 102: True, 103: False, } + policy = ReplayHistory(tokens_per_gen_step=5) mgr._dummy_request_mask = torch.zeros(4, dtype=torch.bool) mgr._dummy_request_mask_host = torch.zeros(4, dtype=torch.bool) - mgr._use_replay_state_update = True - mgr.replay_step_width = 5 - mgr.replay_history_size = 16 - mgr.prev_num_accepted_tokens = torch.full((4,), 13, dtype=torch.int32) - mgr.cache_buf_idx = torch.ones(4, dtype=torch.int32) - mgr.intermediate_state_indices = torch.arange(4, dtype=torch.int32) + policy.prev_num_accepted_tokens = torch.full((4,), 13, dtype=torch.int32) + policy.cache_buf_idx = torch.ones(4, dtype=torch.int32) + policy.intermediate_indices = torch.arange(4, dtype=torch.int32) + policy.intermediate_conv = torch.empty((1, 4, 1)) + mgr._speculative_state = policy mgr.all_ssm_states = [] mgr.all_conv_states = [torch.empty(0)] - mgr.intermediate_conv_states = torch.empty(0) + policy._ssm_states = tuple(mgr.all_ssm_states) + policy._conv_states = tuple(mgr.all_conv_states) monkeypatch.setattr( "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", lambda *args, **kwargs: None, @@ -3926,25 +4043,31 @@ def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): ) assert mgr._dummy_request_mask.tolist() == [False, False, True, True] + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask = torch.zeros(64, dtype=torch.bool) mgr.update_mamba_states( SimpleNamespace(num_seqs=4, num_contexts=1), torch.tensor([1, 3, 3, 3], dtype=torch.int32), state_indices=state_indices, ) - assert mgr.prev_num_accepted_tokens.tolist() == [13, 3, 13, 13] - assert mgr.cache_buf_idx.tolist() == [1, 0, 1, 1] + assert policy.prev_num_accepted_tokens.tolist() == [13, 3, 13, 13] + assert policy.cache_buf_idx.tolist() == [1, 0, 1, 1] def test_v2_hybrid_dynamic_tree_promotes_accepted_leaf_state(monkeypatch): - mgr = object.__new__(MambaHybridCacheManagerV2) + mgr = object.__new__(NemotronHybridCacheManagerV2) mgr.local_num_mamba_layers = 1 - mgr._use_replay_state_update = False - mgr.intermediate_state_indices = torch.arange(2, dtype=torch.int32) + policy = IntermediateState() + policy.intermediate_indices = torch.arange(2, dtype=torch.int32) + policy.intermediate_ssm = torch.empty((1, 2, 8)) + policy.intermediate_conv = torch.empty((1, 2, 8)) + mgr._speculative_state = policy mgr.all_ssm_states = [torch.empty(0)] mgr.all_conv_states = [torch.empty(0)] - mgr.intermediate_ssm_states = torch.empty((1, 2, 8)) - mgr.intermediate_conv_states = torch.empty((1, 2, 8)) + policy._ssm_states = tuple(mgr.all_ssm_states) + policy._conv_states = tuple(mgr.all_conv_states) promoted_positions = [] @@ -3956,6 +4079,9 @@ def capture_promoted_position(_dst, _src, _src_indices, positions, _dst_indices) capture_promoted_position, ) + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask = torch.zeros(64, dtype=torch.bool) mgr.update_mamba_states( SimpleNamespace(num_seqs=3, num_contexts=1), torch.tensor([1, 2, 3], dtype=torch.int32), diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py new file mode 100644 index 000000000000..44a7a751f288 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -0,0 +1,1061 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU contracts for model-owned hybrid managers and shared recurrent state.""" + +from dataclasses import replace +from types import FunctionType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.modules.fla.cache_manager import GDNReplayState, Qwen35HybridCacheManagerV2 +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAReplayState, + KimiK3HybridCacheManagerV2, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + Mamba2State, + NemotronHybridCacheManagerV2, + ReplayHistory, +) +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + PLE_CONV_STATE, + PLE_NGRAM_CONTEXT, + Qwen4ExpHybridCacheManagerV2, + Qwen4ExpPLECacheParams, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + IntermediateState, + MambaAcceptanceBatch, + MambaHybridCacheManager, + MambaRole, + MambaStateLayout, +) +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + GpuCacheTierConfig, + KVCacheManagerConfig, + LayerId, +) + + +@pytest.mark.parametrize( + "manager_cls", + [ + NemotronHybridCacheManagerV2, + Qwen35HybridCacheManagerV2, + Qwen4ExpHybridCacheManagerV2, + KimiK3HybridCacheManagerV2, + ], +) +def test_model_managers_inherit_the_common_lifecycle(manager_cls): + assert manager_cls.__bases__ == (MambaHybridCacheManagerV2,) + for method in ( + "try_commit_blocks", + "prepare_resources", + "_setup_state_indices", + "free_resources", + ): + assert getattr(manager_cls, method) is getattr(MambaHybridCacheManagerV2, method) + + +@pytest.mark.parametrize( + "manager_cls", + [ + MambaHybridCacheManagerV2, + NemotronHybridCacheManagerV2, + Qwen35HybridCacheManagerV2, + Qwen4ExpHybridCacheManagerV2, + KimiK3HybridCacheManagerV2, + ], +) +def test_v2_managers_do_not_expose_legacy_accessors(manager_cls): + for name in ( + "get_intermediate_ssm_states", + "get_intermediate_conv_states", + "get_mamba_ssm_cache_dtype", + "use_replay_state_update", + "_init_speculative_state", + "_validate_model_layout", + ): + assert not hasattr(manager_cls, name) + + +def test_legacy_managers_keep_compatibility_accessors(): + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + CppMambaHybridCacheManager, + MambaCacheManager, + MixedMambaHybridCacheManager, + PythonMambaCacheManager, + ) + + for manager_cls in ( + PythonMambaCacheManager, + MambaCacheManager, + MixedMambaHybridCacheManager, + CppMambaHybridCacheManager, + ): + for name in ( + "get_intermediate_ssm_states", + "get_intermediate_conv_states", + "get_mamba_ssm_cache_dtype", + "use_replay_state_update", + ): + assert hasattr(manager_cls, name) + + +@pytest.mark.parametrize( + "cls", + [ + MambaHybridCacheManager, + MambaHybridCacheManagerV2, + NemotronHybridCacheManagerV2, + Qwen35HybridCacheManagerV2, + Qwen4ExpHybridCacheManagerV2, + KimiK3HybridCacheManagerV2, + Mamba2State, + ReplayHistory, + GDNReplayState, + ], +) +def test_inherited_method_overrides_are_annotated(cls): + for name, attribute in vars(cls).items(): + if isinstance(attribute, (classmethod, staticmethod)): + function = attribute.__func__ + elif isinstance(attribute, property): + function = attribute.fget + else: + function = attribute + if not isinstance(function, FunctionType): + continue + # A standalone concrete class does not implement a project override + # contract merely because object defines __init__. + inherited = any(name in vars(parent) for parent in cls.__mro__[1:] if parent is not object) + assert getattr(function, "__override__", False) == inherited, f"{cls.__name__}.{name}" + + +def _layout(layers=(0, 1), slot_capacity=5): + return MambaStateLayout( + layer_mask=(True, True), + pp_layers=layers, + mamba_pp_layers=layers, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + conv_state_shape=(24, 3), + ssm_state_shape=(2, 4, 4), + conv_state_dtype=torch.float32, + ssm_state_dtype=torch.float32, + n_groups_per_rank=2, + max_batch_size=3, + state_index_capacity=4, + slot_capacity=slot_capacity, + spec_config=SimpleNamespace(tokens_per_gen_step=3), + conv_section_dims=(8, 8, 8), + conv_state_layout="q_k_v", + ) + + +def _state_views(layers=2, slots=5): + states = torch.zeros(layers, slots, 2, 4, 4) + return SimpleNamespace( + all_ssm_states=list(states.unbind()), + all_conv_states=[torch.zeros(slots, 24, 3) for _ in range(layers)], + ssm_state_shape=[2, 4, 4], + ) + + +def _prepare_model_layout(manager): + """Supply constructor-resolved inputs when testing the pre-pool hook alone.""" + manager._state_layout = replace(_layout(), spec_config=manager.spec_config) + manager.local_num_mamba_layers = 2 + manager._global_n_groups = 2 + if isinstance(manager, Qwen4ExpHybridCacheManagerV2): + manager._ple_params = None + manager._is_ple_draft = False + manager._pretrained_config = None + manager._mamba_layer_mask = [True, True] + manager.conv_state_dtype = torch.float32 + + +def test_qwen_managers_share_gdn_algorithm_not_model_inheritance(): + states = [] + for cls in (Qwen35HybridCacheManagerV2, Qwen4ExpHybridCacheManagerV2): + manager = object.__new__(cls) + manager._requested_replay = True + manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) + _prepare_model_layout(manager) + manager._initialize_model_state() + states.append(manager._speculative_state) + assert all(type(state) is GDNReplayState for state in states) + assert states[0] is not states[1] + + +@pytest.mark.parametrize( + "manager_cls", + [ + NemotronHybridCacheManagerV2, + Qwen35HybridCacheManagerV2, + Qwen4ExpHybridCacheManagerV2, + KimiK3HybridCacheManagerV2, + ], +) +def test_model_initialization_validates_selected_replay(manager_cls): + manager = object.__new__(manager_cls) + manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) + manager._requested_replay = True + manager._requested_num_spec = 2 + manager._kda_replay = None + _prepare_model_layout(manager) + manager._state_layout = replace( + manager._state_layout, n_groups_per_rank=0, conv_state_layout="x_b_c" + ) + with pytest.raises(ValueError, match="requires"): + manager._initialize_model_state() + + +@pytest.mark.parametrize( + "requested_num_spec, selected_num_spec, decoding_type", + [ + (2, None, "MTP"), + (None, 2, "MTP"), + (None, 2, "NGram"), + (None, None, "MTP"), + (None, None, None), + ], +) +def test_k3_initializes_only_the_selected_state( + monkeypatch, requested_num_spec, selected_num_spec, decoding_type +): + manager = object.__new__(KimiK3HybridCacheManagerV2) + manager._requested_num_spec = requested_num_spec + manager._kda_replay = None + manager.spec_config = ( + SimpleNamespace(decoding_type=decoding_type, tokens_per_gen_step=3) + if decoding_type is not None + else None + ) + _prepare_model_layout(manager) + intermediate_factory = Mock(wraps=IntermediateState) + replay_factory = Mock(wraps=KDAReplayState) + selector = Mock(return_value=selected_num_spec) + monkeypatch.setitem( + KimiK3HybridCacheManagerV2._initialize_model_state.__globals__, + "IntermediateState", + intermediate_factory, + ) + monkeypatch.setitem( + KimiK3HybridCacheManagerV2._initialize_model_state.__globals__, + "KDAReplayState", + replay_factory, + ) + monkeypatch.setitem( + KimiK3HybridCacheManagerV2._initialize_model_state.__globals__, + "get_kda_replay_num_spec", + selector, + ) + + assert manager._initialize_model_state() is None + num_spec = requested_num_spec if requested_num_spec is not None else selected_num_spec + if num_spec is not None: + intermediate_factory.assert_not_called() + replay_factory.assert_called_once_with(num_spec) + assert not hasattr(manager, "_speculative_state") + state = manager._kda_replay + else: + replay_factory.assert_not_called() + intermediate_factory.assert_called_once_with() + assert manager._kda_replay is None + state = manager._speculative_state + if requested_num_spec is None: + selector.assert_called_once_with(manager.spec_config, manager_supports_replay=True) + else: + selector.assert_not_called() + + manager._state_layout = object() + manager.all_ssm_states, manager.all_conv_states = [], [] + bind, update, reset = Mock(), Mock(), Mock() + monkeypatch.setattr(state, "bind", bind) + monkeypatch.setattr(state, "update", update) + monkeypatch.setattr(state, "reset_slots", reset) + manager._setup_model_state() + bind.assert_called_once_with( + manager._state_layout, manager.all_ssm_states, manager.all_conv_states + ) + batch, slots, host_slots = object(), object(), [0] + manager._update_speculative_state(batch) + update.assert_called_once_with(batch) + manager._reset_model_slots(slots, host_slots) + reset.assert_called_once_with(slots, host_slots) + + shutdown = Mock(wraps=state.shutdown) + monkeypatch.setattr(state, "shutdown", shutdown) + manager._shutdown_model_state() + shutdown.assert_called_once_with() + manager._shutdown_model_state() + assert shutdown.call_count == 2 + if num_spec is not None: + assert not hasattr(manager, "_speculative_state") + + +def test_base_manager_does_not_own_speculative_state(): + forbidden = { + "_speculative_state", + "IntermediateState", + "intermediate_ssm_states", + "intermediate_conv_states", + } + for attribute in vars(MambaHybridCacheManagerV2).values(): + function = attribute.__func__ if isinstance(attribute, classmethod) else attribute + if isinstance(function, FunctionType): + assert not forbidden.intersection(function.__code__.co_names), function.__name__ + + manager = object.__new__(MambaHybridCacheManagerV2) + manager.spec_config = None + manager._initialize_model_state() + manager._setup_model_state() + manager._reset_model_slots(object(), [0]) + manager._shutdown_model_state() + assert not forbidden.intersection(vars(manager)) + + +def test_base_manager_requires_model_for_speculative_decoding(): + manager = object.__new__(MambaHybridCacheManagerV2) + manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) + with pytest.raises(NotImplementedError, match="model-specific"): + manager._initialize_model_state() + with pytest.raises(NotImplementedError, match="model manager"): + manager._update_speculative_state(object()) + + +@pytest.mark.parametrize( + "manager_cls", + [NemotronHybridCacheManagerV2, Qwen35HybridCacheManagerV2, Qwen4ExpHybridCacheManagerV2], +) +@pytest.mark.parametrize("mode", ["plain", "intermediate", "replay"]) +def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): + manager = object.__new__(manager_cls) + manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) if mode != "plain" else None + manager._requested_replay = mode == "replay" + _prepare_model_layout(manager) + manager._initialize_model_state() + state = manager._speculative_state + assert isinstance(state, IntermediateState) == (mode != "replay") + assert isinstance(state, ReplayHistory) == (mode == "replay") + + manager._state_layout = object() + manager.local_num_mamba_layers = 0 + if manager_cls is Qwen4ExpHybridCacheManagerV2: + manager._ple_conv_states = {} + manager._ple_ngram_contexts = {} + manager.mamba_layer_offsets = {42: 0} + conv, ssm, payload = object(), object(), object() + manager.all_conv_states, manager.all_ssm_states = [conv], [ssm] + bind, reset, update, layer_cache = Mock(), Mock(), Mock(), Mock(return_value=payload) + monkeypatch.setattr(state, "bind", bind) + monkeypatch.setattr(state, "reset_slots", reset) + monkeypatch.setattr(state, "update", update) + monkeypatch.setattr(state, "make_layer_cache", layer_cache) + manager._setup_model_state() + bind.assert_called_once_with( + manager._state_layout, manager.all_ssm_states, manager.all_conv_states + ) + cache = manager.mamba_layer_cache(42) + if mode == "plain": + assert cache.conv is conv and cache.temporal is ssm + layer_cache.assert_not_called() + else: + assert cache is payload + layer_cache.assert_called_once_with(0, conv, ssm) + + if mode != "replay": + state.intermediate_ssm = [ssm] if mode == "intermediate" else None + state.intermediate_conv = [conv] if mode != "plain" else None + assert manager.intermediate_ssm_states == ([ssm] if mode == "intermediate" else None) + assert manager.intermediate_conv_states == ([conv] if mode != "plain" else None) + batch, slots, host_slots = object(), object(), [0] + manager._reset_model_slots(slots, host_slots) + reset.assert_called_once_with(slots, host_slots) + manager._update_speculative_state(batch) + update.assert_called_once_with(batch) + + shutdown = Mock(wraps=state.shutdown) + monkeypatch.setattr(state, "shutdown", shutdown) + manager._shutdown_model_state() + shutdown.assert_called_once_with() + assert manager.intermediate_ssm_states is None and state.intermediate_conv is None + if mode == "replay": + assert not hasattr(state, "intermediate_ssm") + + +def test_ple_adds_buffers_without_overriding_snapshot_or_budget_solver(): + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager._mamba_layer_mask = [True, True] + manager.mamba_pp_layers = [0, 1] + manager.ssm_bytes, manager.conv_bytes = 128, 288 + params = Qwen4ExpPLECacheParams( + ple_layer_mask=[False, True], + num_ple_layers=1, + short_conv_channels=16, + short_conv_state_len=6, + ngram_context_len=2, + conv_state_dtype=torch.bfloat16, + ) + manager._init_qwen4_exp_ple_geometry(params, 2, torch.bfloat16) + core = manager._get_recurrent_buffer_configs(0) + ple = manager._get_recurrent_buffer_configs(1) + assert [buffer.role for buffer in core] == [MambaRole.SSM_STATE, MambaRole.CONV_STATE] + assert [buffer.role for buffer in ple[2:]] == [PLE_NGRAM_CONTEXT, PLE_CONV_STATE] + assert manager._mamba_state_bytes_per_slot() == 2 * (128 + 288) + 16 * 6 * 2 + 2 * 8 + + +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("page_index_scale", [1, 2]) +def test_warmup_cleanup_clears_registered_persistent_roles_only( + monkeypatch, replay, page_index_scale +): + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + mamba_cache_manager_v2 as module, + ) + + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager._mamba_layer_mask = [False, True] + manager.pp_layers = manager.mamba_pp_layers = [1] + manager.layer_offsets = manager.mamba_layer_offsets = {1: 0} + manager.local_num_mamba_layers = 1 + manager._state_layout = _layout(layers=(1,)) + manager.max_batch_size = manager._state_layout.max_batch_size + manager.mapping = manager._state_layout.mapping + manager.ssm_state_shape = list(manager._state_layout.ssm_state_shape) + manager.conv_state_shape = list(manager._state_layout.conv_state_shape) + manager.ssm_state_dtype = manager.conv_state_dtype = torch.float32 + manager.ssm_bytes, manager.conv_bytes = 128, 288 + manager._num_reserved_dummy_slots = 0 + manager.kv_cache_config = SimpleNamespace() + manager._minimum_live_gpu_quota = lambda: 0 + manager._init_qwen4_exp_ple_geometry( + Qwen4ExpPLECacheParams( + ple_layer_mask=[False, True], + num_ple_layers=1, + short_conv_channels=16, + short_conv_state_len=6, + ngram_context_len=2, + conv_state_dtype=torch.bfloat16, + ), + 2, + torch.float32, + ) + manager._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), buffers=[BufferConfig(role=Role.KEY, size=256)] + ) + ], + initial_pool_ratio=[1.0], + ) + ) + raw_buffers = {} + + def wrap_buffer(role, dtype, shape): + raw_buffers[role] = torch.full(shape, 7, dtype=dtype) + return raw_buffers[role] + + manager.impl = SimpleNamespace( + get_mem_pool_base_address=lambda layer, role, mode: role, + get_page_index_upper_bound=lambda layer, role: 5 * page_index_scale, + get_page_index_scale=lambda layer, role: page_index_scale, + ) + monkeypatch.setattr(module, "TensorWrapper", wrap_buffer) + monkeypatch.setattr(module, "convert_to_torch_tensor", lambda tensor: tensor) + monkeypatch.setattr(torch.cuda, "current_device", lambda: "cpu") + monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) + manager._setup_states() + manager._speculative_state = GDNReplayState(3) if replay else IntermediateState() + manager._setup_model_state() + state = manager._speculative_state + scratch_before = {} + for name, tensor in vars(state).items(): + if isinstance(tensor, torch.Tensor): + if tensor.is_floating_point(): + tensor.fill_(torch.nan) + scratch_before[name] = tensor.clone() + + # Invalid scratch values must neither trigger the persistent-cache check + # nor be modified by warmup cleanup (especially the arange index buffer). + assert not manager.check_invalid_values_in_kv_cache() + manager._ple_conv_states[1][0, 0, 0] = torch.nan + assert manager.check_invalid_values_in_kv_cache(fill_with_zero=True) + for tensor in ( + manager.all_ssm_states[0], + manager.all_conv_states[0], + manager._ple_conv_states[1], + manager._ple_ngram_contexts[1], + ): + assert torch.count_nonzero(tensor) == 0 + for name, expected in scratch_before.items(): + torch.testing.assert_close(getattr(state, name), expected, equal_nan=True) + if page_index_scale > 1: + # The view must not clear a different layer's interleaved sub-pages. + for raw in raw_buffers.values(): + assert torch.all(raw[1::page_index_scale] == 7) + + +@pytest.mark.parametrize("affine", [False, True]) +def test_gdn_bind_accepts_affine_and_indirect_views(affine): + manager = _state_views(layers=3) + if not affine: + backing = torch.zeros(4, 5, 2, 4, 4) + manager.all_ssm_states = [backing[0], backing[1], backing[3]] + state = GDNReplayState(3) + state.bind(_layout(layers=(0, 1, 2)), manager.all_ssm_states, manager.all_conv_states) + assert (state._state_strides is not None) == affine + assert (state._state_descriptors is not None) != affine + assert state.old_x.shape == (3, 5, 2, 16, 2, 4) + owners = [] + for cls in (Qwen35HybridCacheManagerV2, Qwen4ExpHybridCacheManagerV2): + owner = object.__new__(cls) + owner._speculative_state = state + owners.append(owner) + assert owner.use_gdn_cached_replay_all_layer_commit + state.shutdown() + assert state._state_descriptors is None + assert all(not owner.use_gdn_cached_replay_all_layer_commit for owner in owners) + + +@pytest.mark.parametrize("state", [IntermediateState(), GDNReplayState(3), KDAReplayState(2)]) +def test_no_local_recurrent_layers_allocate_no_algorithm_buffers(state): + state.bind(_layout(layers=(), slot_capacity=0), [], []) + assert not any(isinstance(value, torch.Tensor) for value in vars(state).values()) + + +def test_kda_acceptance_reset_relocation_and_transfer(): + manager = _state_views() + state = KDAReplayState(2) + state.bind(_layout(), manager.all_ssm_states, manager.all_conv_states) + state.prev_num_accepted_tokens.copy_(torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32)) + for buffer in state._layer_replay_buffers(): + for slot in range(5): + buffer[:, slot].fill_(slot + 10) + before = [buffer.clone() for buffer in state._layer_replay_buffers()] + state.relocate_slots([0, 1], [1, 2]) + assert state.prev_num_accepted_tokens.tolist() == [1, 1, 2, 4, 5] + for old, new in zip(before, state._layer_replay_buffers()): + torch.testing.assert_close(new[:, 1], old[:, 0]) + torch.testing.assert_close(new[:, 2], old[:, 1]) + batch = MambaAcceptanceBatch( + attention_metadata=None, + num_contexts=0, + num_generations=2, + num_accepted_tokens=torch.tensor([3, 2]), + accepted_positions=torch.tensor([2, 1]), + source_state_indices=torch.arange(2), + destination_state_indices=torch.tensor([1, 2]), + is_dummy_request=torch.tensor([False, True]), + ) + state.update(batch) + assert state.prev_num_accepted_tokens.tolist() == [1, 2, 2, 4, 5] + state.reset_slots(torch.tensor([1]), [1]) + assert state.prev_num_accepted_tokens[1] == 0 + for conv in manager.all_conv_states: + conv[2].copy_(torch.arange(72).reshape(24, 3)) + state.seed_transferred_slots(torch.tensor([2])) + for layer, conv in enumerate(manager.all_conv_states): + for section, replay in enumerate((state.kda_conv_q, state.kda_conv_k, state.kda_conv_v)): + torch.testing.assert_close( + replay[layer, 2, :, :3], conv[2, section * 8 : (section + 1) * 8] + ) + assert torch.count_nonzero(replay[layer, 2, :, 3:]) == 0 + assert state.prev_num_accepted_tokens[2] == 0 + assert torch.count_nonzero(state.kda_qkg_cache[:, 2]) == 0 + + +def test_kda_scratch_cost_matches_allocated_slot_buffers(): + layout = _layout() + state = KDAReplayState(2) + state.bind(layout, _state_views().all_ssm_states, _state_views().all_conv_states) + buffers = (*state._layer_replay_buffers(), state.prev_num_accepted_tokens) + actual = sum(buffer.numel() * buffer.element_size() for buffer in buffers) + estimated = sum(state.bytes_per_slot(layout, layer_id) for layer_id in layout.mamba_pp_layers) + assert actual == estimated * layout.slot_capacity + + +def test_mamba2_seed_lifecycle_does_not_require_speculative_decoding(): + layout = replace( + _layout(), spec_config=None, stochastic_rounding=True, ssm_state_dtype=torch.float16 + ) + first, second = Mamba2State(), Mamba2State() + first.bind(layout, _state_views().all_ssm_states, _state_views().all_conv_states) + second.bind(layout, _state_views().all_ssm_states, _state_views().all_conv_states) + torch.testing.assert_close(first.rand_seed, second.rand_seed) + assert torch.all(first.rand_seed > 0) + assert first.intermediate_ssm is None + original = first.rand_seed.clone() + first.reset_slots(torch.tensor([1]), [1]) + assert first.rand_seed[1] != original[1] + torch.testing.assert_close(first.rand_seed[[0, 2, 3, 4]], original[[0, 2, 3, 4]]) + + +def test_replay_and_intermediate_are_independent_algorithms(): + assert ReplayHistory.__bases__ == (object,) + assert GDNReplayState.__bases__ == (ReplayHistory,) + assert Mamba2State.__bases__ == (IntermediateState,) + assert not issubclass(ReplayHistory, IntermediateState) + assert not issubclass(GDNReplayState, Mamba2State) + assert not hasattr(ReplayHistory(3), "intermediate_ssm") + assert not hasattr(GDNReplayState(3), "intermediate_ssm") + state = IntermediateState() + assert not {"rand_seed", "_seed_request_counter", "_seed_rank_offset"}.intersection(vars(state)) + assert not hasattr(state, "_promote_conv") + + +@pytest.mark.parametrize("replay_cls", [ReplayHistory, GDNReplayState]) +def test_mamba_seed_reset_matches_between_intermediate_and_replay(replay_cls): + layout = replace(_layout(), stochastic_rounding=True, seed_rank_offset=17) + views = _state_views() + intermediate, replay = Mamba2State(), replay_cls(3) + for state in (intermediate, replay): + state.bind(layout, views.all_ssm_states, views.all_conv_states) + torch.testing.assert_close(intermediate.rand_seed, replay.rand_seed) + original = intermediate.rand_seed.clone() + replay.prev_num_accepted_tokens.fill_(7) + replay.cache_buf_idx.fill_(1) + for buffer in (replay.old_x, replay.old_B, replay.old_dt, replay.old_dA_cumsum): + buffer.fill_(4) + slots = torch.tensor([3, 1]) + for state in (intermediate, replay): + state.reset_slots(slots, [3, 1]) + torch.testing.assert_close(intermediate.rand_seed, replay.rand_seed) + assert torch.all(intermediate.rand_seed[slots] != original[slots]) + torch.testing.assert_close(intermediate.rand_seed[[0, 2, 4]], original[[0, 2, 4]]) + assert replay.prev_num_accepted_tokens.tolist() == [7, 0, 7, 0, 7] + assert replay.cache_buf_idx.tolist() == [1, 0, 1, 0, 1] + for buffer in (replay.old_x, replay.old_B, replay.old_dt, replay.old_dA_cumsum): + assert torch.count_nonzero(buffer[:, slots]) == 0 + assert torch.all(buffer[:, [0, 2, 4]] == 4) + for state in (intermediate, replay): + cache = state.make_layer_cache(0, views.all_conv_states[0], views.all_ssm_states[0]) + assert cache.mamba_ssm_rand_seed is state.rand_seed + state.shutdown() + assert state.rand_seed is None + + +@pytest.mark.parametrize( + "state_cls", [IntermediateState, Mamba2State, ReplayHistory, GDNReplayState] +) +def test_accepted_state_promotion_preserves_algorithm_boundary(monkeypatch, state_cls): + from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager + + replay = state_cls in (ReplayHistory, GDNReplayState) + state = state_cls(3) if replay else state_cls() + views = _state_views() + state.bind(_layout(), views.all_ssm_states, views.all_conv_states) + state.intermediate_conv.copy_( + torch.arange(state.intermediate_conv.numel()).reshape(state.intermediate_conv.shape) + ) + if not replay: + state.intermediate_ssm.fill_(7) + batch = MambaAcceptanceBatch( + attention_metadata=None, + num_contexts=0, + num_generations=2, + num_accepted_tokens=torch.tensor([1, 3]), + accepted_positions=torch.tensor([0, 2]), + source_state_indices=torch.tensor([1, 0]), + destination_state_indices=torch.tensor([4, 2]), + is_dummy_request=None, + ) + calls = [] + + def promote(destination, source, source_indices, positions, destination_indices): + calls.append(destination.data_ptr()) + for src, pos, dst in zip( + source_indices.tolist(), positions.tolist(), destination_indices.tolist() + ): + destination[:, dst].copy_(source[:, src, pos]) + + monkeypatch.setattr(mamba_cache_manager, "_promote_mamba_state_triton", promote) + state.update(batch) + assert len(calls) == (2 if replay else 4) + for layer, conv in enumerate(views.all_conv_states): + torch.testing.assert_close(conv[4], state.intermediate_conv[layer, 1, 0]) + torch.testing.assert_close(conv[2], state.intermediate_conv[layer, 0, 2]) + assert torch.count_nonzero(conv[[0, 1, 3]]) == 0 + for ssm in views.all_ssm_states: + if replay: + assert torch.count_nonzero(ssm) == 0 + else: + assert torch.all(ssm[[4, 2]] == 7) + assert torch.count_nonzero(ssm[[0, 1, 3]]) == 0 + cache = state.make_layer_cache(0, views.all_conv_states[0], views.all_ssm_states[0]) + assert (cache.intermediate_ssm is None) == replay + if replay: + assert state.prev_num_accepted_tokens.tolist() == [0, 0, 3, 0, 1] + assert not hasattr(state, "intermediate_ssm") + state.shutdown() + + +@pytest.mark.parametrize( + "state", + [IntermediateState(), Mamba2State(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)], +) +def test_algorithm_borrows_only_tensor_views_and_releases_them(state): + views = _state_views() + state.bind(_layout(), views.all_ssm_states, views.all_conv_states) + assert state._conv_states[0] is views.all_conv_states[0] + assert not hasattr(state, "manager") + if isinstance(state, (IntermediateState, ReplayHistory)): + assert state._ssm_states[0] is views.all_ssm_states[0] + assert not hasattr(state, "_publish_compatibility_views") + state.shutdown() + state.shutdown() + assert state._conv_states == () + if isinstance(state, (IntermediateState, ReplayHistory)): + assert state._ssm_states == () + assert not any(isinstance(value, torch.Tensor) for value in vars(state).values()) + # Releasing borrowed references must not mutate the persistent pool. + assert views.all_conv_states[0].shape == (5, 24, 3) + + +@pytest.mark.parametrize( + "cls", [IntermediateState, Mamba2State, ReplayHistory, GDNReplayState, KDAReplayState] +) +def test_algorithm_methods_do_not_depend_on_manager(cls): + for function in vars(cls).values(): + if not isinstance(function, FunctionType): + continue + code = function.__code__ + assert "manager" not in code.co_varnames[: code.co_argcount] + assert not { + "_request_id_to_state_index", + "_request_id_to_is_dummy", + "_publish_compatibility_views", + "_publish_replay_views", + }.intersection(code.co_names) + + +def test_qwen4_methods_do_not_access_raw_pool(): + for function in vars(Qwen4ExpHybridCacheManagerV2).values(): + if isinstance(function, FunctionType): + assert not {"impl", "TensorWrapper", "convert_to_torch_tensor"}.intersection( + function.__code__.co_names + ) + + +@pytest.mark.parametrize( + "manager_cls", + [ + NemotronHybridCacheManagerV2, + Qwen35HybridCacheManagerV2, + Qwen4ExpHybridCacheManagerV2, + KimiK3HybridCacheManagerV2, + ], +) +def test_model_compatibility_properties_follow_owned_state(manager_cls): + manager = object.__new__(manager_cls) + manager._speculative_state = IntermediateState() + if manager_cls is KimiK3HybridCacheManagerV2: + manager._kda_replay = None + fields = { + "intermediate_state_indices": "intermediate_indices", + "intermediate_ssm_states": "intermediate_ssm", + "intermediate_conv_states": "intermediate_conv", + } + for property_name, field in fields.items(): + value = torch.zeros(1) + setattr(manager._speculative_state, field, value) + assert getattr(manager, property_name) is value + assert property_name not in vars(manager) + with pytest.raises(AttributeError): + setattr(manager, property_name, value) + manager._speculative_state.shutdown() + assert all(getattr(manager, name) is None for name in fields) + + +def test_k3_resolves_transfer_requests_before_calling_algorithm(): + manager = object.__new__(KimiK3HybridCacheManagerV2) + manager._request_id_to_state_index = {101: 3, 202: 1} + replay = KDAReplayState(2) + replay.prev_num_accepted_tokens = torch.zeros(5, dtype=torch.int32) + replay.seed_transferred_slots = Mock() + manager._kda_replay = replay + manager.on_state_transfer_complete([101, 999, 202, 101]) + slots = replay.seed_transferred_slots.call_args.args[0] + assert slots.tolist() == [1, 3] + assert slots.dtype == torch.long + replay.seed_transferred_slots.reset_mock() + manager.on_state_transfer_complete([999]) + replay.seed_transferred_slots.assert_not_called() + + +@pytest.mark.parametrize("draft_lengths", [[], [0, 0], [1, 0], [1, 1]]) +def test_k3_host_acceptance_handles_empty_mixed_and_completed_requests(draft_lengths): + manager = object.__new__(KimiK3HybridCacheManagerV2) + replay = KDAReplayState(2) + replay.prev_num_accepted_tokens = torch.tensor([5, 6], dtype=torch.int32) + manager._kda_replay = replay + manager._request_id_to_state_index = {101: 0} + manager._request_id_to_is_dummy = {} + requests = [ + SimpleNamespace( + py_request_id=101 + i, py_draft_tokens=[1] * length, py_num_accepted_draft_tokens=1 + ) + for i, length in enumerate(draft_lengths) + ] + batch = SimpleNamespace(generation_requests=requests) + if draft_lengths == [1, 0]: + with pytest.raises(RuntimeError, match="Mixed drafted/undrafted"): + manager._record_replay_request_acceptance(batch) + else: + manager._record_replay_request_acceptance(batch) + assert replay.prev_num_accepted_tokens.tolist() == ( + [1, 6] if draft_lengths == [1, 1] else [5, 6] + ) + + +def _slot_view_manager( + monkeypatch, *, scale=6, role_count=2, page_shape=(4, 2, 3), dtype=torch.float32 +): + manager = object.__new__(KVCacheManagerV2) + manager.layer_offsets = {42: 0} + page_elements = 1 + for dim in page_shape: + page_elements *= dim + page_bytes = page_elements * dtype.itemsize + # First role belongs to a later coalesced layer, not the pool's first page. + layer_offset = 2 if scale > 2 else 0 + backing = torch.arange((3 * scale + layer_offset) * page_elements, dtype=torch.float32).to( + dtype + ) + addresses = { + Role.KEY: 1000 + layer_offset * page_bytes, + Role.VALUE: 1000 + (layer_offset + 1) * page_bytes, + } + converter = SimpleNamespace(scale=scale, expansion=1, layer_offset=layer_offset) + manager.impl = SimpleNamespace( + get_mem_pool_base_address=lambda layer, role, mode: addresses[role], + get_page_stride=lambda layer, role: page_bytes, + get_page_index_converter=lambda layer, role: converter, + get_page_index_upper_bound=lambda layer, role: 3 * scale - layer_offset, + ) + globals_ = KVCacheManagerV2._get_slot_role_view.__globals__ + monkeypatch.setitem(globals_, "TensorWrapper", lambda address, dtype, shape: (address, shape)) + + def convert(wrapper): + address, shape = wrapper + count = 1 + for dim in shape: + count *= dim + offset = (address - 1000) // dtype.itemsize + return backing[offset : offset + count].reshape(shape) + + monkeypatch.setitem(globals_, "convert_to_torch_tensor", convert) + roles = (Role.KEY, Role.VALUE)[:role_count] + return manager, roles, list(page_shape), converter, backing + + +@pytest.mark.parametrize("role_count,scale", [(1, 1), (2, 2), (1, 6), (2, 6)]) +@pytest.mark.parametrize("page_shape", [(4, 2, 3), (2, 4, 3)]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_slot_role_view_preserves_coalesced_stride( + monkeypatch, role_count, scale, page_shape, dtype +): + manager, roles, shape, converter, backing = _slot_view_manager( + monkeypatch, scale=scale, role_count=role_count, page_shape=page_shape, dtype=dtype + ) + view = manager._get_slot_role_view(42, roles, dtype=dtype, page_shape=shape) + assert view.shape == (3, role_count, *page_shape) + assert view.stride(0) == scale * 24 + assert view.stride(1) == 24 + assert view.data_ptr() == backing.data_ptr() + converter.layer_offset * 24 * dtype.itemsize + view[1, 0].fill_(7) + start = (converter.layer_offset + scale) * 24 + assert torch.all(backing[start : start + 24] == 7) + + +@pytest.mark.parametrize("invalid", ["stride", "adjacency", "scale", "expansion", "page_count"]) +def test_slot_role_view_rejects_incompatible_physical_layout(monkeypatch, invalid): + manager, roles, shape, converter, _ = _slot_view_manager(monkeypatch) + if invalid == "stride": + manager.impl.get_page_stride = lambda layer, role: 1 + elif invalid == "adjacency": + manager.impl.get_mem_pool_base_address = lambda layer, role, mode: 1000 + elif invalid == "scale": + converter.scale = 1 + elif invalid == "expansion": + converter.expansion = 2 + else: + manager.impl.get_page_index_upper_bound = lambda layer, role: 1 + with pytest.raises(RuntimeError): + manager._get_slot_role_view(42, roles, dtype=torch.float32, page_shape=shape) + + +def test_qwen4_uses_kv_owned_slot_views(monkeypatch): + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager.qsa_position_layer_id = 42 + manager.tokens_per_block = 4 + manager._get_slot_role_view = Mock(return_value=torch.zeros(3, 1, 4, 3, dtype=torch.int32)) + assert manager.get_qsa_position_buffer().shape == (3, 4, 3) + manager._get_slot_role_view.assert_called_once() + manager.qsa_position_layer_id = None + manager._get_slot_role_view.reset_mock() + assert manager.get_qsa_position_buffer() is None + manager._get_slot_role_view.assert_not_called() + + +@pytest.mark.parametrize("kv_layout", ["NHD", "HND"]) +def test_qwen4_main_kv_view_supplies_only_logical_geometry(monkeypatch, kv_layout): + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + namespace = Qwen4ExpHybridCacheManagerV2.get_buffers.__globals__ + manager._qsa_enabled = True + manager.dtype = next(iter(namespace["QSA_SPARSE_KV_CACHE_DTYPES"])) + manager.kv_cache_type = namespace["CacheTypeCpp"].SELF + manager.layer_offsets = {42: 0} + manager.head_dim_per_layer = [3] + manager.num_kv_heads_per_layer = [2] + manager.tokens_per_block = 4 + monkeypatch.setitem(namespace, "binding_to_torch_dtype", lambda dtype: torch.bfloat16) + sentinel = object() + manager._get_slot_role_view = Mock(return_value=sentinel) + assert manager.get_buffers(42, kv_layout) is sentinel + manager._get_slot_role_view.assert_called_once_with( + 42, + (Role.KEY, Role.VALUE), + dtype=torch.bfloat16, + page_shape=[4, 2, 3] if kv_layout == "NHD" else [2, 4, 3], + ) + assert manager.get_buffers(99, kv_layout) is None + manager.kv_cache_type = namespace["CacheTypeCpp"].SELFKONLY + with pytest.raises(NotImplementedError, match="both K and V"): + manager.get_buffers(42, kv_layout) + + +@pytest.mark.parametrize("invalid", [None, "pool", "scale", "expansion"]) +def test_attention_slot_mapping_validates_backend_contract(invalid): + manager = object.__new__(KVCacheManagerV2) + manager.layer_offsets = {42: 0} + manager.layer_to_pool_mapping_dict = {0: 0} + manager.num_attention_op_pools = 1 + converter = SimpleNamespace(scale=6, expansion=1) + manager.impl = SimpleNamespace(get_page_index_converter=lambda layer, role: converter) + if invalid == "pool": + manager.layer_to_pool_mapping_dict[0] = 1 + elif invalid == "scale": + converter.scale = 0 + elif invalid == "expansion": + converter.expansion = 2 + if invalid is None: + assert manager._get_attention_slot_mapping(42) == (0, 6) + else: + with pytest.raises(RuntimeError): + manager._get_attention_slot_mapping(42) + + +def test_qwen4_requires_one_mapping_but_does_not_read_physical_pool(): + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager.qsa_position_layer_id = 42 + manager.qsa_sparse_layer_ids = [42, 43, 99] + manager.layer_offsets = {42: 0, 43: 1} + manager._get_attention_slot_mapping = Mock(return_value=(0, 6)) + assert manager.get_qsa_attention_pool_layout() == (0, 6) + assert all(call.args != (99,) for call in manager._get_attention_slot_mapping.call_args_list) + manager._get_attention_slot_mapping = lambda layer: (0, 6 if layer == 42 else 4) + with pytest.raises(RuntimeError, match="do not share"): + manager.get_qsa_attention_pool_layout() + + +def test_constructor_validates_before_pool_and_binds_after_slot_capacity(monkeypatch): + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( + mamba_cache_manager_v2 as module, + ) + + events = [] + zeros, arange = torch.zeros, torch.arange + + def cpu_zeros(*args, **kwargs): + kwargs.pop("pin_memory", None) + if kwargs.get("device") == "cuda": + kwargs["device"] = "cpu" + return zeros(*args, **kwargs) + + def cpu_arange(*args, **kwargs): + if kwargs.get("device") == "cuda": + kwargs["device"] = "cpu" + return arange(*args, **kwargs) + + def create_pool(self, *args, **kwargs): + assert events == ["validate"] + assert self._state_layout.slot_capacity is None + assert len(self._get_recurrent_buffer_configs(0)) == 4 + events.append("pool") + self.max_batch_size = kwargs["max_batch_size"] + self.mapping = kwargs["mapping"] + self.layer_offsets = {0: 0, 1: 1} + self.impl = SimpleNamespace( + get_layer_group_id=lambda layer: 0, + get_page_index_scale=lambda layer, role: 1, + get_page_index_upper_bound=lambda layer, role: 5, + ) + + class ObservedQwen4(Qwen4ExpHybridCacheManagerV2): + def _initialize_model_state(self): + super()._initialize_model_state() + events.append("validate") + + def _get_state_buffer(self, local_layer_idx, role, dtype, state_shape): + assert events[:2] == ["validate", "pool"] + assert self._state_layout.slot_capacity == 5 + events.append(role) + return zeros([5, *state_shape], dtype=dtype) + + monkeypatch.setattr(module.KVCacheManagerV2, "__init__", create_pool) + monkeypatch.setattr(module, "get_pp_layers", lambda *args, **kwargs: ([0, 1], None)) + monkeypatch.setattr(torch, "zeros", cpu_zeros) + monkeypatch.setattr(torch, "arange", cpu_arange) + config = SimpleNamespace(enable_block_reuse=False) + config.model_copy = lambda **kwargs: config + params = Qwen4ExpPLECacheParams( + ple_layer_mask=[True, False], + num_ple_layers=1, + short_conv_channels=16, + short_conv_state_len=6, + ngram_context_len=2, + conv_state_dtype=torch.bfloat16, + ) + manager = ObservedQwen4( + 4, + 4, + 2, + 2, + 4, + 1, + [True, False], + torch.float32, + torch.float32, + config, + module.CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=2, + head_dim=4, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=2, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + layer_mask=[False, True], + use_replay_state_update=False, + qwen4_exp_ple_cache_params=params, + ) + assert events == [ + "validate", + "pool", + MambaRole.SSM_STATE, + MambaRole.CONV_STATE, + PLE_CONV_STATE, + PLE_NGRAM_CONTEXT, + ] + assert manager.ple_layer_cache(0)[0].shape == (5, 16, 6) diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index be6c1a3a29f6..6246c08cc90f 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -263,12 +263,12 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( monkeypatch.setattr(_util, "is_kimi_linear", lambda config: False) monkeypatch.setattr(_util, "is_mla", lambda config: False) monkeypatch.setattr(_util, "is_nemotron_hybrid", lambda config: False) - monkeypatch.setattr(_util, "is_qwen3_hybrid", lambda config: True) + monkeypatch.setattr(_util, "is_qwen3_hybrid", lambda config: False) + monkeypatch.setattr(_util, "is_qwen4_exp", lambda config: True) monkeypatch.setattr( _util, "extract_mamba_kv_cache_params", lambda *args, **kwargs: mamba_params ) monkeypatch.setattr(_util, "get_sm_version", lambda: 103) - monkeypatch.setattr(_util, "is_gdn_replay_enabled", lambda: False) monkeypatch.setattr(MambaHybridCacheManagerV2, "__init__", lambda self, *args, **kwargs: None) manager = _create_kv_cache_manager( @@ -296,9 +296,11 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( assert manager.qsa_index_kv_heads == 1 -def test_qsa_cache_manager_requires_sparse_config() -> None: - with pytest.raises(ValueError, match="sparse_attention_config is required"): - QSAMambaHybridCacheManagerV2(layer_mask=[True]) +def test_qwen4_manager_also_supports_dense_attention(monkeypatch) -> None: + monkeypatch.setattr(MambaHybridCacheManagerV2, "__init__", lambda self, *args, **kwargs: None) + manager = QSAMambaHybridCacheManagerV2(layer_mask=[True]) + assert not manager._qsa_enabled + assert manager._extra_buffers_per_layer(tokens_per_block=32) == {} @pytest.mark.parametrize("dtype", (DataType.NVFP4, DataType.FLOAT)) @@ -307,6 +309,7 @@ def test_qsa_cache_manager_delegates_regular_kv_layout( ) -> None: manager = object.__new__(QSAMambaHybridCacheManagerV2) manager.dtype = dtype + manager._qsa_enabled = True sentinel = torch.empty(0, dtype=torch.int8) calls: list[tuple[object, int, str]] = [] diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 5da980b03c4c..d97ed7cd1adb 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -381,7 +381,6 @@ def test_hybrid_and_ple_layout_is_derived_from_config() -> None: from tensorrt_llm._torch.configs import Qwen4ExpTextConfig from tensorrt_llm._torch.pyexecutor.config_utils import ( extract_mamba_kv_cache_params, - extract_qwen4_exp_ple_cache_params, get_qwen3_hybrid_layer_types, ) @@ -394,6 +393,10 @@ def test_hybrid_and_ple_layout_is_derived_from_config() -> None: ] mamba = extract_mamba_kv_cache_params(config) assert mamba.num_mamba_layers == 3 + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + extract_qwen4_exp_ple_cache_params, + ) + ple = extract_qwen4_exp_ple_cache_params(config) assert ple.ple_layer_mask == [False, True, False, False] assert ple.short_conv_channels == 4 * 128 @@ -404,7 +407,7 @@ def test_hybrid_and_ple_layout_is_derived_from_config() -> None: @pytest.mark.parametrize("ple_layer_ids", [[2, 2], [0], [5], [True], [[2]], [1, 2]]) def test_ple_layer_ids_reject_ambiguous_or_out_of_range_values(ple_layer_ids) -> None: from tensorrt_llm._torch.configs import Qwen4ExpTextConfig - from tensorrt_llm._torch.pyexecutor.config_utils import get_qwen4_exp_ple_layer_mask + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import get_qwen4_exp_ple_layer_mask fields = _text_config_dict() fields["ple_layer_ids"] = ple_layer_ids @@ -416,7 +419,9 @@ def test_ple_layer_ids_reject_ambiguous_or_out_of_range_values(ple_layer_ids) -> def test_ple_cache_layout_excludes_separate_mtp_draft() -> None: from tensorrt_llm._torch.configs import Qwen4ExpTextConfig - from tensorrt_llm._torch.pyexecutor._util import _get_qwen4_exp_ple_cache_params + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + get_qwen4_exp_ple_cache_params as _get_qwen4_exp_ple_cache_params, + ) config = Qwen4ExpTextConfig.from_dict(_text_config_dict()) @@ -432,9 +437,9 @@ def test_ple_cache_layout_excludes_separate_mtp_draft() -> None: def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: from tensorrt_llm._torch.configs import Qwen4ExpTextConfig - from tensorrt_llm._torch.pyexecutor.config_utils import extract_qwen4_exp_ple_cache_params - from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - MambaHybridCacheManagerV2, + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + Qwen4ExpHybridCacheManagerV2, + extract_qwen4_exp_ple_cache_params, ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping @@ -447,10 +452,10 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: "max_batch_size": 2, "kv_cache_config": KvCacheConfig(enable_block_reuse=False), } - with_ple = MambaHybridCacheManagerV2.get_cache_size_per_token( + with_ple = Qwen4ExpHybridCacheManagerV2.get_cache_size_per_token( SimpleNamespace(pretrained_config=config, quant_config=None), **common ) - without_ple = MambaHybridCacheManagerV2.get_cache_size_per_token( + without_ple = Qwen4ExpHybridCacheManagerV2.get_cache_size_per_token( SimpleNamespace(pretrained_config=no_ple_config, quant_config=None), **common ) @@ -465,9 +470,10 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state() -> None: def test_ple_states_use_v2_lifecycle_buffers(monkeypatch) -> None: - from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - MambaHybridCacheManagerV2, - MambaRole, + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + PLE_CONV_STATE, + PLE_NGRAM_CONTEXT, + Qwen4ExpHybridCacheManagerV2, ) ngram_context = torch.full((12, 2), 11, dtype=torch.int64) @@ -477,20 +483,20 @@ def test_ple_states_use_v2_lifecycle_buffers(monkeypatch) -> None: def fake_get_state_buffer( self: object, local_layer_idx: int, - role: MambaRole, + role: object, dtype: torch.dtype, state_shape: list[int], ) -> torch.Tensor: del self requested.append((local_layer_idx, role, dtype, state_shape)) - if role == MambaRole.PLE_NGRAM_CONTEXT: + if role == PLE_NGRAM_CONTEXT: return ngram_context - if role == MambaRole.PLE_CONV_STATE: + if role == PLE_CONV_STATE: return conv_state raise AssertionError(f"unexpected role {role}") - monkeypatch.setattr(MambaHybridCacheManagerV2, "_get_state_buffer", fake_get_state_buffer) - manager = object.__new__(MambaHybridCacheManagerV2) + monkeypatch.setattr(Qwen4ExpHybridCacheManagerV2, "_get_state_buffer", fake_get_state_buffer) + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager._ple_layer_ids = [1] manager._ple_ngram_context_shape = [2] manager._ple_conv_state_shape = [16, 6] @@ -505,8 +511,8 @@ def fake_get_state_buffer( assert actual_conv is conv_state assert actual_context is ngram_context assert requested == [ - (0, MambaRole.PLE_CONV_STATE, torch.bfloat16, [16, 6]), - (0, MambaRole.PLE_NGRAM_CONTEXT, torch.int64, [2]), + (0, PLE_CONV_STATE, torch.bfloat16, [16, 6]), + (0, PLE_NGRAM_CONTEXT, torch.int64, [2]), ] diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 56fb548b8793..abc7b55a119e 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -15,6 +15,7 @@ """Unit tests for Mamba2 metadata preparation optimizations.""" from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch @@ -43,7 +44,6 @@ class _GdnReplayCacheManager: - use_replay_state_update = True use_gdn_cached_replay_all_layer_commit = True def __init__(self, prev_num_accepted_tokens, cache_buf_idx): @@ -59,6 +59,54 @@ def get_replay_state_update_metadata(self): ) +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("num_decodes", [0, 2]) +def test_replay_mode_uses_one_metadata_query(monkeypatch, enabled, num_decodes): + replay = SimpleNamespace( + prev_num_accepted_tokens=object(), + cache_buf_idx=object(), + replay_step_width=3, + replay_history_size=MIN_REPLAY_HISTORY_SIZE, + ) + + class Manager: + get_replay_state_update_metadata = Mock(return_value=replay if enabled else None) + + @property + def use_replay_state_update(self): + raise AssertionError("Replay mode must come from metadata, not the legacy flag") + + manager = Manager() + metadata = object.__new__(Mamba2Metadata) + metadata.state_indices = [7, 2, 9] + metadata.replay_work_items = object() + metadata.replay_n_writes = object() + build = Mock() + monkeypatch.setattr(mamba2_metadata, "_build_replay_work_items_torch", build) + Mamba2Metadata._prepare_replay_work_items(metadata, manager, 1 + num_decodes, 1) + manager.get_replay_state_update_metadata.assert_called_once_with() + assert metadata.replay_num_decodes == (num_decodes if enabled else 0) + if enabled and num_decodes: + build.assert_called_once_with( + [2, 9], + replay.prev_num_accepted_tokens, + replay.cache_buf_idx, + metadata.replay_work_items, + metadata.replay_n_writes, + replay.replay_step_width, + replay.replay_history_size, + ) + else: + build.assert_not_called() + + +def test_replay_metadata_optional_for_non_mamba_manager(): + metadata = object.__new__(Mamba2Metadata) + metadata.replay_num_decodes = 10 + metadata._prepare_replay_work_items(SimpleNamespace(), 3, 1) + assert metadata.replay_num_decodes == 0 + + def _torch_reference_work_items(state_indices, prev_num_accepted_tokens, cache_buf_idx): """Run the production ATen path into fresh buffers. @@ -140,8 +188,6 @@ def test_prepare_handles_tensor_cached_tokens(self): def test_prepare_replay_work_items_write_first(self): class ReplayCacheManager: - use_replay_state_update = True - def __init__(self): self.state_indices = [0, 3, 1, 4, 2] self.prev_num_accepted_tokens = torch.tensor( From f7b9d74d4e67d75320ce556a468f479d30832568 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 17 Sep 2026 02:08:46 +0000 Subject: [PATCH 02/25] [None][test] Fix model-owned Mamba manager coverage (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tests/integration/defs/accuracy/test_llm_api_pytorch.py | 8 ++++++-- .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 8 ++++---- 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 5090f29c118f..2b11b5281054 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6632,8 +6632,10 @@ def test_nvfp4_8gpus_mtp(self): enable_padding=True), disable_overlap_scheduler=False, moe_config=MoeConfig(backend="CUTLASS"), - decoding_config=mtp_config, + speculative_config=mtp_config, ) as llm: + assert isinstance(llm.args.speculative_config, MTPDecodingConfig) + assert llm.args.speculative_config.max_draft_len == 3 task = MMLU(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) @@ -6789,8 +6791,10 @@ def test_nvfp4_8gpus_mtp_custom_op(self, monkeypatch): enable_padding=True), disable_overlap_scheduler=False, moe_config=MoeConfig(backend="CUTLASS"), - decoding_config=mtp_config, + speculative_config=mtp_config, ) as llm: + assert isinstance(llm.args.speculative_config, MTPDecodingConfig) + assert llm.args.speculative_config.max_draft_len == 3 task = MMLU(self.MODEL_NAME) task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 4a1d2052f774..9e50610014d7 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -446,8 +446,8 @@ def test_v2_manager_rejects_model_type_kwarg() -> None: @pytest.mark.parametrize( ("use_v2", "enable_block_reuse", "expected"), [ - (True, False, NemotronHybridCacheManagerV2), - (True, True, NemotronHybridCacheManagerV2), + (True, False, Qwen35HybridCacheManagerV2), + (True, True, Qwen35HybridCacheManagerV2), (False, True, CppMambaHybridCacheManager), ("auto", True, CppMambaHybridCacheManager), ], @@ -671,7 +671,7 @@ def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, bac backend=backend, transceiver_runtime="PYTHON" ), ) - is NemotronHybridCacheManagerV2 + is Qwen35HybridCacheManagerV2 ) @@ -865,7 +865,7 @@ def test_qwen3_gdn_replay_uses_v2_preference( ) assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected_v2 - expected_manager = NemotronHybridCacheManagerV2 if expected_v2 else CppMambaHybridCacheManager + expected_manager = Qwen35HybridCacheManagerV2 if expected_v2 else CppMambaHybridCacheManager assert ( get_kv_cache_manager_cls( _hybrid_model_config(), From e8693e1ac5510c844c4449a1131f41ee1f2dbdca Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 17 Sep 2026 02:42:58 +0000 Subject: [PATCH 03/25] [None][fix] Handle globally combined K3 routed MoE output (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 3 + .../_torch/models/modeling_kimi_linear.py | 21 +++++-- .../_torch/moe/fused_moe/impl_contract.py | 3 + .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 17 +++--- .../_torch/moe/fused_moe/moe_resolution.py | 1 + .../_torch/moe/test_kimi_k3_situ_moe.py | 56 +++++++++++++++++++ tests/unittest/_torch/moe/test_moe_backend.py | 28 ++++++++++ 7 files changed, 118 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 2c28757cfee5..540b471583d7 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -251,6 +251,9 @@ class ModelConfig(Generic[TConfig]): # disables the optimization; hardware support is resolved by the backend. skip_correction_threshold: float = 0.0 moe_backend: str = 'CUTLASS' # options can be CUTLASS, TRTLLM + # Model composition reduces only local branches when a fused-communication + # backend returns an already globally combined routed output. + moe_handles_global_routed_output: bool = False # IF true, disables FC2+finalize fusion in CUTLASS MoE backend moe_disable_finalize_fusion: bool = False # If true, use low precision combine in MoE operations (only for NVFP4 quantization) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 38dad076d57f..c8f6fc279a04 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -54,7 +54,7 @@ column-sharded and down is row-sharded. Direct MoE-TP combines the shared hidden-width partial and routed latent partial into one all-reduce after the two streams join, then splits them before the routed norm/up projection. -Communication-backed routed paths keep the shared ``GatedMLP`` reduction, +Fused-communication routed paths reduce only the shared-expert partial, because their routed result is already combined. Under attention DP the shared experts stay replicated. ``lm_head`` uses the stock ``LMHead`` (vocab-sharded + gather), so logits are identical on all ranks. @@ -119,6 +119,7 @@ from ..modules.rms_norm import RMSNorm from ..modules.situ import SituAndMul from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, TRTLLMGenFusedMoE, create_moe +from ..moe.fused_moe.interface import MoESchedulerKind from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..utils import AuxStreamType from .modeling_speculative import SpecDecOneEngineForCausalLM @@ -1166,7 +1167,8 @@ def __init__( # Under attention DP each rank owns different tokens, so the shared # expert is replicated (TP size 1) and must not reduce across ranks. # Direct MoE-TP leaves both branches as partials for one concatenated - # all-reduce. + # all-reduce. A fused-communication backend has already combined the + # routed branch, so only its shared-expert partial needs reduction. use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1 routed_all_reduce = self.routed_experts.all_reduce if use_shared_tp and routed_all_reduce is None: @@ -1174,7 +1176,11 @@ def __init__( "Kimi K3 direct MoE tensor parallelism requires the " "ConfigurableMoE all-reduce even when reduce_results=False." ) - self._use_combined_all_reduce = use_shared_tp + routed_output_is_global = ( + self.routed_experts.backend.scheduler_kind == MoESchedulerKind.FUSED_COMM + ) + self._use_combined_all_reduce = use_shared_tp and not routed_output_is_global + self._reduce_shared_output = use_shared_tp and routed_output_is_global self.shared_experts = GatedMLP( hidden_size=cfg.hidden_size, intermediate_size=shared_intermediate, @@ -1363,6 +1369,7 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) routed_model_config.mapping = routed_mapping routed_model_config.moe_backend = model_config.moe_backend + routed_model_config.moe_handles_global_routed_output = True # MegaMoE uses this value as global DP SymmBuffer capacity, then # divides it by EP size for the per-rank allocation. Other backends # keep the user-configured value as their MoE chunking bound. @@ -1395,7 +1402,11 @@ def forward(self, hidden_states: torch.Tensor, all_rank_num_tokens=None) -> torc """``hidden_states``: ``[num_tokens, hidden_size]`` bf16.""" identity = hidden_states router_logits = self.gate.compute_logits(hidden_states) - moe_all_reduce = self.routed_experts.all_reduce if self._use_combined_all_reduce else None + moe_all_reduce = ( + self.routed_experts.all_reduce + if self._use_combined_all_reduce or self._reduce_shared_output + else None + ) def _routed_output(): # Latent down/up projections via the min-latency fused GEMM op: @@ -1444,6 +1455,8 @@ def _routed_output(): # a dense last dimension. routed_latent = self.routed_expert_norm(routed_latent.contiguous()) routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) + elif self._reduce_shared_output: + shared_out = moe_all_reduce(shared_out) return routed_out + shared_out diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index 98990b593df2..ad63a4bb1348 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -236,6 +236,9 @@ class MoEDeployment: # ``model_config.locality_domain_policy.enabled``. Whether the machine can # actually serve it is ``env.has_dep(MoEDep.LOCALITY_DOMAIN)``. locality_domain_requested: bool = False + # The model wrapper can consume a globally combined routed output without + # reducing that branch again under tensor parallelism. + handles_global_routed_output: bool = False @property def smart_router(self) -> bool: diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index b06322c4e9ef..cc146b39b6c6 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -460,15 +460,18 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: f"MegaMoECuteDsl requires num_slots ({d.num_slots}) " f"divisible by ep_size ({d.ep_size}).", ) - # The fused kernel returns a globally combined routed output. MiniMax-M3 - # accounts for that when composing routed and shared experts; the other - # current model wrappers would apply a second AllReduce under TEP. - if d.parallel_size > 1 and not d.use_dp and not is_minimax_swiglu_bias: + # A TEP wrapper must reduce only its local branches: the fused kernel + # has already combined the routed output across the expert group. + if ( + d.parallel_size > 1 + and not d.use_dp + and not is_minimax_swiglu_bias + and not d.handles_global_routed_output + ): return _reject( MoERejectReason.TOPOLOGY_UNSUPPORTED, - "MegaMoECuteDsl supports TEP only for bias-free MiniMax-style " - "SwigluBias; other model compositions would reduce its already-global " - "routed output twice.", + "MegaMoECuteDsl TEP requires a model composition that handles " + "globally combined routed output without reducing it twice.", ) # ADP wider than EP would need an outer allgather + reducescatter # wrapper that this backend does not have. diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index ecba78855d6b..85df76f94bc6 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -301,6 +301,7 @@ def build_moe_deployment( parallel_size=mapping.tp_size, cluster_size=mapping.moe_cluster_size, use_dp=mapping.enable_attention_dp, + handles_global_routed_output=model_config.moe_handles_global_routed_output, num_slots=num_slots, env=environment if environment is not None else collect_moe_environment(), # Registered balancer only; config alone does not enable EPLB. diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index 57ecf4ae12ea..a8b62ba75fa3 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -535,6 +535,62 @@ def test_kimi_k3_routed_config_preserves_explicit_backend(backend): assert routed_model_config.moe_backend == backend assert model_config.moe_backend == backend + assert routed_model_config.moe_handles_global_routed_output + assert not model_config.moe_handles_global_routed_output + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("global_routed_output", [False, True]) +def test_kimi_k3_tep_reduces_each_branch_once(monkeypatch, global_routed_output): + """Local and already-combined routed outputs must give the same model result.""" + shared_local = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) + shared_peer = torch.tensor([[4.0, 3.0, 2.0, 1.0]]) + routed_local = torch.tensor([[2.0, 3.0]]) + routed_peer = torch.tensor([[5.0, 7.0]]) + collective_inputs = [] + + def all_reduce(value): + collective_inputs.append(value.clone()) + peer = shared_peer if global_routed_output else torch.cat((shared_peer, routed_peer), -1) + return value + peer + + class RoutedExperts(torch.nn.Module): + def forward(self, *args, **kwargs): + return routed_local + routed_peer if global_routed_output else routed_local + + runtime = KimiK3MoERuntime.__new__(KimiK3MoERuntime) + torch.nn.Module.__init__(runtime) + runtime.hidden_size = 4 + runtime.moe_hidden_size = 2 + runtime._use_combined_all_reduce = not global_routed_output + runtime._reduce_shared_output = global_routed_output + runtime.gate = SimpleNamespace(compute_logits=lambda value: value) + runtime.routed_experts = RoutedExperts() + runtime.routed_experts.all_reduce = all_reduce + runtime.shared_experts = lambda value: shared_local + runtime.routed_expert_down_proj = torch.nn.Identity() + runtime.routed_expert_norm = torch.nn.RMSNorm(2, eps=0.1) + runtime.routed_expert_up_proj = torch.nn.Linear(2, 4, bias=False) + runtime.routed_expert_up_proj.weight.data.copy_( + torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, -1.0]]) + ) + runtime.moe_main_event = runtime.moe_shared_event = runtime.shared_expert_stream = None + monkeypatch.setattr(modeling_kimi_linear, "_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", True) + monkeypatch.setattr( + modeling_kimi_linear, + "maybe_execute_in_parallel", + lambda first, second, *args, **kwargs: (first(), second()), + ) + + result = runtime(torch.zeros(1, 4)) + expected = ( + shared_local + + shared_peer + + runtime.routed_expert_up_proj(runtime.routed_expert_norm(routed_local + routed_peer)) + ) + torch.testing.assert_close(result, expected) + assert len(collective_inputs) == 1 + assert collective_inputs[0].shape[-1] == (4 if global_routed_output else 6) @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index bd96fa844aa7..ee03f8f5583f 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -1182,6 +1182,34 @@ def test_megamoe_cutedsl_tep_requires_minimax_m3_activation() -> None: assert plain_swiglu_verdict.reject_reason is MoERejectReason.TOPOLOGY_UNSUPPORTED +@pytest.mark.cpu_only +@pytest.mark.parametrize("handles_global_output", [False, True]) +def test_megamoe_cutedsl_situ_tep_requires_global_output_composition(handles_global_output): + deployment = build_moe_deployment( + ModelConfig( + mapping=Mapping(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8), + moe_handles_global_routed_output=handles_global_output, + ), + num_experts=8, + environment=_make_megamoe_cutedsl_test_deployment().env, + ) + verdict = MegaMoECuteDsl.can_implement( + MoEProblem( + quant=QuantAlgo.NVFP4.value, + dtype_act=torch.bfloat16, + hidden_size=512, + intermediate_size=512, + num_experts=8, + top_k=2, + activation=ActivationType.SiTu.name, + ), + deployment, + ) + assert verdict.eligible is handles_global_output + if not handles_global_output: + assert verdict.reject_reason is MoERejectReason.TOPOLOGY_UNSUPPORTED + + @pytest.mark.cpu_only def test_megamoe_cutedsl_accepts_uniform_minimax_m3_swiglu_bias() -> None: moe = _make_megamoe_cutedsl_for_ctor_test( From d30c806ad37a5ef20b9991d767161255f9525365 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 17 Sep 2026 06:03:04 +0000 Subject: [PATCH 04/25] [None][refactor] Remove unrelated K3 MoE compatibility changes (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tensorrt_llm/_torch/model_config.py | 3 - .../_torch/models/modeling_kimi_linear.py | 21 ++----- .../_torch/moe/fused_moe/impl_contract.py | 3 - .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 17 +++--- .../_torch/moe/fused_moe/moe_resolution.py | 1 - .../_torch/moe/test_kimi_k3_situ_moe.py | 56 ------------------- tests/unittest/_torch/moe/test_moe_backend.py | 28 ---------- 7 files changed, 11 insertions(+), 118 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 540b471583d7..2c28757cfee5 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -251,9 +251,6 @@ class ModelConfig(Generic[TConfig]): # disables the optimization; hardware support is resolved by the backend. skip_correction_threshold: float = 0.0 moe_backend: str = 'CUTLASS' # options can be CUTLASS, TRTLLM - # Model composition reduces only local branches when a fused-communication - # backend returns an already globally combined routed output. - moe_handles_global_routed_output: bool = False # IF true, disables FC2+finalize fusion in CUTLASS MoE backend moe_disable_finalize_fusion: bool = False # If true, use low precision combine in MoE operations (only for NVFP4 quantization) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index c8f6fc279a04..38dad076d57f 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -54,7 +54,7 @@ column-sharded and down is row-sharded. Direct MoE-TP combines the shared hidden-width partial and routed latent partial into one all-reduce after the two streams join, then splits them before the routed norm/up projection. -Fused-communication routed paths reduce only the shared-expert partial, +Communication-backed routed paths keep the shared ``GatedMLP`` reduction, because their routed result is already combined. Under attention DP the shared experts stay replicated. ``lm_head`` uses the stock ``LMHead`` (vocab-sharded + gather), so logits are identical on all ranks. @@ -119,7 +119,6 @@ from ..modules.rms_norm import RMSNorm from ..modules.situ import SituAndMul from ..moe.fused_moe import ConfigurableMoE, SiTuActivation, TRTLLMGenFusedMoE, create_moe -from ..moe.fused_moe.interface import MoESchedulerKind from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..utils import AuxStreamType from .modeling_speculative import SpecDecOneEngineForCausalLM @@ -1167,8 +1166,7 @@ def __init__( # Under attention DP each rank owns different tokens, so the shared # expert is replicated (TP size 1) and must not reduce across ranks. # Direct MoE-TP leaves both branches as partials for one concatenated - # all-reduce. A fused-communication backend has already combined the - # routed branch, so only its shared-expert partial needs reduction. + # all-reduce. use_shared_tp = not attention_dp and model_config.mapping.tp_size > 1 routed_all_reduce = self.routed_experts.all_reduce if use_shared_tp and routed_all_reduce is None: @@ -1176,11 +1174,7 @@ def __init__( "Kimi K3 direct MoE tensor parallelism requires the " "ConfigurableMoE all-reduce even when reduce_results=False." ) - routed_output_is_global = ( - self.routed_experts.backend.scheduler_kind == MoESchedulerKind.FUSED_COMM - ) - self._use_combined_all_reduce = use_shared_tp and not routed_output_is_global - self._reduce_shared_output = use_shared_tp and routed_output_is_global + self._use_combined_all_reduce = use_shared_tp self.shared_experts = GatedMLP( hidden_size=cfg.hidden_size, intermediate_size=shared_intermediate, @@ -1369,7 +1363,6 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) routed_model_config.mapping = routed_mapping routed_model_config.moe_backend = model_config.moe_backend - routed_model_config.moe_handles_global_routed_output = True # MegaMoE uses this value as global DP SymmBuffer capacity, then # divides it by EP size for the per-rank allocation. Other backends # keep the user-configured value as their MoE chunking bound. @@ -1402,11 +1395,7 @@ def forward(self, hidden_states: torch.Tensor, all_rank_num_tokens=None) -> torc """``hidden_states``: ``[num_tokens, hidden_size]`` bf16.""" identity = hidden_states router_logits = self.gate.compute_logits(hidden_states) - moe_all_reduce = ( - self.routed_experts.all_reduce - if self._use_combined_all_reduce or self._reduce_shared_output - else None - ) + moe_all_reduce = self.routed_experts.all_reduce if self._use_combined_all_reduce else None def _routed_output(): # Latent down/up projections via the min-latency fused GEMM op: @@ -1455,8 +1444,6 @@ def _routed_output(): # a dense last dimension. routed_latent = self.routed_expert_norm(routed_latent.contiguous()) routed_out = self._routed_projection(routed_latent, self.routed_expert_up_proj) - elif self._reduce_shared_output: - shared_out = moe_all_reduce(shared_out) return routed_out + shared_out diff --git a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py index ad63a4bb1348..98990b593df2 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py +++ b/tensorrt_llm/_torch/moe/fused_moe/impl_contract.py @@ -236,9 +236,6 @@ class MoEDeployment: # ``model_config.locality_domain_policy.enabled``. Whether the machine can # actually serve it is ``env.has_dep(MoEDep.LOCALITY_DOMAIN)``. locality_domain_requested: bool = False - # The model wrapper can consume a globally combined routed output without - # reducing that branch again under tensor parallelism. - handles_global_routed_output: bool = False @property def smart_router(self) -> bool: diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index cc146b39b6c6..b06322c4e9ef 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -460,18 +460,15 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility: f"MegaMoECuteDsl requires num_slots ({d.num_slots}) " f"divisible by ep_size ({d.ep_size}).", ) - # A TEP wrapper must reduce only its local branches: the fused kernel - # has already combined the routed output across the expert group. - if ( - d.parallel_size > 1 - and not d.use_dp - and not is_minimax_swiglu_bias - and not d.handles_global_routed_output - ): + # The fused kernel returns a globally combined routed output. MiniMax-M3 + # accounts for that when composing routed and shared experts; the other + # current model wrappers would apply a second AllReduce under TEP. + if d.parallel_size > 1 and not d.use_dp and not is_minimax_swiglu_bias: return _reject( MoERejectReason.TOPOLOGY_UNSUPPORTED, - "MegaMoECuteDsl TEP requires a model composition that handles " - "globally combined routed output without reducing it twice.", + "MegaMoECuteDsl supports TEP only for bias-free MiniMax-style " + "SwigluBias; other model compositions would reduce its already-global " + "routed output twice.", ) # ADP wider than EP would need an outer allgather + reducescatter # wrapper that this backend does not have. diff --git a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py index 85df76f94bc6..ecba78855d6b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py +++ b/tensorrt_llm/_torch/moe/fused_moe/moe_resolution.py @@ -301,7 +301,6 @@ def build_moe_deployment( parallel_size=mapping.tp_size, cluster_size=mapping.moe_cluster_size, use_dp=mapping.enable_attention_dp, - handles_global_routed_output=model_config.moe_handles_global_routed_output, num_slots=num_slots, env=environment if environment is not None else collect_moe_environment(), # Registered balancer only; config alone does not enable EPLB. diff --git a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py index a8b62ba75fa3..57ecf4ae12ea 100644 --- a/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/moe/test_kimi_k3_situ_moe.py @@ -535,62 +535,6 @@ def test_kimi_k3_routed_config_preserves_explicit_backend(backend): assert routed_model_config.moe_backend == backend assert model_config.moe_backend == backend - assert routed_model_config.moe_handles_global_routed_output - assert not model_config.moe_handles_global_routed_output - - -@pytest.mark.cpu_only -@pytest.mark.parametrize("global_routed_output", [False, True]) -def test_kimi_k3_tep_reduces_each_branch_once(monkeypatch, global_routed_output): - """Local and already-combined routed outputs must give the same model result.""" - shared_local = torch.tensor([[1.0, 2.0, 3.0, 4.0]]) - shared_peer = torch.tensor([[4.0, 3.0, 2.0, 1.0]]) - routed_local = torch.tensor([[2.0, 3.0]]) - routed_peer = torch.tensor([[5.0, 7.0]]) - collective_inputs = [] - - def all_reduce(value): - collective_inputs.append(value.clone()) - peer = shared_peer if global_routed_output else torch.cat((shared_peer, routed_peer), -1) - return value + peer - - class RoutedExperts(torch.nn.Module): - def forward(self, *args, **kwargs): - return routed_local + routed_peer if global_routed_output else routed_local - - runtime = KimiK3MoERuntime.__new__(KimiK3MoERuntime) - torch.nn.Module.__init__(runtime) - runtime.hidden_size = 4 - runtime.moe_hidden_size = 2 - runtime._use_combined_all_reduce = not global_routed_output - runtime._reduce_shared_output = global_routed_output - runtime.gate = SimpleNamespace(compute_logits=lambda value: value) - runtime.routed_experts = RoutedExperts() - runtime.routed_experts.all_reduce = all_reduce - runtime.shared_experts = lambda value: shared_local - runtime.routed_expert_down_proj = torch.nn.Identity() - runtime.routed_expert_norm = torch.nn.RMSNorm(2, eps=0.1) - runtime.routed_expert_up_proj = torch.nn.Linear(2, 4, bias=False) - runtime.routed_expert_up_proj.weight.data.copy_( - torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [2.0, -1.0]]) - ) - runtime.moe_main_event = runtime.moe_shared_event = runtime.shared_expert_stream = None - monkeypatch.setattr(modeling_kimi_linear, "_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", True) - monkeypatch.setattr( - modeling_kimi_linear, - "maybe_execute_in_parallel", - lambda first, second, *args, **kwargs: (first(), second()), - ) - - result = runtime(torch.zeros(1, 4)) - expected = ( - shared_local - + shared_peer - + runtime.routed_expert_up_proj(runtime.routed_expert_norm(routed_local + routed_peer)) - ) - torch.testing.assert_close(result, expected) - assert len(collective_inputs) == 1 - assert collective_inputs[0].shape[-1] == (4 if global_routed_output else 6) @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index ee03f8f5583f..bd96fa844aa7 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -1182,34 +1182,6 @@ def test_megamoe_cutedsl_tep_requires_minimax_m3_activation() -> None: assert plain_swiglu_verdict.reject_reason is MoERejectReason.TOPOLOGY_UNSUPPORTED -@pytest.mark.cpu_only -@pytest.mark.parametrize("handles_global_output", [False, True]) -def test_megamoe_cutedsl_situ_tep_requires_global_output_composition(handles_global_output): - deployment = build_moe_deployment( - ModelConfig( - mapping=Mapping(world_size=8, tp_size=8, moe_tp_size=1, moe_ep_size=8), - moe_handles_global_routed_output=handles_global_output, - ), - num_experts=8, - environment=_make_megamoe_cutedsl_test_deployment().env, - ) - verdict = MegaMoECuteDsl.can_implement( - MoEProblem( - quant=QuantAlgo.NVFP4.value, - dtype_act=torch.bfloat16, - hidden_size=512, - intermediate_size=512, - num_experts=8, - top_k=2, - activation=ActivationType.SiTu.name, - ), - deployment, - ) - assert verdict.eligible is handles_global_output - if not handles_global_output: - assert verdict.reject_reason is MoERejectReason.TOPOLOGY_UNSUPPORTED - - @pytest.mark.cpu_only def test_megamoe_cutedsl_accepts_uniform_minimax_m3_swiglu_bias() -> None: moe = _make_megamoe_cutedsl_for_ctor_test( From ac515d3db959cccc4a17d34174cd4f7ae2ad6d17 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 17 Sep 2026 06:41:09 +0000 Subject: [PATCH 05/25] [None][refactor] Delegate hybrid manager state through a shared contract (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/modules/fla/cache_manager.py | 52 +------- .../_torch/modules/kimi_kda/cache_manager.py | 84 +++---------- .../_torch/modules/mamba/cache_manager.py | 62 ++-------- .../_torch/modules/qwen4_exp/cache_manager.py | 51 ++------ .../kv_cache/mamba_cache_manager/common.py | 38 +++++- .../mamba_cache_manager_v2.py | 47 +++++++- .../kv_cache/test_mamba_cache_manager.py | 13 +- .../kv_cache/test_mamba_model_managers.py | 112 ++++++++++++++---- 8 files changed, 212 insertions(+), 247 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index f7b82a6f5923..8932e27b49b3 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -21,7 +21,6 @@ MambaAcceptanceBatch, MambaLayerCache, MambaStateLayout, - ReplayStateUpdateMetadata, _mamba_effective_tp_size, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( @@ -269,60 +268,17 @@ def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> None: - self._speculative_state = create_gdn_state(self, self._requested_replay) + def _initialize_model_state(self) -> IntermediateState | GDNReplayState: + state = create_gdn_state(self, self._requested_replay) - validate_gdn_layout(self, self._speculative_state) - - @override - def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: - state = self._speculative_state - return state.get_replay_metadata() if isinstance(state, GDNReplayState) else None - - @property - def intermediate_state_indices(self) -> torch.Tensor | None: - return self._speculative_state.intermediate_indices - - @property - def intermediate_ssm_states(self) -> torch.Tensor | None: - state = self._speculative_state - return None if isinstance(state, GDNReplayState) else state.intermediate_ssm - - @property - def intermediate_conv_states(self) -> torch.Tensor | None: - return self._speculative_state.intermediate_conv + validate_gdn_layout(self, state) + return state @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: state = self._speculative_state return isinstance(state, GDNReplayState) and state.has_bound_states - @override - def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: - if self.spec_config is None: - return super().mamba_layer_cache(layer_idx) - return self._speculative_state.make_layer_cache( - self.mamba_layer_offsets[layer_idx], - self.get_conv_states(layer_idx), - self.get_ssm_states(layer_idx), - ) - - @override - def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - self._speculative_state.reset_slots(slots, host_slots) - - @override - def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: - self._speculative_state.update(batch) - - @override - def _setup_model_state(self) -> None: - self._speculative_state.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) - - @override - def _shutdown_model_state(self) -> None: - self._speculative_state.shutdown() - def get_gdn_cache_params(config, *, spec_config=None, quant_config=None): """Shared Qwen3.5/Qwen4 GDN geometry; attention and PLE stay model-owned.""" diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index 6d252641a606..d8c44c48ee43 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -24,6 +24,7 @@ MambaAcceptanceBatch, MambaLayerCache, MambaStateLayout, + ReplayStateUpdateMetadata, SpeculativeMambaLayerCache, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( @@ -135,6 +136,17 @@ def allocate_conv_cache() -> torch.Tensor: class KDAReplayState: """Own all V2 state used by the KDA fused multi-token verifier.""" + @property + def intermediate_ssm(self) -> torch.Tensor | None: + return None + + @property + def intermediate_conv(self) -> torch.Tensor | None: + return None + + def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: + return None + def __init__( self, num_speculative_tokens: int, @@ -461,17 +473,15 @@ def __init__( super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> None: + def _initialize_model_state(self) -> IntermediateState | KDAReplayState: num_spec = self._requested_num_spec if num_spec is None: num_spec = get_kda_replay_num_spec(self.spec_config, manager_supports_replay=True) if num_spec is not None: self._kda_replay = KDAReplayState(num_spec) - else: - self._speculative_state = IntermediateState() - - if self._kda_replay is not None: self._kda_replay.validate(self._state_layout) + return self._kda_replay + return IntermediateState() @property def use_kda_replay_update(self) -> bool: @@ -486,56 +496,6 @@ def _extra_scratch_bytes_per_slot(self) -> int: for layer_id in self.mamba_pp_layers ) - @override - def _setup_model_state(self) -> None: - if self._kda_replay is None: - self._speculative_state.bind( - self._state_layout, self.all_ssm_states, self.all_conv_states - ) - else: - self._kda_replay.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) - - @property - def intermediate_state_indices(self) -> torch.Tensor | None: - if self._kda_replay is not None: - return self._kda_replay.intermediate_indices - return self._speculative_state.intermediate_indices - - @property - def intermediate_ssm_states(self) -> torch.Tensor | None: - if self._kda_replay is not None: - return None - return self._speculative_state.intermediate_ssm - - @property - def intermediate_conv_states(self) -> torch.Tensor | None: - if self._kda_replay is not None: - return None - return self._speculative_state.intermediate_conv - - @override - def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: - if self._kda_replay is None: - if self.spec_config is None: - return super().mamba_layer_cache(layer_idx) - return self._speculative_state.make_layer_cache( - self.mamba_layer_offsets[layer_idx], - self.get_conv_states(layer_idx), - self.get_ssm_states(layer_idx), - ) - return self._kda_replay.make_layer_cache( - self.mamba_layer_offsets[layer_idx], - self.get_conv_states(layer_idx), - self.get_ssm_states(layer_idx), - ) - - @override - def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - if self._kda_replay is None: - self._speculative_state.reset_slots(slots, host_slots) - else: - self._kda_replay.reset_slots(slots, host_slots) - @override def _on_state_slots_relocated(self, old_slots: list[int], new_slots: list[int]) -> None: if self._kda_replay is None: @@ -548,13 +508,6 @@ def _on_state_slots_relocated(self, old_slots: list[int], new_slots: list[int]) ) self._kda_replay.reset_slots(slots, fresh_slots) - @override - def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: - if self._kda_replay is None: - self._speculative_state.update(batch) - else: - self._kda_replay.update(batch) - @override def update_resources( self, scheduled_batch, attn_metadata=None, kv_cache_dtype_byte_size=None @@ -630,13 +583,6 @@ def seed_kda_replay_caches_for_disagg_gen(self, request_ids: list[int]) -> None: """Compatibility entry point; the executor uses the generic transfer hook.""" self.on_state_transfer_complete(request_ids) - @override - def _shutdown_model_state(self) -> None: - if self._kda_replay is not None: - self._kda_replay.shutdown() - else: - self._speculative_state.shutdown() - def get_kimi_cache_params(config, *, spec_config=None, quant_config=None): """KDA uses equal Q/K/V sections and FP32 recurrent state alongside MLA.""" diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py index cbcc9ca66018..1500c020041f 100644 --- a/tensorrt_llm/_torch/modules/mamba/cache_manager.py +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -191,6 +191,10 @@ def shutdown(self) -> None: class ReplayHistory: """Own compact Mamba2 replay history, seeds, and bookkeeping.""" + @property + def intermediate_ssm(self) -> torch.Tensor | None: + return None + def __init__(self, tokens_per_gen_step: int) -> None: self._ssm_states: tuple[torch.Tensor, ...] = () self._conv_states: tuple[torch.Tensor, ...] = () @@ -398,6 +402,8 @@ def select_mamba2_state( class NemotronHybridCacheManagerV2(MambaHybridCacheManagerV2): """Nemotron/Mamba2 geometry, replay selection and seed ownership.""" + _speculative_state: Mamba2State | ReplayHistory + @override def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) -> None: self._requested_replay = use_replay_state_update @@ -405,19 +411,19 @@ def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> None: + def _initialize_model_state(self) -> Mamba2State | ReplayHistory: if self._requested_replay is not None: - self._speculative_state = create_mamba2_state( + state = create_mamba2_state( spec_config=self.spec_config, use_replay=self._requested_replay ) else: - self._speculative_state = select_mamba2_state( + state = select_mamba2_state( spec_config=self.spec_config, ssm_cache_dtype=self.ssm_state_dtype, stochastic_rounding=self._mamba_ssm_stochastic_rounding, ) - if isinstance(self._speculative_state, ReplayHistory): + if isinstance(state, ReplayHistory): from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( _mamba_effective_tp_size, ) @@ -426,56 +432,12 @@ def _initialize_model_state(self) -> None: self._state_layout.mapping ): raise ValueError("Replay state groups must be divisible by the effective TP size") - self._speculative_state.validate(self._state_layout) + state.validate(self._state_layout) + return state - @override - def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: - state = self._speculative_state - return state.get_replay_metadata() if isinstance(state, ReplayHistory) else None - - @override def get_mamba_ssm_rand_seed(self) -> torch.Tensor | None: return self._speculative_state.rand_seed - @property - def intermediate_state_indices(self) -> torch.Tensor | None: - return self._speculative_state.intermediate_indices - - @property - def intermediate_ssm_states(self) -> torch.Tensor | None: - state = self._speculative_state - return None if isinstance(state, ReplayHistory) else state.intermediate_ssm - - @property - def intermediate_conv_states(self) -> torch.Tensor | None: - return self._speculative_state.intermediate_conv - - @override - def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: - if self.spec_config is None: - return super().mamba_layer_cache(layer_idx) - return self._speculative_state.make_layer_cache( - self.mamba_layer_offsets[layer_idx], - self.get_conv_states(layer_idx), - self.get_ssm_states(layer_idx), - ) - - @override - def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - self._speculative_state.reset_slots(slots, host_slots) - - @override - def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: - self._speculative_state.update(batch) - - @override - def _setup_model_state(self) -> None: - self._speculative_state.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) - - @override - def _shutdown_model_state(self) -> None: - self._speculative_state.shutdown() - def get_nemotron_cache_params(config, *, spec_config=None, quant_config=None): """Derive Mamba2 state geometry from the Nemotron hybrid pattern.""" diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index 13495601eb9e..29f1f5681eda 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -33,9 +33,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( - MambaAcceptanceBatch, - MambaLayerCache, - ReplayStateUpdateMetadata, + IntermediateState, _get_local_mamba_cache_layout, ) from tensorrt_llm._torch.pyexecutor.resource_manager import get_pp_layers @@ -435,10 +433,10 @@ def _ple_buffer_configs(self) -> list[BufferConfig]: ] @override - def _initialize_model_state(self) -> None: - self._speculative_state = create_gdn_state(self, self._requested_replay) + def _initialize_model_state(self) -> IntermediateState | GDNReplayState: + state = create_gdn_state(self, self._requested_replay) - validate_gdn_layout(self, self._speculative_state) + validate_gdn_layout(self, state) params = self._ple_params if self._is_ple_draft: params = None @@ -451,6 +449,7 @@ def _initialize_model_state(self) -> None: self._init_qwen4_exp_ple_geometry( params, len(self._mamba_layer_mask), self.conv_state_dtype ) + return state @override def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConfig]: @@ -461,15 +460,10 @@ def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConf @override def _setup_model_state(self) -> None: - self._speculative_state.bind(self._state_layout, self.all_ssm_states, self.all_conv_states) + super()._setup_model_state() if self.local_num_mamba_layers: self._setup_ple_states(self._state_layout.slot_capacity) - @override - def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: - state = self._speculative_state - return state.get_replay_metadata() if isinstance(state, GDNReplayState) else None - @classmethod @override def _extra_state_bytes_per_rank(cls, model_config, mapping, **kwargs) -> int: @@ -494,44 +488,13 @@ def _extra_state_bytes_per_rank(cls, model_config, mapping, **kwargs) -> int: def _shutdown_model_state(self) -> None: self._ple_conv_states.clear() self._ple_ngram_contexts.clear() - self._speculative_state.shutdown() - - @property - def intermediate_state_indices(self) -> torch.Tensor | None: - return self._speculative_state.intermediate_indices - - @property - def intermediate_ssm_states(self) -> torch.Tensor | None: - state = self._speculative_state - return None if isinstance(state, GDNReplayState) else state.intermediate_ssm - - @property - def intermediate_conv_states(self) -> torch.Tensor | None: - return self._speculative_state.intermediate_conv + super()._shutdown_model_state() @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: state = self._speculative_state return isinstance(state, GDNReplayState) and state.has_bound_states - @override - def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: - if self.spec_config is None: - return super().mamba_layer_cache(layer_idx) - return self._speculative_state.make_layer_cache( - self.mamba_layer_offsets[layer_idx], - self.get_conv_states(layer_idx), - self.get_ssm_states(layer_idx), - ) - - @override - def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - self._speculative_state.reset_slots(slots, host_slots) - - @override - def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: - self._speculative_state.update(batch) - def get_qwen4_exp_ple_cache_params(config, *, total_layers: int, is_draft: bool): """Align target-only PLE state with a target/draft cache layout.""" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py index eece2102d4c3..0df7259d4b4a 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -131,6 +131,38 @@ class MambaAcceptanceBatch: is_dummy_request: torch.Tensor | None +class MambaState(Protocol): + """Model-selected recurrent state; persistent views remain pool-owned.""" + + @property + def intermediate_indices(self) -> torch.Tensor | None: ... + + @property + def intermediate_ssm(self) -> torch.Tensor | None: ... + + @property + def intermediate_conv(self) -> torch.Tensor | None: ... + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: ... + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: ... + + def update(self, batch: MambaAcceptanceBatch) -> None: ... + + def shutdown(self) -> None: ... + + def make_layer_cache( + self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor + ) -> MambaLayerCache: ... + + def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: ... + + class BaseMambaCacheManager(ABC): """Abstract interface for accessing Mamba/recurrent state caches.""" @@ -337,9 +369,6 @@ def get_ssm_states(self, layer_idx: int) -> torch.Tensor: def get_conv_states(self, layer_idx: int) -> torch.Tensor: return self.all_conv_states[self.mamba_layer_offsets[layer_idx]] - def get_mamba_ssm_rand_seed(self) -> torch.Tensor | None: - return None - def _get_mamba_hybrid_pool_size(max_batch_size: int, mapping: Mapping) -> int: """Return the internal Mamba state pool size for MixedMambaHybridCacheManager.""" @@ -831,6 +860,9 @@ def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | Non def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: """Plain intermediate buffers have no request-persistent bookkeeping.""" + def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: + return None + def update(self, batch: MambaAcceptanceBatch) -> None: if self.intermediate_indices is None: return diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 4702734f29e2..eace94805e03 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -77,7 +77,9 @@ MambaHybridCacheManager, MambaLayerCache, MambaRole, + MambaState, MambaStateLayout, + ReplayStateUpdateMetadata, _estimate_mamba_hybrid_cache_cost, _get_num_cuda_graph_padding_dummy_slots, _mamba_effective_tp_size, @@ -269,7 +271,7 @@ def __init__( conv_section_dims=tuple(self.conv_section_dims), conv_state_layout=conv_state_layout, ) - self._initialize_model_state() + self._speculative_state: MambaState | None = self._initialize_model_state() if isinstance(num_kv_heads, int): per_layer_kv_heads = [num_kv_heads] * total_layers @@ -316,12 +318,13 @@ def __init__( self._setup_state_pool() - def _initialize_model_state(self) -> None: - """Initialize and validate model-owned state before constructing the pool.""" + def _initialize_model_state(self) -> MambaState | None: + """Return validated model-owned state before constructing the pool.""" if self.spec_config is not None: raise NotImplementedError( "Speculative decoding requires a model-specific Mamba cache manager" ) + return None @classmethod @override @@ -710,9 +713,39 @@ def _setup_states(self) -> None: def _setup_model_state(self) -> None: """Bind views only after layout, pool and slot capacity are final.""" + if self._speculative_state is not None: + self._speculative_state.bind( + self._state_layout, self.all_ssm_states, self.all_conv_states + ) + + @property + def intermediate_state_indices(self) -> torch.Tensor | None: + state = self._speculative_state + return state.intermediate_indices if state is not None else None + + @property + def intermediate_ssm_states(self) -> torch.Tensor | None: + state = self._speculative_state + return state.intermediate_ssm if state is not None else None + + @property + def intermediate_conv_states(self) -> torch.Tensor | None: + state = self._speculative_state + return state.intermediate_conv if state is not None else None + + @override + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + state = self._speculative_state + return state.get_replay_metadata() if state is not None else None @override def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: + if self._speculative_state is not None and self.spec_config is not None: + return self._speculative_state.make_layer_cache( + self.mamba_layer_offsets[layer_idx], + self.get_conv_states(layer_idx), + self.get_ssm_states(layer_idx), + ) return MambaLayerCache( conv=self.get_conv_states(layer_idx), temporal=self.get_ssm_states(layer_idx) ) @@ -922,6 +955,8 @@ def _reset_context_mamba_slots(self, num_contexts: int) -> None: def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: """Reset any model-owned request state for the supplied slots.""" + if self._speculative_state is not None: + self._speculative_state.reset_slots(slots, host_slots) def _refresh_dummy_request_mask(self, is_dummy: list[bool]) -> None: count = len(is_dummy) @@ -1012,7 +1047,9 @@ def update_mamba_states( ) def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: - raise NotImplementedError("The model manager must implement speculative state updates") + if self._speculative_state is None: + raise NotImplementedError("The model manager must provide speculative state") + self._speculative_state.update(batch) @override def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: @@ -1309,6 +1346,8 @@ def shutdown(self): def _shutdown_model_state(self) -> None: """Release model-owned state before the common pool is destroyed.""" + if self._speculative_state is not None: + self._speculative_state.shutdown() __all__ = ["MambaHybridCacheManagerV2"] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 9e50610014d7..088334d9c0bf 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -354,7 +354,7 @@ def test_kimi_explicit_v2_manager_enables_owned_kda_replay_policy( mgr.spec_config = spec_config mgr._requested_num_spec = None mgr._state_layout = SimpleNamespace(spec_config=spec_config, mamba_pp_layers=()) - mgr._initialize_model_state() + mgr._speculative_state = mgr._initialize_model_state() policy = mgr._kda_replay assert isinstance(policy, KDAReplayState) assert policy.num_speculative_tokens == spec_config.tokens_per_gen_step - 1 @@ -569,7 +569,7 @@ def __init__(self, *args, **kwargs): mgr._state_layout = SimpleNamespace( spec_config=mgr.spec_config, mamba_pp_layers=(), n_groups_per_rank=1 ) - mgr._initialize_model_state() + mgr._speculative_state = mgr._initialize_model_state() assert isinstance(mgr._speculative_state, GDNReplayState) assert captured_v2["max_num_tokens"] == 256 assert "model_type" not in captured_v2 @@ -2294,15 +2294,15 @@ def _build_v2_hybrid_with_mamba_layer( class ModelStateManager(model_cls): def _initialize_model_state(self): if speculative_state is None: - super()._initialize_model_state() + return super()._initialize_model_state() elif isinstance(speculative_state, KDAReplayState): self._kda_replay = speculative_state self._kda_replay.validate(self._state_layout) else: from tensorrt_llm._torch.modules.fla.cache_manager import validate_gdn_layout - self._speculative_state = speculative_state - validate_gdn_layout(self, self._speculative_state) + validate_gdn_layout(self, speculative_state) + return speculative_state manager_cls = ModelStateManager return manager_cls( @@ -3620,6 +3620,7 @@ def test_v2_kda_replay_policy_records_acceptance_and_skips_dummy_rows(): policy.intermediate_indices = torch.arange(3, dtype=torch.int32) mgr._dummy_request_mask = torch.tensor([False, False, True]) mgr._kda_replay = policy + mgr._speculative_state = policy mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) if not hasattr(mgr, "_dummy_request_mask"): @@ -3647,6 +3648,7 @@ def test_v2_kda_replay_policy_resets_context_slots(): ) policy.prev_num_accepted_tokens = torch.tensor([5, 6, 7], dtype=torch.int32) mgr._kda_replay = policy + mgr._speculative_state = policy mgr.cuda_state_indices = torch.tensor([2, 1], dtype=torch.int32) mgr._host_state_indices = torch.tensor([2, 1], dtype=torch.int32) @@ -3667,6 +3669,7 @@ def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): ) policy.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) mgr._kda_replay = policy + mgr._speculative_state = policy base_update = MagicMock() monkeypatch.setattr(KVCacheManagerV2, "update_resources", base_update) active = SimpleNamespace( diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index eb6887b626aa..58e1ea777775 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -188,7 +188,7 @@ def test_qwen_managers_share_gdn_algorithm_not_model_inheritance(): manager._requested_replay = True manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) _prepare_model_layout(manager) - manager._initialize_model_state() + manager._speculative_state = manager._initialize_model_state() states.append(manager._speculative_state) assert all(type(state) is GDNReplayState for state in states) assert states[0] is not states[1] @@ -258,12 +258,12 @@ def test_k3_initializes_only_the_selected_state( selector, ) - assert manager._initialize_model_state() is None + manager._speculative_state = manager._initialize_model_state() num_spec = requested_num_spec if requested_num_spec is not None else selected_num_spec if num_spec is not None: intermediate_factory.assert_not_called() replay_factory.assert_called_once_with(num_spec) - assert not hasattr(manager, "_speculative_state") + assert manager._speculative_state is manager._kda_replay state = manager._kda_replay else: replay_factory.assert_not_called() @@ -298,28 +298,90 @@ def test_k3_initializes_only_the_selected_state( manager._shutdown_model_state() assert shutdown.call_count == 2 if num_spec is not None: - assert not hasattr(manager, "_speculative_state") - + assert manager._speculative_state is manager._kda_replay -def test_base_manager_does_not_own_speculative_state(): - forbidden = { - "_speculative_state", - "IntermediateState", - "intermediate_ssm_states", - "intermediate_conv_states", - } - for attribute in vars(MambaHybridCacheManagerV2).values(): - function = attribute.__func__ if isinstance(attribute, classmethod) else attribute - if isinstance(function, FunctionType): - assert not forbidden.intersection(function.__code__.co_names), function.__name__ +def test_base_manager_without_model_state(): manager = object.__new__(MambaHybridCacheManagerV2) manager.spec_config = None - manager._initialize_model_state() + manager._speculative_state = manager._initialize_model_state() + assert manager._speculative_state is None manager._setup_model_state() manager._reset_model_slots(object(), [0]) manager._shutdown_model_state() - assert not forbidden.intersection(vars(manager)) + assert manager.intermediate_state_indices is None + assert manager.intermediate_ssm_states is None + assert manager.intermediate_conv_states is None + assert manager.get_replay_state_update_metadata() is None + + +@pytest.mark.parametrize( + "manager_cls", + [ + MambaHybridCacheManager, + MambaHybridCacheManagerV2, + Qwen35HybridCacheManagerV2, + Qwen4ExpHybridCacheManagerV2, + KimiK3HybridCacheManagerV2, + ], +) +def test_seed_accessor_is_not_a_common_manager_capability(manager_cls): + assert not hasattr(manager_cls, "get_mamba_ssm_rand_seed") + + +def test_mamba2_seed_lifecycle_without_speculative_decoding(): + manager = object.__new__(NemotronHybridCacheManagerV2) + manager.spec_config = None + manager._requested_replay = False + manager._state_layout = replace(_layout(), spec_config=None, stochastic_rounding=True) + views = _state_views() + manager.all_ssm_states = views.all_ssm_states + manager.all_conv_states = views.all_conv_states + manager._speculative_state = manager._initialize_model_state() + manager._setup_model_state() + seeds = manager.get_mamba_ssm_rand_seed() + before = seeds.clone() + manager._reset_model_slots(torch.tensor([1]), [1]) + assert manager.get_mamba_ssm_rand_seed() is seeds + assert seeds[1] != before[1] + torch.testing.assert_close(seeds[[0, 2, 3, 4]], before[[0, 2, 3, 4]]) + assert manager.intermediate_ssm_states is None + manager._shutdown_model_state() + assert manager.get_mamba_ssm_rand_seed() is None + + +@pytest.mark.parametrize( + "state", [IntermediateState(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)] +) +def test_common_manager_consumes_model_state_contract(state): + manager = object.__new__(MambaHybridCacheManagerV2) + manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) + manager._speculative_state = state + manager._state_layout = _layout() + views = _state_views() + manager.all_ssm_states = views.all_ssm_states + manager.all_conv_states = views.all_conv_states + manager.mamba_layer_offsets = {42: 0} + manager._setup_model_state() + assert manager.intermediate_state_indices is state.intermediate_indices + assert manager.intermediate_ssm_states is state.intermediate_ssm + assert manager.intermediate_conv_states is state.intermediate_conv + payload = manager.mamba_layer_cache(42) + assert payload.conv is manager.all_conv_states[0] + assert payload.temporal is manager.all_ssm_states[0] + if isinstance(state, KDAReplayState): + assert payload.kda_qkg_cache is not None + elif isinstance(state, ReplayHistory): + assert payload.old_x is not None + else: + assert payload.intermediate_ssm is not None + assert (manager.get_replay_state_update_metadata() is not None) == isinstance( + state, ReplayHistory + ) + manager._shutdown_model_state() + assert manager.intermediate_state_indices is None + assert manager.intermediate_ssm_states is None + assert manager.intermediate_conv_states is None def test_base_manager_requires_model_for_speculative_decoding(): @@ -327,6 +389,7 @@ def test_base_manager_requires_model_for_speculative_decoding(): manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) with pytest.raises(NotImplementedError, match="model-specific"): manager._initialize_model_state() + manager._speculative_state = None with pytest.raises(NotImplementedError, match="model manager"): manager._update_speculative_state(object()) @@ -341,7 +404,7 @@ def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) if mode != "plain" else None manager._requested_replay = mode == "replay" _prepare_model_layout(manager) - manager._initialize_model_state() + manager._speculative_state = manager._initialize_model_state() state = manager._speculative_state assert isinstance(state, IntermediateState) == (mode != "replay") assert isinstance(state, ReplayHistory) == (mode == "replay") @@ -388,7 +451,7 @@ def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): shutdown.assert_called_once_with() assert manager.intermediate_ssm_states is None and state.intermediate_conv is None if mode == "replay": - assert not hasattr(state, "intermediate_ssm") + assert state.intermediate_ssm is None def test_ple_adds_buffers_without_overriding_snapshot_or_budget_solver(): @@ -608,8 +671,8 @@ def test_replay_and_intermediate_are_independent_algorithms(): assert Mamba2State.__bases__ == (IntermediateState,) assert not issubclass(ReplayHistory, IntermediateState) assert not issubclass(GDNReplayState, Mamba2State) - assert not hasattr(ReplayHistory(3), "intermediate_ssm") - assert not hasattr(GDNReplayState(3), "intermediate_ssm") + assert ReplayHistory(3).intermediate_ssm is None + assert GDNReplayState(3).intermediate_ssm is None state = IntermediateState() assert not {"rand_seed", "_seed_request_counter", "_seed_rank_offset"}.intersection(vars(state)) assert not hasattr(state, "_promote_conv") @@ -697,7 +760,7 @@ def promote(destination, source, source_indices, positions, destination_indices) assert (cache.intermediate_ssm is None) == replay if replay: assert state.prev_num_accepted_tokens.tolist() == [0, 0, 3, 0, 1] - assert not hasattr(state, "intermediate_ssm") + assert state.intermediate_ssm is None state.shutdown() @@ -1005,8 +1068,9 @@ def create_pool(self, *args, **kwargs): class ObservedQwen4(Qwen4ExpHybridCacheManagerV2): def _initialize_model_state(self): - super()._initialize_model_state() + state = super()._initialize_model_state() events.append("validate") + return state def _get_state_buffer(self, local_layer_idx, role, dtype, state_shape): assert events[:2] == ["validate", "pool"] From b5d7747cae27cf29bc42a1623b0c68b6d2424941 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Sun, 20 Sep 2026 06:55:40 +0000 Subject: [PATCH 06/25] [None][refactor] Address hybrid cache manager review (By Agent) Move Qwen4 pool views into the model manager and inline its attention mapping. Remove the QSA compatibility alias and redundant legacy warnings. Resolve legacy replay policy during construction and handle empty recurrent layouts safely. Validation: 384 targeted tests passed on B300; Qwen4 FP8 ADP4 MTP3 prefix-cache integration passed before the review follow-ups. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../attention/backends/sparse/qsa/__init__.py | 1 - .../backends/sparse/qsa/cache_manager.py | 11 -- .../attention/backends/sparse/qsa/metadata.py | 7 +- .../attention/backends/sparse/registry.py | 5 +- .../_torch/modules/mamba/cache_manager.py | 2 + .../_torch/modules/qwen4_exp/cache_manager.py | 81 ++++++++- tensorrt_llm/_torch/pyexecutor/_util.py | 67 ++------ .../kv_cache/kv_cache_manager_v2.py | 69 -------- .../kv_cache/mamba_cache_manager/legacy.py | 67 +++++++- .../attention/sparse/qsa/test_qsa_sparse.py | 9 +- .../kv_cache/test_mamba_cache_manager.py | 25 ++- .../kv_cache/test_mamba_model_managers.py | 160 ++++++++++++++---- .../modeling/test_qsa_runtime_wiring.py | 13 +- 13 files changed, 322 insertions(+), 195 deletions(-) delete mode 100644 tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py index 443edf7415df..bb11190acf00 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py @@ -7,7 +7,6 @@ _EXPORTS = { "QSATrtllmAttention": "backend", - "QSAMambaHybridCacheManagerV2": "cache_manager", "QSAAttentionMetadata": "metadata", "QSASparseMetadataParams": "params", "QSASparseParams": "params", diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py deleted file mode 100644 index 5e8bf2e1e8e7..000000000000 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ /dev/null @@ -1,11 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Compatibility imports for the Qwen4-owned QSA cache manager.""" - -from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import QSA_INDEX_POSITION -from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( - Qwen4ExpHybridCacheManagerV2 as QSAMambaHybridCacheManagerV2, -) - -__all__ = ["QSA_INDEX_POSITION", "QSAMambaHybridCacheManagerV2"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index 43e9b5318662..35003ed163e4 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -9,7 +9,6 @@ from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata -from .cache_manager import QSAMambaHybridCacheManagerV2 from .constants import QSA_KEY_ROLE_INDEX from .kernels import triton_qsa_unscale_block_table from .params import QSASparseMetadataParams @@ -45,9 +44,11 @@ def __init__(self, *args: object, **kwargs: object) -> None: super().__init__(*args, **kwargs) def __post_init__(self) -> None: + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 + super().__post_init__() - if not isinstance(self.kv_cache_manager, QSAMambaHybridCacheManagerV2): - raise TypeError("QSA sparse attention requires QSAMambaHybridCacheManagerV2") + if not isinstance(self.kv_cache_manager, Qwen4ExpHybridCacheManagerV2): + raise TypeError("QSA sparse attention requires Qwen4ExpHybridCacheManagerV2") if not isinstance(self.sparse_metadata_params, QSASparseMetadataParams): raise TypeError("QSA sparse attention requires QSASparseMetadataParams") self.qsa_has_local_layers = self.kv_cache_manager.qsa_position_layer_id is not None diff --git a/tensorrt_llm/_torch/attention/backends/sparse/registry.py b/tensorrt_llm/_torch/attention/backends/sparse/registry.py index 5b0520e77392..2ced059aa75d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/registry.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/registry.py @@ -27,7 +27,6 @@ def get_sparse_attn_kv_cache_manager( from .deepseek_v4 import DeepseekV4CacheManager from .dsa import DSACacheManager, DSACacheManagerV2 from .minimax_m3 import MiniMaxM3KVCacheManagerV2 - from .qsa import QSAMambaHybridCacheManagerV2 from .rocket import RocketKVCacheManager if sparse_attention_config.algorithm == "rocket": @@ -43,7 +42,9 @@ def get_sparse_attn_kv_cache_manager( elif sparse_attention_config.algorithm == "qsa": if not use_kv_cache_manager_v2: raise ValueError("QSA sparse attention requires KV cache manager V2") - return QSAMambaHybridCacheManagerV2 + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 + + return Qwen4ExpHybridCacheManagerV2 else: raise ValueError( f"Unsupported sparse attention algorithm: {sparse_attention_config.algorithm}" diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py index 1500c020041f..f245ad4edc2a 100644 --- a/tensorrt_llm/_torch/modules/mamba/cache_manager.py +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -301,6 +301,8 @@ def make_layer_cache( ) def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + if self.prev_num_accepted_tokens is None: + return self.prev_num_accepted_tokens.index_fill_(0, slots, 0) self.cache_buf_idx.index_fill_(0, slots, 0) for buffer in (self.old_x, self.old_B, self.old_dt, self.old_dA_cumsum): diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index 29f1f5681eda..5649d732bbf2 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -37,10 +37,10 @@ _get_local_mamba_cache_layout, ) from tensorrt_llm._torch.pyexecutor.resource_manager import get_pp_layers -from tensorrt_llm._utils import binding_to_torch_dtype +from tensorrt_llm._utils import TensorWrapper, binding_to_torch_dtype, convert_to_torch_tensor from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp from tensorrt_llm.logger import logger -from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole if TYPE_CHECKING: @@ -265,6 +265,60 @@ def get_index_k_buffer( kv_layout=kv_layout, ) + def _get_view_by_role_and_layer( + self, + layer_idx: int, + roles: tuple[DataRole, ...], + *, + dtype: torch.dtype, + page_shape: list[int], + ) -> torch.Tensor: + """Return a zero-copy strided view of one layer's selected roles across the pool. + + The view has shape [slots, roles, *page_shape] and preserves the physical + slot stride when layers or roles share a pool, e.g. Qwen DSA. Selected + roles must be adjacent and equal-sized; expanded page mappings are unsupported. + """ + if not roles: + raise ValueError("A slot view requires at least one role") + layer_offset = self.layer_offsets[layer_idx] + first_role = roles[0] + address = self.impl.get_mem_pool_base_address( + layer_offset, first_role, PageIndexMode.SHARED + ) + page_stride = self.impl.get_page_stride(layer_offset, first_role) + expected_stride = math.prod(page_shape) * dtype.itemsize + if page_stride != expected_stride: + raise RuntimeError( + f"Slot-view page stride mismatch: expected {expected_stride}, got {page_stride}" + ) + for offset, role in enumerate(roles[1:], start=1): + role_address = self.impl.get_mem_pool_base_address( + layer_offset, role, PageIndexMode.SHARED + ) + if ( + role_address != address + offset * page_stride + or self.impl.get_page_stride(layer_offset, role) != page_stride + ): + raise RuntimeError(f"Slot-view roles are not adjacent equal-sized pages: {roles}") + + converter = self.impl.get_page_index_converter(layer_offset, first_role) + scale = int(converter.scale) + if scale < len(roles): + raise RuntimeError( + f"Invalid slot-view page-index scale: {scale} for {len(roles)} roles" + ) + if int(converter.expansion) != 1: + raise RuntimeError("Slot views do not support expanded page indices") + page_upper = self.impl.get_page_index_upper_bound(layer_offset, first_role) + num_pages_with_offset = page_upper + int(converter.layer_offset) + if num_pages_with_offset % scale != 0: + raise RuntimeError("Slot-view page mapping is inconsistent") + view = convert_to_torch_tensor( + TensorWrapper(address, dtype, [num_pages_with_offset // scale, scale, *page_shape]) + ) + return view[:, : len(roles)] + @override def get_buffers( self, @@ -295,7 +349,7 @@ def get_buffers( page_shape = [self.tokens_per_block, num_heads, head_dim] else: page_shape = [num_heads, self.tokens_per_block, head_dim] - return self._get_slot_role_view( + return self._get_view_by_role_and_layer( layer_idx, (Role.KEY, Role.VALUE), dtype=torch_dtype, page_shape=page_shape ) @@ -304,7 +358,7 @@ def get_qsa_position_buffer(self) -> Optional[torch.Tensor]: """Return per-token three-axis RoPE/mRoPE position coordinates.""" if self.qsa_position_layer_id is None: return None - return self._get_slot_role_view( + return self._get_view_by_role_and_layer( self.qsa_position_layer_id, (QSA_INDEX_POSITION,), dtype=QSA_POSITION_CACHE_DTYPE, @@ -315,11 +369,26 @@ def get_qsa_attention_pool_layout(self) -> tuple[int, int]: """Require one block-table mapping for all local sparse layers.""" if self.qsa_position_layer_id is None: raise RuntimeError("QSA cache manager has no local sparse layer") - mapping = self._get_attention_slot_mapping(self.qsa_position_layer_id) + local_idx = self.layer_offsets[self.qsa_position_layer_id] + pool_id = self.layer_to_pool_mapping_dict[local_idx] + converter = self.impl.get_page_index_converter(local_idx, Role.KEY) + mapping = (pool_id, int(converter.scale)) for layer_id in self.qsa_sparse_layer_ids: if layer_id not in self.layer_offsets: continue - candidate = self._get_attention_slot_mapping(layer_id) + local_idx = self.layer_offsets[layer_id] + pool_id = self.layer_to_pool_mapping_dict[local_idx] + if pool_id >= self.num_attention_op_pools: + raise RuntimeError( + f"KV pool {pool_id} is not represented in the attention block table" + ) + converter = self.impl.get_page_index_converter(local_idx, Role.KEY) + scale = int(converter.scale) + if scale <= 0: + raise RuntimeError(f"Invalid attention page-index scale: {scale}") + if int(converter.expansion) != 1: + raise RuntimeError("Attention slot mapping does not support expanded page indices") + candidate = (pool_id, scale) if candidate != mapping: raise RuntimeError( "QSA local layers do not share one attention page mapping: " diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index cef484fa73d9..5a968b40c052 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -65,7 +65,6 @@ MambaHybridCacheManagerV2, MixedMambaHybridCacheManager, use_py_mamba_cache_manager) -from .kv_cache.mamba_cache_manager.legacy import warn_legacy_mamba_cache_manager from .llm_request import ExecutorResponse, LlmRequestState from .model_engine import PyTorchModelEngine from .py_executor import PyExecutor @@ -120,13 +119,6 @@ def _resolve_disagg_transceiver_route( return backend, runtime -def _select_legacy_mamba_manager(manager_cls): - implementation = ("python" if issubclass( - manager_cls, MixedMambaHybridCacheManager) else "cpp") - warn_legacy_mamba_cache_manager(implementation, manager_cls.__name__) - return manager_cls - - def get_kv_cache_manager_cls( model_config: ModelConfig, kv_cache_config: KvCacheConfig, @@ -199,10 +191,10 @@ def get_kv_cache_manager_cls( logger.info( "Using CppMambaHybridCacheManager for Kimi K3 hybrid " "model (block reuse enabled)") - return _select_legacy_mamba_manager(CppMambaHybridCacheManager) + return CppMambaHybridCacheManager logger.info( "Using MixedMambaHybridCacheManager for Kimi K3 hybrid model") - return _select_legacy_mamba_manager(MixedMambaHybridCacheManager) + return MixedMambaHybridCacheManager # Skip Softmax only changes attention kernels. Hybrid models still # need a Mamba-capable cache manager for recurrent state. @@ -238,14 +230,12 @@ def get_kv_cache_manager_cls( "and transceiver_runtime='PYTHON' requires " "use_kv_cache_manager_v2=True.") if kv_cache_config.enable_block_reuse: - return _select_legacy_mamba_manager( - CppMambaHybridCacheManager) + return CppMambaHybridCacheManager if runtime == "PYTHON" and backend == "NIXL": logger.info("Python transceiver detected; using " "MixedMambaHybridCacheManager for hybrid model") - return _select_legacy_mamba_manager( - MixedMambaHybridCacheManager) - return _select_legacy_mamba_manager(CppMambaHybridCacheManager) + return MixedMambaHybridCacheManager + return CppMambaHybridCacheManager if use_py_mamba_cache_manager() and not is_disagg: if use_v2: @@ -260,7 +250,7 @@ def get_kv_cache_manager_cls( "TRTLLM_USE_PY_MAMBA to use the configured cache manager.") logger.info( "Using MixedMambaHybridCacheManager for hybrid mamba model") - return _select_legacy_mamba_manager(MixedMambaHybridCacheManager) + return MixedMambaHybridCacheManager env_override = os.environ.get('TLLM_MAMBA_MANAGER_PREFERENCE', None) if env_override is not None: env_override = env_override.upper() @@ -280,8 +270,7 @@ def get_kv_cache_manager_cls( "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=MIXED " "overrides the default Mamba cache manager to " "MixedMambaHybridCacheManager.") - return _select_legacy_mamba_manager( - MixedMambaHybridCacheManager) + return MixedMambaHybridCacheManager if env_override == 'CPP': if use_v2: raise ValueError( @@ -291,14 +280,14 @@ def get_kv_cache_manager_cls( "Environment variable TLLM_MAMBA_MANAGER_PREFERENCE=CPP " "overrides the default Mamba cache manager to " "CppMambaHybridCacheManager.") - return _select_legacy_mamba_manager(CppMambaHybridCacheManager) + return CppMambaHybridCacheManager logger.warning( f"Unrecognized value for TLLM_MAMBA_MANAGER_PREFERENCE: {env_override}. " "Expected 'CPP' or 'MIXED'. Using the configured " "KV cache manager default.") if not use_v2: - return _select_legacy_mamba_manager(CppMambaHybridCacheManager) + return CppMambaHybridCacheManager if (kv_cache_config.enable_block_reuse and kv_cache_config.enable_kv_pool_rebalance): @@ -2389,6 +2378,7 @@ def _get_mamba_cache_layer_masks( _CONV_STATE_LAYOUT_BY_MODEL_TYPE = { "nemotron_hybrid": "x_b_c", "qwen3_next": "q_k_v", + "kimi_linear": "q_k_v", } @@ -2640,19 +2630,8 @@ def _create_kv_cache_manager( # manager's own internal gate (`tp_size = 1 if enable_attention_dp # else tp_size`, then num_heads / n_groups / conv_dim divide by # it), so the params pass through unscaled. - from ..modules.kimi_kda.cache_manager import select_kda_replay_state - - kimi_extra_kwargs = {} - if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - replay = select_kda_replay_state(spec_config=spec_config, - manager_cls=kv_cache_manager_cls) - if replay is not None: - kimi_extra_kwargs[ - "kda_replay_num_spec"] = replay.num_speculative_tokens - # KDA's conv state is a [Q | K | V] concatenation whose three sections - # have identical width, i.e. the qwen3_next section layout. - kimi_extra_kwargs.update( - _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next")) + kimi_extra_kwargs = _mamba_conv_layout_kwargs(kv_cache_manager_cls, + "kimi_linear") kv_cache_manager = kv_cache_manager_cls( # mamba (KDA) cache parameters mamba_params.state_size, @@ -2738,15 +2717,6 @@ def _create_kv_cache_manager( stochastic_rounding = getattr( quant_config, 'mamba_ssm_stochastic_rounding', False) if quant_config is not None else False - from ..modules.mamba.cache_manager import select_mamba2_state - - use_replay = None - if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - use_replay = select_mamba2_state( - spec_config=spec_config, - ssm_cache_dtype=mamba_params.mamba_ssm_cache_dtype, - stochastic_rounding=stochastic_rounding).uses_replay - # Stochastic-rounding seeds must live on the cache manager (not be # re-created with torch.randint per forward) whenever SR can fire # on the fp16 SSM cache. This mirrors the predicate the mixer uses @@ -2786,7 +2756,7 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, - use_replay_state_update=use_replay, + use_replay_state_update=None, mamba_ssm_stochastic_rounding=mamba_ssm_stochastic_rounding, **mamba_manager_extra_kwargs, ) @@ -2813,15 +2783,6 @@ def _create_kv_cache_manager( )) num_mamba_layers = (0 if is_draft and mamba_params.num_draft_layers > 0 else mamba_params.num_mamba_layers) - from ..modules.fla.cache_manager import select_gdn_replay_state - - use_replay = None - if not issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2): - use_replay = select_gdn_replay_state( - spec_config=spec_config, - ssm_cache_dtype=mamba_params.mamba_ssm_cache_dtype, - manager_cls=kv_cache_manager_cls) is not None - mamba_manager_extra_kwargs = dict(manager_extra_kwargs) mamba_manager_extra_kwargs.update( _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next")) @@ -2858,7 +2819,7 @@ def _create_kv_cache_manager( spec_config=spec_config, is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, - use_replay_state_update=use_replay, + use_replay_state_update=None, **mamba_manager_extra_kwargs, ) else: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 7b17ccd33f46..57a3140595ce 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -2975,75 +2975,6 @@ def get_buffers(self, layer_idx: int, kv_layout: str = "NHD") -> Optional[torch. ) ) - def _get_slot_role_view( - self, - layer_idx: int, - roles: tuple[DataRole, ...], - *, - dtype: torch.dtype, - page_shape: list[int], - ) -> torch.Tensor: - """View adjacent equal-sized roles as [slots, roles, *page_shape]. - - Unlike a packed page view, this retains the physical lifecycle-slot - stride when other layers or roles share the same pool. The caller - supplies the registered role shape and dtype; expanded page mappings - cannot be represented by this view. - """ - if not roles: - raise ValueError("A slot view requires at least one role") - layer_offset = self.layer_offsets[layer_idx] - first_role = roles[0] - address = self.impl.get_mem_pool_base_address( - layer_offset, first_role, PageIndexMode.SHARED - ) - page_stride = self.impl.get_page_stride(layer_offset, first_role) - expected_stride = math.prod(page_shape) * dtype.itemsize - if page_stride != expected_stride: - raise RuntimeError( - f"Slot-view page stride mismatch: expected {expected_stride}, got {page_stride}" - ) - for offset, role in enumerate(roles[1:], start=1): - role_address = self.impl.get_mem_pool_base_address( - layer_offset, role, PageIndexMode.SHARED - ) - if ( - role_address != address + offset * page_stride - or self.impl.get_page_stride(layer_offset, role) != page_stride - ): - raise RuntimeError(f"Slot-view roles are not adjacent equal-sized pages: {roles}") - - converter = self.impl.get_page_index_converter(layer_offset, first_role) - scale = int(converter.scale) - if scale < len(roles): - raise RuntimeError( - f"Invalid slot-view page-index scale: {scale} for {len(roles)} roles" - ) - if int(converter.expansion) != 1: - raise RuntimeError("Slot views do not support expanded page indices") - page_upper = self.impl.get_page_index_upper_bound(layer_offset, first_role) - num_pages_with_offset = page_upper + int(converter.layer_offset) - if num_pages_with_offset % scale != 0: - raise RuntimeError("Slot-view page mapping is inconsistent") - view = convert_to_torch_tensor( - TensorWrapper(address, dtype, [num_pages_with_offset // scale, scale, *page_shape]) - ) - return view[:, : len(roles)] - - def _get_attention_slot_mapping(self, layer_idx: int) -> tuple[int, int]: - """Return an attention layer's block-table pool and unexpanded slot scale.""" - local_idx = self.layer_offsets[layer_idx] - pool_id = self.layer_to_pool_mapping_dict[local_idx] - if pool_id >= self.num_attention_op_pools: - raise RuntimeError(f"KV pool {pool_id} is not represented in the attention block table") - converter = self.impl.get_page_index_converter(local_idx, Role.KEY) - scale = int(converter.scale) - if scale <= 0: - raise RuntimeError(f"Invalid attention page-index scale: {scale}") - if int(converter.expansion) != 1: - raise RuntimeError("Attention slot mapping does not support expanded page indices") - return pool_id, scale - def get_index_k_buffer( self, layer_idx: int, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py index e89c039c78b7..b9782e976740 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py @@ -91,7 +91,7 @@ def _estimate_mamba_hybrid_cache_cost(model_config, mapping, **kwargs): def warn_legacy_mamba_cache_manager( implementation: Literal["python", "cpp"], manager_name: str ) -> None: - """Warn once when a V1 implementation is selected or constructed.""" + """Warn once when a V1 implementation is constructed.""" display_name = "Python" if implementation == "python" else "C++" logger.warning_once( f"{manager_name} uses the deprecated {display_name} Mamba cache " @@ -101,6 +101,50 @@ def warn_legacy_mamba_cache_manager( ) +def _resolve_legacy_replay_options( + manager_cls: type, + model_type: str, + spec_config: object | None, + ssm_cache_dtype: torch.dtype, + stochastic_rounding: bool, + use_replay: bool | None, + kda_num_spec: int | None = None, +) -> tuple[str, bool, int | None]: + """Resolve legacy replay policy before allocating model state. + + Kimi uses the Qwen convolution layout but needs its own replay selection. + An explicit Mamba2 replay flag retains the existing caller override. + """ + if model_type == "kimi_linear": + from tensorrt_llm._torch.modules.kimi_kda.cache_manager import select_kda_replay_state + + if kda_num_spec is None: + state = select_kda_replay_state(spec_config=spec_config, manager_cls=manager_cls) + if state is not None: + kda_num_spec = state.num_speculative_tokens + model_type = "qwen3_next" + elif model_type == "nemotron_hybrid" and use_replay is None: + from tensorrt_llm._torch.modules.mamba.cache_manager import select_mamba2_state + + use_replay = select_mamba2_state( + spec_config=spec_config, + ssm_cache_dtype=ssm_cache_dtype, + stochastic_rounding=stochastic_rounding, + ).uses_replay + elif model_type == "qwen3_next" and use_replay is None: + from tensorrt_llm._torch.modules.fla.cache_manager import select_gdn_replay_state + + use_replay = ( + select_gdn_replay_state( + spec_config=spec_config, + ssm_cache_dtype=ssm_cache_dtype, + manager_cls=manager_cls, + ) + is not None + ) + return model_type, bool(use_replay), kda_num_spec + + class PythonMambaCacheManager(BaseResourceManager): """Pure-Python mamba state manager with speculative decoding support. @@ -1397,7 +1441,7 @@ def __init__( execution_stream: Optional[torch.cuda.Stream] = None, model_type: str = "nemotron_hybrid", is_draft: bool = False, - use_replay_state_update: bool = False, + use_replay_state_update: bool | None = False, mamba_ssm_stochastic_rounding: bool = False, kda_replay_num_spec: Optional[int] = None, # Per-pool configurations forwarded to the C++ KVCacheManager ctor. @@ -1407,6 +1451,15 @@ def __init__( ) -> None: # mamba hybrid cache requires block reuse to be disabled in KV cache config warn_legacy_mamba_cache_manager("python", type(self).__name__) + model_type, use_replay_state_update, kda_replay_num_spec = _resolve_legacy_replay_options( + type(self), + model_type, + spec_config, + mamba_ssm_cache_dtype, + mamba_ssm_stochastic_rounding, + use_replay_state_update, + kda_replay_num_spec, + ) assert not kv_cache_config.enable_block_reuse, ( "mamba hybrid cache requires block reuse to be disabled in KV cache config" ) @@ -1593,7 +1646,7 @@ def __init__( layer_mask: Optional[List[bool]] = None, # this is the full attention layer mask is_estimating_kv_cache: bool = False, is_draft: bool = False, - use_replay_state_update: bool = False, + use_replay_state_update: bool | None = False, mamba_ssm_stochastic_rounding: bool = False, model_type: str = "nemotron_hybrid", **kwargs, @@ -1603,6 +1656,14 @@ def __init__( # 2) Full attention layers (full_attention_layer_mask is True) # 3) Not managed layers (both masks are False) warn_legacy_mamba_cache_manager("cpp", type(self).__name__) + model_type, use_replay_state_update, _ = _resolve_legacy_replay_options( + type(self), + model_type, + spec_config, + mamba_ssm_cache_dtype, + mamba_ssm_stochastic_rounding, + use_replay_state_update, + ) total_layers = len(mamba_layer_mask) if layer_mask is None: full_attention_layer_mask = [False] * total_layers diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index c79ce8d0b4fc..00561e853005 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -7,7 +7,6 @@ import torch from tensorrt_llm._torch.attention.backends.sparse.qsa import ( - QSAMambaHybridCacheManagerV2, QSASparseMetadataParams, QSASparseParams, ) @@ -35,6 +34,7 @@ select_qsa_paged_tokens, select_qsa_tokens, ) +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode @@ -1247,7 +1247,8 @@ def test_fused_qsa_prefill_bounds_sparse_attention_to_visible_tokens() -> None: def test_qsa_side_buffers_use_exact_geometry_and_one_position_role() -> None: - manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager._qsa_enabled = True manager.qsa_index_dim = 128 manager.qsa_index_kv_heads = 1 manager.qsa_sparse_layer_ids = [7, 11] @@ -1272,13 +1273,13 @@ def get_page_stride(self, layer_idx, role): del layer_idx, role return 0 - manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager.qsa_position_layer_id = 7 manager.layer_offsets = {7: 2} manager.impl = _Impl() manager.tokens_per_block = 128 - with pytest.raises(RuntimeError, match="position-cache page stride mismatch"): + with pytest.raises(RuntimeError, match="Slot-view page stride mismatch"): manager.get_qsa_position_buffer() assert len(calls) == 1 diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 088334d9c0bf..7ccb1bf2421d 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -375,12 +375,10 @@ def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout( assert "model_type" not in kwargs -def test_kimi_v1_manager_still_selects_qwen3_next_model_type( +def test_kimi_v1_manager_receives_model_identity_for_replay_selection( monkeypatch: pytest.MonkeyPatch, ) -> None: - """The V1 managers have no `conv_state_layout` parameter; they must keep - getting `model_type='qwen3_next'` and must not receive V2 KDA replay - parameters (TRTLLM-15216 regression guard).""" + """V1 resolves Kimi replay and its Q/K/V layout inside the constructor.""" captured: dict[str, object] = {} class RecordingV1Manager(CppMambaHybridCacheManager): @@ -410,7 +408,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: is_draft=False, ) kwargs = captured["kwargs"] - assert kwargs["model_type"] == "qwen3_next" + assert kwargs["model_type"] == "kimi_linear" assert "conv_state_layout" not in kwargs assert "kda_replay_num_spec" not in kwargs @@ -553,10 +551,10 @@ def __init__(self, *args, **kwargs): **common_kwargs, ) - assert captured_cpp["use_replay_state_update"] is True + assert captured_cpp["use_replay_state_update"] is None assert captured_cpp["model_type"] == "qwen3_next" assert captured_cpp["max_num_tokens"] == 256 - assert captured_mixed["use_replay_state_update"] is False + assert captured_mixed["use_replay_state_update"] is None assert captured_mixed["model_type"] == "qwen3_next" assert captured_mixed["max_num_tokens"] == 256 assert captured_v2["use_replay_state_update"] is None @@ -574,6 +572,19 @@ def __init__(self, *args, **kwargs): assert captured_v2["max_num_tokens"] == 256 assert "model_type" not in captured_v2 assert captured_v2["conv_state_layout"] == "q_k_v" + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.legacy import ( + _resolve_legacy_replay_options, + ) + + for manager_cls, enabled in ((RecordingCppManager, True), (RecordingMixedManager, False)): + assert _resolve_legacy_replay_options( + manager_cls, + "qwen3_next", + common_kwargs["spec_config"], + torch.bfloat16, + False, + None, + ) == ("qwen3_next", enabled, None) fallback_logs = [str(call.args[0]) for call in info_log.call_args_list] assert any("RecordingMixedManager was selected" in log for log in fallback_logs) assert not any("RecordingV2Manager was selected" in log for log in fallback_logs) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index 58e1ea777775..a640fb4e1dde 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -25,7 +25,7 @@ Qwen4ExpHybridCacheManagerV2, Qwen4ExpPLECacheParams, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( IntermediateState, @@ -592,9 +592,12 @@ def test_gdn_bind_accepts_affine_and_indirect_views(affine): assert all(not owner.use_gdn_cached_replay_all_layer_commit for owner in owners) -@pytest.mark.parametrize("state", [IntermediateState(), GDNReplayState(3), KDAReplayState(2)]) +@pytest.mark.parametrize( + "state", [IntermediateState(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)] +) def test_no_local_recurrent_layers_allocate_no_algorithm_buffers(state): state.bind(_layout(layers=(), slot_capacity=0), [], []) + state.reset_slots(torch.tensor([0]), [0]) assert not any(isinstance(value, torch.Tensor) for value in vars(state).values()) @@ -803,14 +806,6 @@ def test_algorithm_methods_do_not_depend_on_manager(cls): }.intersection(code.co_names) -def test_qwen4_methods_do_not_access_raw_pool(): - for function in vars(Qwen4ExpHybridCacheManagerV2).values(): - if isinstance(function, FunctionType): - assert not {"impl", "TensorWrapper", "convert_to_torch_tensor"}.intersection( - function.__code__.co_names - ) - - @pytest.mark.parametrize( "manager_cls", [ @@ -885,7 +880,7 @@ def test_k3_host_acceptance_handles_empty_mixed_and_completed_requests(draft_len def _slot_view_manager( monkeypatch, *, scale=6, role_count=2, page_shape=(4, 2, 3), dtype=torch.float32 ): - manager = object.__new__(KVCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager.layer_offsets = {42: 0} page_elements = 1 for dim in page_shape: @@ -907,7 +902,7 @@ def _slot_view_manager( get_page_index_converter=lambda layer, role: converter, get_page_index_upper_bound=lambda layer, role: 3 * scale - layer_offset, ) - globals_ = KVCacheManagerV2._get_slot_role_view.__globals__ + globals_ = Qwen4ExpHybridCacheManagerV2._get_view_by_role_and_layer.__globals__ monkeypatch.setitem(globals_, "TensorWrapper", lambda address, dtype, shape: (address, shape)) def convert(wrapper): @@ -932,7 +927,7 @@ def test_slot_role_view_preserves_coalesced_stride( manager, roles, shape, converter, backing = _slot_view_manager( monkeypatch, scale=scale, role_count=role_count, page_shape=page_shape, dtype=dtype ) - view = manager._get_slot_role_view(42, roles, dtype=dtype, page_shape=shape) + view = manager._get_view_by_role_and_layer(42, roles, dtype=dtype, page_shape=shape) assert view.shape == (3, role_count, *page_shape) assert view.stride(0) == scale * 24 assert view.stride(1) == 24 @@ -956,20 +951,22 @@ def test_slot_role_view_rejects_incompatible_physical_layout(monkeypatch, invali else: manager.impl.get_page_index_upper_bound = lambda layer, role: 1 with pytest.raises(RuntimeError): - manager._get_slot_role_view(42, roles, dtype=torch.float32, page_shape=shape) + manager._get_view_by_role_and_layer(42, roles, dtype=torch.float32, page_shape=shape) -def test_qwen4_uses_kv_owned_slot_views(monkeypatch): +def test_qwen4_position_uses_model_owned_slot_view(monkeypatch): manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager.qsa_position_layer_id = 42 manager.tokens_per_block = 4 - manager._get_slot_role_view = Mock(return_value=torch.zeros(3, 1, 4, 3, dtype=torch.int32)) + manager._get_view_by_role_and_layer = Mock( + return_value=torch.zeros(3, 1, 4, 3, dtype=torch.int32) + ) assert manager.get_qsa_position_buffer().shape == (3, 4, 3) - manager._get_slot_role_view.assert_called_once() + manager._get_view_by_role_and_layer.assert_called_once() manager.qsa_position_layer_id = None - manager._get_slot_role_view.reset_mock() + manager._get_view_by_role_and_layer.reset_mock() assert manager.get_qsa_position_buffer() is None - manager._get_slot_role_view.assert_not_called() + manager._get_view_by_role_and_layer.assert_not_called() @pytest.mark.parametrize("kv_layout", ["NHD", "HND"]) @@ -985,9 +982,9 @@ def test_qwen4_main_kv_view_supplies_only_logical_geometry(monkeypatch, kv_layou manager.tokens_per_block = 4 monkeypatch.setitem(namespace, "binding_to_torch_dtype", lambda dtype: torch.bfloat16) sentinel = object() - manager._get_slot_role_view = Mock(return_value=sentinel) + manager._get_view_by_role_and_layer = Mock(return_value=sentinel) assert manager.get_buffers(42, kv_layout) is sentinel - manager._get_slot_role_view.assert_called_once_with( + manager._get_view_by_role_and_layer.assert_called_once_with( 42, (Role.KEY, Role.VALUE), dtype=torch.bfloat16, @@ -1000,9 +997,11 @@ def test_qwen4_main_kv_view_supplies_only_logical_geometry(monkeypatch, kv_layou @pytest.mark.parametrize("invalid", [None, "pool", "scale", "expansion"]) -def test_attention_slot_mapping_validates_backend_contract(invalid): - manager = object.__new__(KVCacheManagerV2) +def test_qwen4_attention_pool_layout_validates_backend_contract(invalid): + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager.layer_offsets = {42: 0} + manager.qsa_position_layer_id = 42 + manager.qsa_sparse_layer_ids = [42] manager.layer_to_pool_mapping_dict = {0: 0} manager.num_attention_op_pools = 1 converter = SimpleNamespace(scale=6, expansion=1) @@ -1014,21 +1013,29 @@ def test_attention_slot_mapping_validates_backend_contract(invalid): elif invalid == "expansion": converter.expansion = 2 if invalid is None: - assert manager._get_attention_slot_mapping(42) == (0, 6) + assert manager.get_qsa_attention_pool_layout() == (0, 6) else: with pytest.raises(RuntimeError): - manager._get_attention_slot_mapping(42) + manager.get_qsa_attention_pool_layout() -def test_qwen4_requires_one_mapping_but_does_not_read_physical_pool(): +def test_qwen4_requires_one_mapping_and_skips_nonlocal_layers(): manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager.qsa_position_layer_id = 42 manager.qsa_sparse_layer_ids = [42, 43, 99] manager.layer_offsets = {42: 0, 43: 1} - manager._get_attention_slot_mapping = Mock(return_value=(0, 6)) + manager.layer_to_pool_mapping_dict = {0: 0, 1: 0} + manager.num_attention_op_pools = 1 + converters = { + 0: SimpleNamespace(scale=6, expansion=1), + 1: SimpleNamespace(scale=6, expansion=1), + } + manager.impl = SimpleNamespace( + get_page_index_converter=Mock(side_effect=lambda layer, role: converters[layer]) + ) assert manager.get_qsa_attention_pool_layout() == (0, 6) - assert all(call.args != (99,) for call in manager._get_attention_slot_mapping.call_args_list) - manager._get_attention_slot_mapping = lambda layer: (0, 6 if layer == 42 else 4) + assert {call.args[0] for call in manager.impl.get_page_index_converter.call_args_list} == {0, 1} + converters[1].scale = 4 with pytest.raises(RuntimeError, match="do not share"): manager.get_qsa_attention_pool_layout() @@ -1124,3 +1131,98 @@ def _get_state_buffer(self, local_layer_idx, role, dtype, state_shape): PLE_NGRAM_CONTEXT, ] assert manager.ple_layer_cache(0)[0].shape == (5, 16, 6) + + +@pytest.mark.parametrize("model_type", ["nemotron_hybrid", "qwen3_next", "kimi_linear"]) +@pytest.mark.parametrize("enabled", [False, True]) +def test_legacy_constructor_resolves_model_replay(monkeypatch, model_type, enabled): + from tensorrt_llm._torch.modules.fla import cache_manager as gdn + from tensorrt_llm._torch.modules.kimi_kda import cache_manager as kda + from tensorrt_llm._torch.modules.mamba import cache_manager as mamba + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import legacy + + spec = SimpleNamespace(tokens_per_gen_step=4) + selectors = [ + Mock(return_value=SimpleNamespace(uses_replay=enabled)), + Mock(return_value=object() if enabled else None), + Mock(return_value=SimpleNamespace(num_speculative_tokens=3) if enabled else None), + ] + monkeypatch.setattr(mamba, "select_mamba2_state", selectors[0]) + monkeypatch.setattr(gdn, "select_gdn_replay_state", selectors[1]) + monkeypatch.setattr(kda, "select_kda_replay_state", selectors[2]) + mamba_init, kv_init = Mock(), Mock() + monkeypatch.setattr(legacy.MambaCacheManager, "__init__", mamba_init) + monkeypatch.setattr(legacy.KVCacheManager, "__init__", kv_init) + legacy.MixedMambaHybridCacheManager( + 4, + 4, + 2, + 2, + 4, + 1, + [True, False], + torch.float32, + torch.float32, + SimpleNamespace(enable_block_reuse=False), + legacy.CacheTypeCpp.SELF, + num_layers=1, + layer_mask=[False, True], + num_kv_heads=2, + head_dim=4, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=2, + mapping=Mapping(world_size=1, tp_size=1, pp_size=1), + spec_config=spec, + model_type=model_type, + use_replay_state_update=None, + ) + expected_selector = ["nemotron_hybrid", "qwen3_next", "kimi_linear"].index(model_type) + assert [selector.call_count for selector in selectors] == [ + int(i == expected_selector) for i in range(3) + ] + kwargs = mamba_init.call_args.kwargs + assert kwargs["model_type"] == ( + "nemotron_hybrid" if model_type == "nemotron_hybrid" else "qwen3_next" + ) + assert kwargs["use_replay_state_update"] == (enabled and model_type != "kimi_linear") + assert kwargs["kda_replay_num_spec"] == (3 if enabled and model_type == "kimi_linear" else None) + kv_init.assert_called_once() + + +@pytest.mark.parametrize("requested", [False, True]) +def test_legacy_replay_explicit_override_bypasses_auto_selection(monkeypatch, requested): + from tensorrt_llm._torch.modules.mamba import cache_manager as mamba + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import legacy + + selector = Mock(side_effect=AssertionError("Explicit replay flag must bypass selection")) + monkeypatch.setattr(mamba, "select_mamba2_state", selector) + assert legacy._resolve_legacy_replay_options( + legacy.CppMambaHybridCacheManager, + "nemotron_hybrid", + object(), + torch.float32, + False, + requested, + ) == ("nemotron_hybrid", requested, None) + selector.assert_not_called() + + +def test_legacy_kda_replay_retains_cpp_manager_gate(monkeypatch): + from tensorrt_llm._torch.modules.kimi_kda import _kda_kernels + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import legacy + + monkeypatch.setattr(_kda_kernels, "is_kda_mtp_verify_available", lambda: True) + spec = SimpleNamespace(tokens_per_gen_step=4) + for manager_cls, expected in ( + (legacy.MixedMambaHybridCacheManager, 3), + (legacy.CppMambaHybridCacheManager, None), + ): + assert legacy._resolve_legacy_replay_options( + manager_cls, + "kimi_linear", + spec, + torch.float32, + False, + False, + ) == ("qwen3_next", False, expected) diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index fe8067319606..d9ddc9bd6ce4 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -10,13 +10,13 @@ from tensorrt_llm._torch.attention.backends.sparse.hooks import get_sparse_attention_hooks from tensorrt_llm._torch.attention.backends.sparse.qsa import ( QSAAttentionMetadata, - QSAMambaHybridCacheManagerV2, QSASparseMetadataParams, QSASparseParams, ) from tensorrt_llm._torch.attention.backends.sparse.qsa.module import QSASparseHooks from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import MambaKVCacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 @@ -167,7 +167,7 @@ def test_qsa_empty_batch_keeps_the_regular_backend_path() -> None: def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers(monkeypatch) -> None: monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) - manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager.qsa_position_layer_id = None metadata = object.__new__(QSAAttentionMetadata) metadata.kv_cache_manager = manager @@ -194,7 +194,7 @@ def test_qsa_hybrid_routes_to_sparse_v2_cache_manager(monkeypatch) -> None: manager_cls = get_kv_cache_manager_cls(model_config, kv_cache_config) - assert manager_cls is QSAMambaHybridCacheManagerV2 + assert manager_cls is Qwen4ExpHybridCacheManagerV2 def test_qsa_hybrid_rejects_kv_cache_manager_v1(monkeypatch) -> None: @@ -269,12 +269,11 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( monkeypatch.setattr( _util, "extract_mamba_kv_cache_params", lambda *args, **kwargs: mamba_params ) - monkeypatch.setattr(_util, "get_sm_version", lambda: 103) monkeypatch.setattr(MambaHybridCacheManagerV2, "__init__", lambda self, *args, **kwargs: None) manager = _create_kv_cache_manager( model_engine=None, - kv_cache_manager_cls=QSAMambaHybridCacheManagerV2, + kv_cache_manager_cls=Qwen4ExpHybridCacheManagerV2, mapping=SimpleNamespace(enable_attention_dp=False), kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=True), tokens_per_block=128, @@ -299,7 +298,7 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( def test_qwen4_manager_also_supports_dense_attention(monkeypatch) -> None: monkeypatch.setattr(MambaHybridCacheManagerV2, "__init__", lambda self, *args, **kwargs: None) - manager = QSAMambaHybridCacheManagerV2(layer_mask=[True]) + manager = Qwen4ExpHybridCacheManagerV2(layer_mask=[True]) assert not manager._qsa_enabled assert manager._extra_buffers_per_layer(tokens_per_block=32) == {} @@ -308,7 +307,7 @@ def test_qwen4_manager_also_supports_dense_attention(monkeypatch) -> None: def test_qsa_cache_manager_delegates_regular_kv_layout( monkeypatch: pytest.MonkeyPatch, dtype: DataType ) -> None: - manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManagerV2) manager.dtype = dtype manager._qsa_enabled = True sentinel = torch.empty(0, dtype=torch.int8) From ad8fe366fb342a68ed3d3d22d51010038d5b0e77 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 21 Sep 2026 06:24:03 +0000 Subject: [PATCH 07/25] [None][refactor] Simplify Mamba layer cache payload hierarchy (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/modules/fla/cache_manager.py | 23 +----- .../_torch/modules/kimi_kda/cache_manager.py | 6 +- .../_torch/modules/mamba/cache_manager.py | 15 ++-- .../kv_cache/mamba_cache_manager/__init__.py | 6 +- .../kv_cache/mamba_cache_manager/common.py | 6 +- .../kv_cache/test_mamba_cache_manager.py | 8 +- .../kv_cache/test_mamba_model_managers.py | 82 +++++++++++++++++++ 7 files changed, 104 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index 8a343f3521cc..e14d46cca02d 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -6,7 +6,6 @@ from __future__ import annotations import sys -from dataclasses import dataclass import torch @@ -15,11 +14,10 @@ else: from typing_extensions import override -from tensorrt_llm._torch.modules.mamba.cache_manager import Mamba2ReplayLayerCache, ReplayHistory +from tensorrt_llm._torch.modules.mamba.cache_manager import ReplayHistory from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( IntermediateState, MambaAcceptanceBatch, - MambaLayerCache, MambaStateLayout, _mamba_effective_tp_size, ) @@ -36,11 +34,6 @@ ) -@dataclass(frozen=True, kw_only=True) -class GDNReplayLayerCache(Mamba2ReplayLayerCache): - """GDN cached-replay tensors for one recurrent layer.""" - - class GDNReplayState(ReplayHistory): """Own GDN all-layer cached-replay checkpoint commits.""" @@ -124,19 +117,6 @@ def update(self, batch: MambaAcceptanceBatch) -> None: self.commit_all_layers(batch.attention_metadata, num_decodes) super().update(batch) - @override - def make_layer_cache( - self, - layer_offset: int, - conv: torch.Tensor, - temporal: torch.Tensor, - ) -> MambaLayerCache: - return GDNReplayLayerCache( - conv=conv, - temporal=temporal, - **self._layer_cache_fields(layer_offset), - ) - def commit_all_layers(self, attention_metadata: object, num_decodes: int) -> None: """Commit every local GDN checkpoint using one partitioned launch.""" states = self._ssm_states @@ -301,7 +281,6 @@ def get_gdn_cache_params(config, *, spec_config=None, quant_config=None): __all__ = [ "GDNReplayState", - "GDNReplayLayerCache", "Qwen35HybridCacheManagerV2", "create_gdn_state", "select_gdn_replay_state", diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index 8443bfbbe9af..77dc857b5ffa 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -25,7 +25,6 @@ MambaLayerCache, MambaStateLayout, ReplayStateUpdateMetadata, - SpeculativeMambaLayerCache, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, @@ -34,7 +33,7 @@ @dataclass(frozen=True, kw_only=True) -class KDAReplayLayerCache(SpeculativeMambaLayerCache): +class KDAReplayLayerCache(MambaLayerCache): """Per-layer KDA replay tensors consumed by fused multi-token verify.""" prev_num_accepted_tokens: torch.Tensor | None = field( @@ -264,9 +263,6 @@ def bind( def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: return { - "intermediate_conv_window": None, - "intermediate_ssm": None, - "mamba_ssm_rand_seed": self.rand_seed, "prev_num_accepted_tokens": self.prev_num_accepted_tokens, "kda_conv_q": None if self.kda_conv_q is None else self.kda_conv_q[layer_offset], "kda_conv_k": None if self.kda_conv_k is None else self.kda_conv_k[layer_offset], diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py index 444648191142..d2f2709404bf 100644 --- a/tensorrt_llm/_torch/modules/mamba/cache_manager.py +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -22,7 +22,6 @@ MambaLayerCache, MambaStateLayout, ReplayStateUpdateMetadata, - SpeculativeMambaLayerCache, _promote_intermediate_states, _stack_state_views, ) @@ -44,9 +43,15 @@ @dataclass(frozen=True, kw_only=True) -class Mamba2ReplayLayerCache(SpeculativeMambaLayerCache): - """Mamba2 compact-replay tensors for one recurrent layer.""" +class ReplayLayerCache(MambaLayerCache): + """Shared Mamba2/GDN replay views consumed by their recurrent mixers.""" + intermediate_conv_window: torch.Tensor | None = None + intermediate_ssm: torch.Tensor | None = None + mamba_ssm_rand_seed: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) prev_num_accepted_tokens: torch.Tensor | None = field( default=None, metadata={"slot_shared": True}, @@ -294,7 +299,7 @@ def make_layer_cache( conv: torch.Tensor, temporal: torch.Tensor, ) -> MambaLayerCache: - return Mamba2ReplayLayerCache( + return ReplayLayerCache( conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset), @@ -460,7 +465,7 @@ def get_nemotron_cache_params(config, *, spec_config=None, quant_config=None): __all__ = [ "Mamba2State", "ReplayHistory", - "Mamba2ReplayLayerCache", + "ReplayLayerCache", "NemotronHybridCacheManagerV2", "create_mamba2_state", "select_mamba2_state", diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py index 9a1f85fd1e4a..2a6a83544775 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py @@ -15,11 +15,11 @@ if TYPE_CHECKING: from .common import ( BaseMambaCacheManager, + IntermediateLayerCache, MambaHybridCacheManager, MambaLayerCache, MambaRole, ReplayStateUpdateMetadata, - SpeculativeMambaLayerCache, ) from .legacy import ( CppMambaHybridCacheManager, @@ -46,7 +46,7 @@ "MixedMambaHybridCacheManager": ("legacy", "MixedMambaHybridCacheManager"), "PythonMambaCacheManager": ("legacy", "PythonMambaCacheManager"), "ReplayStateUpdateMetadata": ("common", "ReplayStateUpdateMetadata"), - "SpeculativeMambaLayerCache": ("common", "SpeculativeMambaLayerCache"), + "IntermediateLayerCache": ("common", "IntermediateLayerCache"), "_advance_replay_state": ("common", "_advance_replay_state"), "_allocate_mamba_seed_buffer": ("common", "_allocate_mamba_seed_buffer"), "_compute_deterministic_mamba_seed": ( @@ -93,6 +93,6 @@ def __dir__() -> list[str]: "MixedMambaHybridCacheManager", "PythonMambaCacheManager", "ReplayStateUpdateMetadata", - "SpeculativeMambaLayerCache", + "IntermediateLayerCache", "use_py_mamba_cache_manager", ] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py index 0df7259d4b4a..4e9ca3ded1ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -79,7 +79,7 @@ def at_layer_idx(self, layer: int) -> "MambaLayerCache": @dataclass(frozen=True, kw_only=True) -class SpeculativeMambaLayerCache(MambaLayerCache): +class IntermediateLayerCache(MambaLayerCache): """Generic intermediate-state view used during speculative decoding.""" intermediate_conv_window: torch.Tensor | None = None @@ -717,7 +717,7 @@ def _estimate_mamba_hybrid_cache_cost( "BaseMambaCacheManager", "MambaHybridCacheManager", "MambaLayerCache", - "SpeculativeMambaLayerCache", + "IntermediateLayerCache", "MambaStateLayout", "MambaAcceptanceBatch", "IntermediateState", @@ -885,6 +885,6 @@ def shutdown(self) -> None: def make_layer_cache( self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor ) -> MambaLayerCache: - return SpeculativeMambaLayerCache( + return IntermediateLayerCache( conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset) ) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 28616508a0b1..6f2c967cb35a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3679,8 +3679,8 @@ def test_v2_kda_replay_policy_allocates_logical_slot_caches(): assert mgr.get_replay_state_update_metadata() is None layer_cache = mgr.mamba_layer_cache(0) - assert layer_cache.intermediate_ssm is None - assert layer_cache.intermediate_conv_window is None + assert not hasattr(layer_cache, "intermediate_ssm") + assert not hasattr(layer_cache, "intermediate_conv_window") cache_size = layer_cache.temporal.shape[0] assert layer_cache.kda_conv_q.shape == (cache_size, 48, 6) assert layer_cache.kda_conv_k.shape == (cache_size, 48, 6) @@ -3705,8 +3705,8 @@ def test_v2_kda_replay_policy_allocates_logical_slot_caches(): assert layer_cache.prev_num_accepted_tokens.data_ptr() == ( policy.prev_num_accepted_tokens.data_ptr() ) - assert layer_cache.intermediate_ssm is None - assert layer_cache.intermediate_conv_window is None + assert not hasattr(layer_cache, "intermediate_ssm") + assert not hasattr(layer_cache, "intermediate_conv_window") finally: mgr.shutdown() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index c0731ebf637b..4e008f1e1c38 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -11,6 +11,7 @@ from tensorrt_llm._torch.modules.fla.cache_manager import GDNReplayState, Qwen35HybridCacheManagerV2 from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAReplayLayerCache, KDAReplayState, KimiK3HybridCacheManagerV2, ) @@ -18,6 +19,7 @@ Mamba2State, NemotronHybridCacheManagerV2, ReplayHistory, + ReplayLayerCache, ) from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( PLE_CONV_STATE, @@ -28,6 +30,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + IntermediateLayerCache, IntermediateState, MambaAcceptanceBatch, MambaHybridCacheManager, @@ -653,6 +656,85 @@ def test_kda_scratch_cost_matches_allocated_slot_buffers(): assert actual == estimated * layout.slot_capacity +@pytest.mark.parametrize("layer", [0, 1]) +@pytest.mark.parametrize("kind", ["intermediate", "mamba2", "replay", "gdn", "kda"]) +def test_layer_payload_preserves_storage_and_strides(kind, layer): + states = { + "intermediate": IntermediateState, + "mamba2": Mamba2State, + "replay": lambda: ReplayHistory(3), + "gdn": lambda: GDNReplayState(3), + "kda": lambda: KDAReplayState(2), + } + state = states[kind]() + layout = replace(_layout(), stochastic_rounding=True) + # Interleave persistent slots to exercise borrowed, non-contiguous pool views. + conv = torch.zeros(2, 10, 24, 3)[:, ::2] + ssm = torch.zeros(2, 10, 2, 4, 4)[:, ::2] + state.bind(layout, list(ssm.unbind()), list(conv.unbind())) + payload = state.make_layer_cache(layer, conv[layer], ssm[layer]) + expected_type = { + "intermediate": IntermediateLayerCache, + "mamba2": IntermediateLayerCache, + "replay": ReplayLayerCache, + "gdn": ReplayLayerCache, + "kda": KDAReplayLayerCache, + }[kind] + assert type(payload) is expected_type + pairs = [(payload.conv, conv[layer]), (payload.temporal, ssm[layer])] + if kind == "kda": + assert payload.has_kda_replay_caches + assert not hasattr(payload, "intermediate_ssm") + assert not hasattr(payload, "intermediate_conv_window") + assert not hasattr(payload, "mamba_ssm_rand_seed") + for name in ( + "kda_conv_q", + "kda_conv_k", + "kda_conv_v", + "kda_qkg_cache", + "kda_v_cache", + "kda_beta_cache", + ): + pairs.append((getattr(payload, name), getattr(state, name)[layer])) + assert payload.prev_num_accepted_tokens is state.prev_num_accepted_tokens + else: + pairs.append((payload.intermediate_conv_window, state.intermediate_conv[layer])) + if kind in ("replay", "gdn"): + assert payload.intermediate_ssm is None + assert payload.cache_buf_idx is state.cache_buf_idx + assert payload.prev_num_accepted_tokens is state.prev_num_accepted_tokens + for name in ("old_x", "old_B", "old_dt", "old_dA_cumsum"): + pairs.append((getattr(payload, name), getattr(state, name)[layer])) + else: + pairs.append((payload.intermediate_ssm, state.intermediate_ssm[layer])) + if kind == "intermediate": + assert payload.mamba_ssm_rand_seed is None + else: + assert payload.mamba_ssm_rand_seed is state.rand_seed + for view, source in pairs: + assert view.data_ptr() == source.data_ptr() + assert view.shape == source.shape + assert view.stride() == source.stride() + view.fill_(7) + assert torch.all(source == 7) + state.shutdown() + + +def test_kda_seed_lifecycle_is_independent_of_layer_payload(): + state = KDAReplayState(2) + layout = replace(_layout(), stochastic_rounding=True, seed_rank_offset=17) + views = _state_views() + state.bind(layout, views.all_ssm_states, views.all_conv_states) + original = state.rand_seed.clone() + payload = state.make_layer_cache(0, views.all_conv_states[0], views.all_ssm_states[0]) + assert not hasattr(payload, "mamba_ssm_rand_seed") + state.reset_slots(torch.tensor([1, 3]), [1, 3]) + assert torch.all(state.rand_seed[[1, 3]] != original[[1, 3]]) + torch.testing.assert_close(state.rand_seed[[0, 2, 4]], original[[0, 2, 4]]) + state.shutdown() + assert state.rand_seed is None + + def test_mamba2_seed_lifecycle_does_not_require_speculative_decoding(): layout = replace( _layout(), spec_config=None, stochastic_rounding=True, ssm_state_dtype=torch.float16 From 3f4f80814e2997d275d72136badc5506fe10c1ce Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 21 Sep 2026 06:55:11 +0000 Subject: [PATCH 08/25] [None][refactor] Decouple GDN and KDA caches from Mamba2 (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/modules/fla/cache_manager.py | 2 +- .../_torch/modules/kimi_kda/cache_manager.py | 8 +- .../_torch/modules/mamba/cache_manager.py | 292 +----------------- .../kv_cache/mamba_cache_manager/common.py | 12 +- .../kv_cache/mamba_cache_manager/replay.py | 246 +++++++++++++++ .../kv_cache/mamba_cache_manager/seeds.py | 66 ++++ .../kv_cache/test_mamba_cache_manager.py | 12 +- .../kv_cache/test_mamba_model_managers.py | 8 +- 8 files changed, 341 insertions(+), 305 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py create mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index e14d46cca02d..90ab9de62076 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -14,7 +14,6 @@ else: from typing_extensions import override -from tensorrt_llm._torch.modules.mamba.cache_manager import ReplayHistory from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( IntermediateState, MambaAcceptanceBatch, @@ -24,6 +23,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayHistory from tensorrt_llm._torch.utils import is_gdn_replay_enabled from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index 77dc857b5ffa..1677b936534d 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -15,10 +15,6 @@ else: from typing_extensions import override -from tensorrt_llm._torch.modules.mamba.cache_manager import ( - allocate_mamba_seed_buffer, - compute_deterministic_mamba_seed, -) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( IntermediateState, MambaAcceptanceBatch, @@ -29,6 +25,10 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.seeds import ( + allocate_mamba_seed_buffer, + compute_deterministic_mamba_seed, +) from tensorrt_llm.logger import logger diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py index d2f2709404bf..bd8a7451ad77 100644 --- a/tensorrt_llm/_torch/modules/mamba/cache_manager.py +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -7,7 +7,6 @@ import os import sys -from dataclasses import dataclass, field import torch @@ -18,128 +17,26 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( IntermediateState, - MambaAcceptanceBatch, - MambaLayerCache, MambaStateLayout, - ReplayStateUpdateMetadata, - _promote_intermediate_states, - _stack_state_views, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) -from tensorrt_llm._utils import get_sm_version, prefer_pinned +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayHistory +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.seeds import ( + _reset_mamba_seed_buffer, + allocate_mamba_seed_buffer, +) +from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -MIN_REPLAY_HISTORY_SIZE = 16 - -_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 -_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 -_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 -_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D -_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 -_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 - - -@dataclass(frozen=True, kw_only=True) -class ReplayLayerCache(MambaLayerCache): - """Shared Mamba2/GDN replay views consumed by their recurrent mixers.""" - - intermediate_conv_window: torch.Tensor | None = None - intermediate_ssm: torch.Tensor | None = None - mamba_ssm_rand_seed: torch.Tensor | None = field( - default=None, - metadata={"slot_shared": True}, - ) - prev_num_accepted_tokens: torch.Tensor | None = field( - default=None, - metadata={"slot_shared": True}, - ) - cache_buf_idx: torch.Tensor | None = field( - default=None, - metadata={"slot_shared": True}, - ) - old_x: torch.Tensor | None = None - old_B: torch.Tensor | None = None - old_dt: torch.Tensor | None = None - old_dA_cumsum: torch.Tensor | None = None - - -def _splitmix64(value: int) -> int: - """Return the SplitMix64 finalizer for ``value``.""" - value = (value + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK - value ^= value >> 30 - value = (value * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK - value ^= value >> 27 - value = (value * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK - value ^= value >> 31 - return value & _MAMBA_SSM_UINT64_MASK - - -def compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: - """Return a reproducible, nonzero int64 Philox seed.""" - folded = ( - _MAMBA_SSM_SEED_BASE - + counter * _MAMBA_SSM_SEED_MIX_COUNTER - + slot * _MAMBA_SSM_SEED_MIX_SLOT - + rank_offset * _MAMBA_SSM_SEED_MIX_RANK - ) & _MAMBA_SSM_UINT64_MASK - seed = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK - return seed or 1 - - -def allocate_mamba_seed_buffer( - cache_size: int, rank_offset: int, device: torch.device -) -> torch.Tensor: - """Allocate the graph-stable per-slot Philox seed buffer.""" - seeds = [compute_deterministic_mamba_seed(0, slot, rank_offset) for slot in range(cache_size)] - return torch.tensor(seeds, dtype=torch.int64, device=device) - def mamba_seed_rank_offset(mapping: Mapping) -> int: """Return a distinct deterministic seed offset for this parallel rank.""" return mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + mapping.rank * 1_009 -def advance_replay_state( - replay_metadata: ReplayStateUpdateMetadata, - state_indices: torch.Tensor, - accepted_tokens: torch.Tensor, - is_dummy_request: torch.Tensor | None = None, -) -> None: - """Advance double-buffered replay bookkeeping after a verify step.""" - slots = state_indices.long() - accepted_tokens = accepted_tokens.to(replay_metadata.prev_num_accepted_tokens.dtype) - previous = replay_metadata.prev_num_accepted_tokens[slots] - wrote_checkpoint = ( - previous + replay_metadata.replay_step_width > replay_metadata.replay_history_size - ) - next_accepted = torch.where(wrote_checkpoint, accepted_tokens, previous + accepted_tokens) - cache_buffer = replay_metadata.cache_buf_idx[slots] - next_cache_buffer = torch.where(wrote_checkpoint, 1 - cache_buffer, cache_buffer) - if is_dummy_request is not None: - next_accepted = torch.where(is_dummy_request, previous, next_accepted) - next_cache_buffer = torch.where(is_dummy_request, cache_buffer, next_cache_buffer) - replay_metadata.prev_num_accepted_tokens[slots] = next_accepted - replay_metadata.cache_buf_idx[slots] = next_cache_buffer - - -def _reset_mamba_seed_buffer( - seed_buffer: torch.Tensor, - slots: torch.Tensor, - host_slots: list[int], - *, - counter: int, - rank_offset: int, -) -> None: - """Refresh selected slots in the graph-stable int64 Philox buffer.""" - seeds = [compute_deterministic_mamba_seed(counter, slot, rank_offset) for slot in host_slots] - seed_buffer[slots] = torch.tensor(seeds, dtype=torch.int64, pin_memory=prefer_pinned()).to( - seed_buffer.device, non_blocking=True - ) - - class Mamba2State(IntermediateState): """Intermediate promotion plus Mamba2 stochastic-rounding seeds.""" @@ -193,177 +90,6 @@ def shutdown(self) -> None: self.rand_seed = None -class ReplayHistory: - """Own compact Mamba2 replay history, seeds, and bookkeeping.""" - - @property - def intermediate_ssm(self) -> torch.Tensor | None: - return None - - def __init__(self, tokens_per_gen_step: int) -> None: - self._ssm_states: tuple[torch.Tensor, ...] = () - self._conv_states: tuple[torch.Tensor, ...] = () - self._stacked_conv: torch.Tensor | None = None - self.intermediate_conv: torch.Tensor | None = None - self.intermediate_indices: torch.Tensor | None = None - self.rand_seed: torch.Tensor | None = None - self._seed_request_counter = 0 - self._seed_rank_offset = 0 - self.replay_step_width = tokens_per_gen_step - self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step) - self.prev_num_accepted_tokens: torch.Tensor | None = None - self.cache_buf_idx: torch.Tensor | None = None - self.old_x: torch.Tensor | None = None - self.old_B: torch.Tensor | None = None - self.old_dt: torch.Tensor | None = None - self.old_dA_cumsum: torch.Tensor | None = None - - @property - def uses_replay(self) -> bool: - return True - - def validate(self, context: MambaStateLayout) -> None: - if context.spec_config is None: - raise ValueError("Mamba2 replay requires speculative decoding") - if context.n_groups_per_rank <= 0 and context.mamba_pp_layers: - raise ValueError("Mamba2 replay requires at least one state group per rank") - - def bind( - self, - context: MambaStateLayout, - ssm_states: list[torch.Tensor], - conv_states: list[torch.Tensor], - ) -> None: - self._ssm_states = tuple(ssm_states) - self._conv_states = tuple(conv_states) - self._stacked_conv = _stack_state_views(conv_states) - self._seed_rank_offset = context.seed_rank_offset - if not context.mamba_pp_layers: - return - - states = ssm_states - cache_size = states[0].shape[0] - device = states[0].device - tokens_per_step = context.spec_config.tokens_per_gen_step - common_intermediate_shape = [ - len(context.mamba_pp_layers), - context.max_batch_size, - tokens_per_step, - ] - self.intermediate_conv = torch.zeros( - common_intermediate_shape + list(context.conv_state_shape), - dtype=context.conv_state_dtype, - device=device, - ) - self.intermediate_indices = torch.arange( - context.max_batch_size, dtype=torch.int32, device=device - ) - self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) - - num_heads, head_dim, d_state = context.ssm_state_shape - common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] - self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) - self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) - self.old_x = torch.zeros( - common_replay_shape + [self.replay_history_size, num_heads, head_dim], - dtype=context.conv_state_dtype, - device=device, - ) - self.old_B = torch.zeros( - common_replay_shape + [self.replay_history_size, context.n_groups_per_rank, d_state], - dtype=context.conv_state_dtype, - device=device, - ) - self.old_dt = torch.zeros( - common_replay_shape + [num_heads, self.replay_history_size], - dtype=torch.float32, - device=device, - ) - self.old_dA_cumsum = torch.zeros_like(self.old_dt) - - def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: - return { - "mamba_ssm_rand_seed": self.rand_seed, - "intermediate_conv_window": self.intermediate_conv[layer_offset], - "prev_num_accepted_tokens": self.prev_num_accepted_tokens, - "cache_buf_idx": self.cache_buf_idx, - "old_x": self.old_x[layer_offset], - "old_B": self.old_B[layer_offset], - "old_dt": self.old_dt[layer_offset], - "old_dA_cumsum": self.old_dA_cumsum[layer_offset], - } - - def make_layer_cache( - self, - layer_offset: int, - conv: torch.Tensor, - temporal: torch.Tensor, - ) -> MambaLayerCache: - return ReplayLayerCache( - conv=conv, - temporal=temporal, - **self._layer_cache_fields(layer_offset), - ) - - def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - if self.prev_num_accepted_tokens is None: - return - self.prev_num_accepted_tokens.index_fill_(0, slots, 0) - self.cache_buf_idx.index_fill_(0, slots, 0) - for buffer in (self.old_x, self.old_B, self.old_dt, self.old_dA_cumsum): - buffer.index_fill_(1, slots, 0) - if self.rand_seed is not None: - self._seed_request_counter += 1 - _reset_mamba_seed_buffer( - self.rand_seed, - slots, - host_slots, - counter=self._seed_request_counter, - rank_offset=self._seed_rank_offset, - ) - - def update(self, batch: MambaAcceptanceBatch) -> None: - metadata = self.get_replay_metadata() - if metadata is None: - raise RuntimeError("Mamba2 replay buffers are not bound") - # Import through the facade to retain the established monkeypatch point. - from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager - - mamba_cache_manager._advance_replay_state( - metadata, - batch.destination_state_indices, - batch.num_accepted_tokens, - batch.is_dummy_request, - ) - _promote_intermediate_states( - self._conv_states, self.intermediate_conv, batch, self._stacked_conv - ) - - def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: - if self.prev_num_accepted_tokens is None or self.cache_buf_idx is None: - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=self.prev_num_accepted_tokens, - cache_buf_idx=self.cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size, - ) - - def shutdown(self) -> None: - self._ssm_states = () - self._conv_states = () - self._stacked_conv = None - self.intermediate_conv = None - self.intermediate_indices = None - self.rand_seed = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - - def create_mamba2_state( *, spec_config: object | None, use_replay: bool ) -> Mamba2State | ReplayHistory: @@ -464,15 +190,9 @@ def get_nemotron_cache_params(config, *, spec_config=None, quant_config=None): __all__ = [ "Mamba2State", - "ReplayHistory", - "ReplayLayerCache", "NemotronHybridCacheManagerV2", "create_mamba2_state", "select_mamba2_state", "get_nemotron_cache_params", - "allocate_mamba_seed_buffer", - "compute_deterministic_mamba_seed", - "advance_replay_state", "mamba_seed_rank_offset", - "MIN_REPLAY_HISTORY_SIZE", ] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py index 4e9ca3ded1ba..a6f79a091a0e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -203,8 +203,8 @@ def _advance_replay_state( accepted_tokens: torch.Tensor, is_dummy_request: torch.Tensor | None = None, ) -> None: - """Compatibility wrapper for Mamba2-owned replay bookkeeping.""" - from tensorrt_llm._torch.modules.mamba.cache_manager import advance_replay_state + """Compatibility wrapper for shared replay bookkeeping.""" + from .replay import advance_replay_state advance_replay_state(replay_metadata, state_indices, accepted_tokens, is_dummy_request) @@ -212,15 +212,15 @@ def _advance_replay_state( def _allocate_mamba_seed_buffer( cache_size: int, rank_offset: int, device: torch.device ) -> torch.Tensor: - """Compatibility wrapper for Mamba2-owned seed allocation.""" - from tensorrt_llm._torch.modules.mamba.cache_manager import allocate_mamba_seed_buffer + """Compatibility wrapper for shared per-slot seed allocation.""" + from .seeds import allocate_mamba_seed_buffer return allocate_mamba_seed_buffer(cache_size, rank_offset, device) def _compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: - """Compatibility wrapper for Mamba2-owned deterministic seeding.""" - from tensorrt_llm._torch.modules.mamba.cache_manager import compute_deterministic_mamba_seed + """Compatibility wrapper for shared deterministic seeding.""" + from .seeds import compute_deterministic_mamba_seed return compute_deterministic_mamba_seed(counter, slot, rank_offset) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py new file mode 100644 index 000000000000..2d857c89e26c --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py @@ -0,0 +1,246 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared double-buffered replay storage for recurrent cache algorithms. + +Model modules choose this storage and supply any model-specific checkpoint commit. +This module does not import model implementations or dispatch their kernels. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from .common import ( + MIN_REPLAY_HISTORY_SIZE, + MambaAcceptanceBatch, + MambaLayerCache, + MambaStateLayout, + ReplayStateUpdateMetadata, + _promote_intermediate_states, + _stack_state_views, +) +from .seeds import _reset_mamba_seed_buffer, allocate_mamba_seed_buffer + + +@dataclass(frozen=True, kw_only=True) +class ReplayLayerCache(MambaLayerCache): + """Shared Mamba2/GDN replay views consumed by their recurrent mixers.""" + + intermediate_conv_window: torch.Tensor | None = None + intermediate_ssm: torch.Tensor | None = None + mamba_ssm_rand_seed: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + prev_num_accepted_tokens: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + cache_buf_idx: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + old_x: torch.Tensor | None = None + old_B: torch.Tensor | None = None + old_dt: torch.Tensor | None = None + old_dA_cumsum: torch.Tensor | None = None + + +def advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: torch.Tensor | None = None, +) -> None: + """Advance double-buffered replay bookkeeping after a verify step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to(replay_metadata.prev_num_accepted_tokens.dtype) + previous = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = ( + previous + replay_metadata.replay_step_width > replay_metadata.replay_history_size + ) + next_accepted = torch.where(wrote_checkpoint, accepted_tokens, previous + accepted_tokens) + cache_buffer = replay_metadata.cache_buf_idx[slots] + next_cache_buffer = torch.where(wrote_checkpoint, 1 - cache_buffer, cache_buffer) + if is_dummy_request is not None: + next_accepted = torch.where(is_dummy_request, previous, next_accepted) + next_cache_buffer = torch.where(is_dummy_request, cache_buffer, next_cache_buffer) + replay_metadata.prev_num_accepted_tokens[slots] = next_accepted + replay_metadata.cache_buf_idx[slots] = next_cache_buffer + + +class ReplayHistory: + """Own shared replay history, seeds, bookkeeping, and conv promotion.""" + + @property + def intermediate_ssm(self) -> torch.Tensor | None: + return None + + def __init__(self, tokens_per_gen_step: int) -> None: + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self._stacked_conv: torch.Tensor | None = None + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + self.rand_seed: torch.Tensor | None = None + self._seed_request_counter = 0 + self._seed_rank_offset = 0 + self.replay_step_width = tokens_per_gen_step + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step) + self.prev_num_accepted_tokens: torch.Tensor | None = None + self.cache_buf_idx: torch.Tensor | None = None + self.old_x: torch.Tensor | None = None + self.old_B: torch.Tensor | None = None + self.old_dt: torch.Tensor | None = None + self.old_dA_cumsum: torch.Tensor | None = None + + @property + def uses_replay(self) -> bool: + return True + + def validate(self, context: MambaStateLayout) -> None: + if context.spec_config is None: + raise ValueError("Replay requires speculative decoding") + if context.n_groups_per_rank <= 0 and context.mamba_pp_layers: + raise ValueError("Replay requires at least one state group per rank") + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + self._stacked_conv = _stack_state_views(conv_states) + self._seed_rank_offset = context.seed_rank_offset + if not context.mamba_pp_layers: + return + + states = ssm_states + cache_size = states[0].shape[0] + device = states[0].device + tokens_per_step = context.spec_config.tokens_per_gen_step + common_intermediate_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_conv = torch.zeros( + common_intermediate_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) + + num_heads, head_dim, d_state = context.ssm_state_shape + common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.old_x = torch.zeros( + common_replay_shape + [self.replay_history_size, num_heads, head_dim], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_replay_shape + [self.replay_history_size, context.n_groups_per_rank, d_state], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_replay_shape + [num_heads, self.replay_history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros_like(self.old_dt) + + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + return { + "mamba_ssm_rand_seed": self.rand_seed, + "intermediate_conv_window": self.intermediate_conv[layer_offset], + "prev_num_accepted_tokens": self.prev_num_accepted_tokens, + "cache_buf_idx": self.cache_buf_idx, + "old_x": self.old_x[layer_offset], + "old_B": self.old_B[layer_offset], + "old_dt": self.old_dt[layer_offset], + "old_dA_cumsum": self.old_dA_cumsum[layer_offset], + } + + def make_layer_cache( + self, + layer_offset: int, + conv: torch.Tensor, + temporal: torch.Tensor, + ) -> MambaLayerCache: + return ReplayLayerCache( + conv=conv, + temporal=temporal, + **self._layer_cache_fields(layer_offset), + ) + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + if self.prev_num_accepted_tokens is None: + return + self.prev_num_accepted_tokens.index_fill_(0, slots, 0) + self.cache_buf_idx.index_fill_(0, slots, 0) + for buffer in (self.old_x, self.old_B, self.old_dt, self.old_dA_cumsum): + buffer.index_fill_(1, slots, 0) + if self.rand_seed is not None: + self._seed_request_counter += 1 + _reset_mamba_seed_buffer( + self.rand_seed, + slots, + host_slots, + counter=self._seed_request_counter, + rank_offset=self._seed_rank_offset, + ) + + def update(self, batch: MambaAcceptanceBatch) -> None: + metadata = self.get_replay_metadata() + if metadata is None: + raise RuntimeError("Replay buffers are not bound") + # Import through the facade to retain the established monkeypatch point. + from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager + + mamba_cache_manager._advance_replay_state( + metadata, + batch.destination_state_indices, + batch.num_accepted_tokens, + batch.is_dummy_request, + ) + _promote_intermediate_states( + self._conv_states, self.intermediate_conv, batch, self._stacked_conv + ) + + def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: + if self.prev_num_accepted_tokens is None or self.cache_buf_idx is None: + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size, + ) + + def shutdown(self) -> None: + self._ssm_states = () + self._conv_states = () + self._stacked_conv = None + self.intermediate_conv = None + self.intermediate_indices = None + self.rand_seed = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + +__all__ = ["ReplayHistory", "ReplayLayerCache", "advance_replay_state"] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py new file mode 100644 index 000000000000..c92ddc5d747d --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Deterministic per-slot seeds shared by recurrent cache algorithms.""" + +from __future__ import annotations + +import torch + +from tensorrt_llm._utils import prefer_pinned + +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +def _splitmix64(value: int) -> int: + """Return the SplitMix64 finalizer for ``value``.""" + value = (value + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 30 + value = (value * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 27 + value = (value * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 31 + return value & _MAMBA_SSM_UINT64_MASK + + +def compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: + """Return a reproducible, nonzero int64 Philox seed.""" + folded = ( + _MAMBA_SSM_SEED_BASE + + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK + ) & _MAMBA_SSM_UINT64_MASK + seed = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + return seed or 1 + + +def allocate_mamba_seed_buffer( + cache_size: int, rank_offset: int, device: torch.device +) -> torch.Tensor: + """Allocate the graph-stable per-slot Philox seed buffer.""" + seeds = [compute_deterministic_mamba_seed(0, slot, rank_offset) for slot in range(cache_size)] + return torch.tensor(seeds, dtype=torch.int64, device=device) + + +def _reset_mamba_seed_buffer( + seed_buffer: torch.Tensor, + slots: torch.Tensor, + host_slots: list[int], + *, + counter: int, + rank_offset: int, +) -> None: + """Refresh selected slots in the graph-stable int64 Philox buffer.""" + seeds = [compute_deterministic_mamba_seed(counter, slot, rank_offset) for slot in host_slots] + seed_buffer[slots] = torch.tensor(seeds, dtype=torch.int64, pin_memory=prefer_pinned()).to( + seed_buffer.device, non_blocking=True + ) + + +__all__ = ["allocate_mamba_seed_buffer", "compute_deterministic_mamba_seed"] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 6f2c967cb35a..07ddddc7a2bf 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -22,11 +22,7 @@ KDAReplayState, KimiK3HybridCacheManagerV2, ) -from tensorrt_llm._torch.modules.mamba.cache_manager import ( - IntermediateState, - NemotronHybridCacheManagerV2, - ReplayHistory, -) +from tensorrt_llm._torch.modules.mamba.cache_manager import NemotronHybridCacheManagerV2 from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( KvCacheCreator, @@ -58,7 +54,11 @@ _mamba_snapshot_rule_counts, _promote_mamba_state_triton, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import _stack_state_views +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + IntermediateState, + _stack_state_views, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayHistory from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index 4e008f1e1c38..4ab8e1fab1a8 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -18,8 +18,6 @@ from tensorrt_llm._torch.modules.mamba.cache_manager import ( Mamba2State, NemotronHybridCacheManagerV2, - ReplayHistory, - ReplayLayerCache, ) from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( PLE_CONV_STATE, @@ -37,6 +35,10 @@ MambaRole, MambaStateLayout, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ( + ReplayHistory, + ReplayLayerCache, +) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, @@ -754,6 +756,8 @@ def test_mamba2_seed_lifecycle_does_not_require_speculative_decoding(): def test_replay_and_intermediate_are_independent_algorithms(): assert ReplayHistory.__bases__ == (object,) assert GDNReplayState.__bases__ == (ReplayHistory,) + assert ReplayHistory.__module__.endswith("mamba_cache_manager.replay") + assert not any("modules.mamba" in cls.__module__ for cls in GDNReplayState.__mro__) assert Mamba2State.__bases__ == (IntermediateState,) assert not issubclass(ReplayHistory, IntermediateState) assert not issubclass(GDNReplayState, Mamba2State) From eaf26e4d556bae6deac376509532f3fb5ec33481 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:32:07 +0000 Subject: [PATCH 09/25] [None][refactor] Keep replay implementations in their model modules (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/modules/fla/cache_manager.py | 183 +++++++++++++++-- .../_torch/modules/mamba/cache_manager.py | 180 ++++++++++++++++- .../kv_cache/mamba_cache_manager/replay.py | 189 +----------------- .../kv_cache/test_mamba_cache_manager.py | 6 +- .../kv_cache/test_mamba_model_managers.py | 20 +- 5 files changed, 366 insertions(+), 212 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index 90ab9de62076..6367ea92bee8 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -15,15 +15,24 @@ from typing_extensions import override from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + MIN_REPLAY_HISTORY_SIZE, IntermediateState, MambaAcceptanceBatch, + MambaLayerCache, MambaStateLayout, + ReplayStateUpdateMetadata, _mamba_effective_tp_size, + _promote_intermediate_states, + _stack_state_views, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayHistory +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayLayerCache +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.seeds import ( + _reset_mamba_seed_buffer, + allocate_mamba_seed_buffer, +) from tensorrt_llm._torch.utils import is_gdn_replay_enabled from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger @@ -34,23 +43,99 @@ ) -class GDNReplayState(ReplayHistory): - """Own GDN all-layer cached-replay checkpoint commits.""" +class GDNReplayState: + """Own GDN replay buffers, bookkeeping, and checkpoint commits.""" - @override def __init__(self, tokens_per_gen_step: int) -> None: - super().__init__(tokens_per_gen_step) + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self._stacked_conv: torch.Tensor | None = None + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + self.rand_seed: torch.Tensor | None = None + self._seed_request_counter = 0 + self._seed_rank_offset = 0 + self.replay_step_width = tokens_per_gen_step + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step) + self.prev_num_accepted_tokens: torch.Tensor | None = None + self.cache_buf_idx: torch.Tensor | None = None + self.old_x: torch.Tensor | None = None + self.old_B: torch.Tensor | None = None + self.old_dt: torch.Tensor | None = None + self.old_dA_cumsum: torch.Tensor | None = None self._state_descriptors: torch.Tensor | None = None self._state_strides: tuple[int, int, int] | None = None - @override + @property + def intermediate_ssm(self) -> torch.Tensor | None: + return None + + @property + def uses_replay(self) -> bool: + return True + + @property + def has_bound_states(self) -> bool: + return self._state_descriptors is not None or self._state_strides is not None + + def validate(self, context: MambaStateLayout) -> None: + if context.spec_config is None: + raise ValueError("GDN replay requires speculative decoding") + if context.n_groups_per_rank <= 0 and context.mamba_pp_layers: + raise ValueError("GDN replay requires at least one state group per rank") + def bind( self, context: MambaStateLayout, ssm_states: list[torch.Tensor], conv_states: list[torch.Tensor], ) -> None: - super().bind(context, ssm_states, conv_states) + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + self._stacked_conv = _stack_state_views(conv_states) + self._seed_rank_offset = context.seed_rank_offset + if not context.mamba_pp_layers: + return + + states = ssm_states + cache_size = states[0].shape[0] + device = states[0].device + tokens_per_step = context.spec_config.tokens_per_gen_step + common_intermediate_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_conv = torch.zeros( + common_intermediate_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) + + num_heads, head_dim, d_state = context.ssm_state_shape + common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.old_x = torch.zeros( + common_replay_shape + [self.replay_history_size, num_heads, head_dim], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_replay_shape + [self.replay_history_size, context.n_groups_per_rank, d_state], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_replay_shape + [num_heads, self.replay_history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros_like(self.old_dt) states = self._ssm_states if not states: return @@ -106,16 +191,66 @@ def bind( key="gdn_cached_replay_v2_commit_mode_fused", ) - @property - def has_bound_states(self) -> bool: - return self._state_descriptors is not None or self._state_strides is not None + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + return { + "mamba_ssm_rand_seed": self.rand_seed, + "intermediate_conv_window": self.intermediate_conv[layer_offset], + "prev_num_accepted_tokens": self.prev_num_accepted_tokens, + "cache_buf_idx": self.cache_buf_idx, + "old_x": self.old_x[layer_offset], + "old_B": self.old_B[layer_offset], + "old_dt": self.old_dt[layer_offset], + "old_dA_cumsum": self.old_dA_cumsum[layer_offset], + } + + def make_layer_cache( + self, + layer_offset: int, + conv: torch.Tensor, + temporal: torch.Tensor, + ) -> MambaLayerCache: + return ReplayLayerCache( + conv=conv, + temporal=temporal, + **self._layer_cache_fields(layer_offset), + ) + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + if self.prev_num_accepted_tokens is None: + return + self.prev_num_accepted_tokens.index_fill_(0, slots, 0) + self.cache_buf_idx.index_fill_(0, slots, 0) + for buffer in (self.old_x, self.old_B, self.old_dt, self.old_dA_cumsum): + buffer.index_fill_(1, slots, 0) + if self.rand_seed is not None: + self._seed_request_counter += 1 + _reset_mamba_seed_buffer( + self.rand_seed, + slots, + host_slots, + counter=self._seed_request_counter, + rank_offset=self._seed_rank_offset, + ) - @override def update(self, batch: MambaAcceptanceBatch) -> None: num_decodes = batch.num_generations if num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: self.commit_all_layers(batch.attention_metadata, num_decodes) - super().update(batch) + metadata = self.get_replay_metadata() + if metadata is None: + raise RuntimeError("GDN replay buffers are not bound") + # Import through the facade to retain the established monkeypatch point. + from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager + + mamba_cache_manager._advance_replay_state( + metadata, + batch.destination_state_indices, + batch.num_accepted_tokens, + batch.is_dummy_request, + ) + _promote_intermediate_states( + self._conv_states, self.intermediate_conv, batch, self._stacked_conv + ) def commit_all_layers(self, attention_metadata: object, num_decodes: int) -> None: """Commit every local GDN checkpoint using one partitioned launch.""" @@ -153,9 +288,29 @@ def commit_all_layers(self, attention_metadata: object, num_decodes: int) -> Non history_size=self.replay_history_size, ) - @override + def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: + if self.prev_num_accepted_tokens is None or self.cache_buf_idx is None: + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size, + ) + def shutdown(self) -> None: - super().shutdown() + self._ssm_states = () + self._conv_states = () + self._stacked_conv = None + self.intermediate_conv = None + self.intermediate_indices = None + self.rand_seed = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None self._state_descriptors = None self._state_strides = None diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py index bd8a7451ad77..7a4553287033 100644 --- a/tensorrt_llm/_torch/modules/mamba/cache_manager.py +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -16,13 +16,19 @@ from typing_extensions import override from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( + MIN_REPLAY_HISTORY_SIZE, IntermediateState, + MambaAcceptanceBatch, + MambaLayerCache, MambaStateLayout, + ReplayStateUpdateMetadata, + _promote_intermediate_states, + _stack_state_views, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayHistory +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayLayerCache from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.seeds import ( _reset_mamba_seed_buffer, allocate_mamba_seed_buffer, @@ -90,6 +96,177 @@ def shutdown(self) -> None: self.rand_seed = None +class ReplayHistory: + """Own Mamba2 replay history, seeds, bookkeeping, and conv promotion.""" + + @property + def intermediate_ssm(self) -> torch.Tensor | None: + return None + + def __init__(self, tokens_per_gen_step: int) -> None: + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self._stacked_conv: torch.Tensor | None = None + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + self.rand_seed: torch.Tensor | None = None + self._seed_request_counter = 0 + self._seed_rank_offset = 0 + self.replay_step_width = tokens_per_gen_step + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step) + self.prev_num_accepted_tokens: torch.Tensor | None = None + self.cache_buf_idx: torch.Tensor | None = None + self.old_x: torch.Tensor | None = None + self.old_B: torch.Tensor | None = None + self.old_dt: torch.Tensor | None = None + self.old_dA_cumsum: torch.Tensor | None = None + + @property + def uses_replay(self) -> bool: + return True + + def validate(self, context: MambaStateLayout) -> None: + if context.spec_config is None: + raise ValueError("Mamba2 replay requires speculative decoding") + if context.n_groups_per_rank <= 0 and context.mamba_pp_layers: + raise ValueError("Mamba2 replay requires at least one state group per rank") + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + self._stacked_conv = _stack_state_views(conv_states) + self._seed_rank_offset = context.seed_rank_offset + if not context.mamba_pp_layers: + return + + states = ssm_states + cache_size = states[0].shape[0] + device = states[0].device + tokens_per_step = context.spec_config.tokens_per_gen_step + common_intermediate_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_conv = torch.zeros( + common_intermediate_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) + + num_heads, head_dim, d_state = context.ssm_state_shape + common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] + self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) + self.old_x = torch.zeros( + common_replay_shape + [self.replay_history_size, num_heads, head_dim], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_B = torch.zeros( + common_replay_shape + [self.replay_history_size, context.n_groups_per_rank, d_state], + dtype=context.conv_state_dtype, + device=device, + ) + self.old_dt = torch.zeros( + common_replay_shape + [num_heads, self.replay_history_size], + dtype=torch.float32, + device=device, + ) + self.old_dA_cumsum = torch.zeros_like(self.old_dt) + + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + return { + "mamba_ssm_rand_seed": self.rand_seed, + "intermediate_conv_window": self.intermediate_conv[layer_offset], + "prev_num_accepted_tokens": self.prev_num_accepted_tokens, + "cache_buf_idx": self.cache_buf_idx, + "old_x": self.old_x[layer_offset], + "old_B": self.old_B[layer_offset], + "old_dt": self.old_dt[layer_offset], + "old_dA_cumsum": self.old_dA_cumsum[layer_offset], + } + + def make_layer_cache( + self, + layer_offset: int, + conv: torch.Tensor, + temporal: torch.Tensor, + ) -> MambaLayerCache: + return ReplayLayerCache( + conv=conv, + temporal=temporal, + **self._layer_cache_fields(layer_offset), + ) + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + if self.prev_num_accepted_tokens is None: + return + self.prev_num_accepted_tokens.index_fill_(0, slots, 0) + self.cache_buf_idx.index_fill_(0, slots, 0) + for buffer in (self.old_x, self.old_B, self.old_dt, self.old_dA_cumsum): + buffer.index_fill_(1, slots, 0) + if self.rand_seed is not None: + self._seed_request_counter += 1 + _reset_mamba_seed_buffer( + self.rand_seed, + slots, + host_slots, + counter=self._seed_request_counter, + rank_offset=self._seed_rank_offset, + ) + + def update(self, batch: MambaAcceptanceBatch) -> None: + metadata = self.get_replay_metadata() + if metadata is None: + raise RuntimeError("Mamba2 replay buffers are not bound") + # Import through the facade to retain the established monkeypatch point. + from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager + + mamba_cache_manager._advance_replay_state( + metadata, + batch.destination_state_indices, + batch.num_accepted_tokens, + batch.is_dummy_request, + ) + _promote_intermediate_states( + self._conv_states, self.intermediate_conv, batch, self._stacked_conv + ) + + def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: + if self.prev_num_accepted_tokens is None or self.cache_buf_idx is None: + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size, + ) + + def shutdown(self) -> None: + self._ssm_states = () + self._conv_states = () + self._stacked_conv = None + self.intermediate_conv = None + self.intermediate_indices = None + self.rand_seed = None + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + + def create_mamba2_state( *, spec_config: object | None, use_replay: bool ) -> Mamba2State | ReplayHistory: @@ -190,6 +367,7 @@ def get_nemotron_cache_params(config, *, spec_config=None, quant_config=None): __all__ = [ "Mamba2State", + "ReplayHistory", "NemotronHybridCacheManagerV2", "create_mamba2_state", "select_mamba2_state", diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py index 2d857c89e26c..047b2443171e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py @@ -1,10 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Shared double-buffered replay storage for recurrent cache algorithms. +"""Replay payload contract and double-buffer bookkeeping helpers. -Model modules choose this storage and supply any model-specific checkpoint commit. -This module does not import model implementations or dispatch their kernels. +Buffer ownership, allocation, reset, and update policies belong to model modules. """ from __future__ import annotations @@ -13,16 +12,7 @@ import torch -from .common import ( - MIN_REPLAY_HISTORY_SIZE, - MambaAcceptanceBatch, - MambaLayerCache, - MambaStateLayout, - ReplayStateUpdateMetadata, - _promote_intermediate_states, - _stack_state_views, -) -from .seeds import _reset_mamba_seed_buffer, allocate_mamba_seed_buffer +from .common import MambaLayerCache, ReplayStateUpdateMetadata @dataclass(frozen=True, kw_only=True) @@ -72,175 +62,4 @@ def advance_replay_state( replay_metadata.cache_buf_idx[slots] = next_cache_buffer -class ReplayHistory: - """Own shared replay history, seeds, bookkeeping, and conv promotion.""" - - @property - def intermediate_ssm(self) -> torch.Tensor | None: - return None - - def __init__(self, tokens_per_gen_step: int) -> None: - self._ssm_states: tuple[torch.Tensor, ...] = () - self._conv_states: tuple[torch.Tensor, ...] = () - self._stacked_conv: torch.Tensor | None = None - self.intermediate_conv: torch.Tensor | None = None - self.intermediate_indices: torch.Tensor | None = None - self.rand_seed: torch.Tensor | None = None - self._seed_request_counter = 0 - self._seed_rank_offset = 0 - self.replay_step_width = tokens_per_gen_step - self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, tokens_per_gen_step) - self.prev_num_accepted_tokens: torch.Tensor | None = None - self.cache_buf_idx: torch.Tensor | None = None - self.old_x: torch.Tensor | None = None - self.old_B: torch.Tensor | None = None - self.old_dt: torch.Tensor | None = None - self.old_dA_cumsum: torch.Tensor | None = None - - @property - def uses_replay(self) -> bool: - return True - - def validate(self, context: MambaStateLayout) -> None: - if context.spec_config is None: - raise ValueError("Replay requires speculative decoding") - if context.n_groups_per_rank <= 0 and context.mamba_pp_layers: - raise ValueError("Replay requires at least one state group per rank") - - def bind( - self, - context: MambaStateLayout, - ssm_states: list[torch.Tensor], - conv_states: list[torch.Tensor], - ) -> None: - self._ssm_states = tuple(ssm_states) - self._conv_states = tuple(conv_states) - self._stacked_conv = _stack_state_views(conv_states) - self._seed_rank_offset = context.seed_rank_offset - if not context.mamba_pp_layers: - return - - states = ssm_states - cache_size = states[0].shape[0] - device = states[0].device - tokens_per_step = context.spec_config.tokens_per_gen_step - common_intermediate_shape = [ - len(context.mamba_pp_layers), - context.max_batch_size, - tokens_per_step, - ] - self.intermediate_conv = torch.zeros( - common_intermediate_shape + list(context.conv_state_shape), - dtype=context.conv_state_dtype, - device=device, - ) - self.intermediate_indices = torch.arange( - context.max_batch_size, dtype=torch.int32, device=device - ) - self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) - - num_heads, head_dim, d_state = context.ssm_state_shape - common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] - self.prev_num_accepted_tokens = torch.zeros(cache_size, dtype=torch.int32, device=device) - self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) - self.old_x = torch.zeros( - common_replay_shape + [self.replay_history_size, num_heads, head_dim], - dtype=context.conv_state_dtype, - device=device, - ) - self.old_B = torch.zeros( - common_replay_shape + [self.replay_history_size, context.n_groups_per_rank, d_state], - dtype=context.conv_state_dtype, - device=device, - ) - self.old_dt = torch.zeros( - common_replay_shape + [num_heads, self.replay_history_size], - dtype=torch.float32, - device=device, - ) - self.old_dA_cumsum = torch.zeros_like(self.old_dt) - - def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: - return { - "mamba_ssm_rand_seed": self.rand_seed, - "intermediate_conv_window": self.intermediate_conv[layer_offset], - "prev_num_accepted_tokens": self.prev_num_accepted_tokens, - "cache_buf_idx": self.cache_buf_idx, - "old_x": self.old_x[layer_offset], - "old_B": self.old_B[layer_offset], - "old_dt": self.old_dt[layer_offset], - "old_dA_cumsum": self.old_dA_cumsum[layer_offset], - } - - def make_layer_cache( - self, - layer_offset: int, - conv: torch.Tensor, - temporal: torch.Tensor, - ) -> MambaLayerCache: - return ReplayLayerCache( - conv=conv, - temporal=temporal, - **self._layer_cache_fields(layer_offset), - ) - - def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - if self.prev_num_accepted_tokens is None: - return - self.prev_num_accepted_tokens.index_fill_(0, slots, 0) - self.cache_buf_idx.index_fill_(0, slots, 0) - for buffer in (self.old_x, self.old_B, self.old_dt, self.old_dA_cumsum): - buffer.index_fill_(1, slots, 0) - if self.rand_seed is not None: - self._seed_request_counter += 1 - _reset_mamba_seed_buffer( - self.rand_seed, - slots, - host_slots, - counter=self._seed_request_counter, - rank_offset=self._seed_rank_offset, - ) - - def update(self, batch: MambaAcceptanceBatch) -> None: - metadata = self.get_replay_metadata() - if metadata is None: - raise RuntimeError("Replay buffers are not bound") - # Import through the facade to retain the established monkeypatch point. - from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager - - mamba_cache_manager._advance_replay_state( - metadata, - batch.destination_state_indices, - batch.num_accepted_tokens, - batch.is_dummy_request, - ) - _promote_intermediate_states( - self._conv_states, self.intermediate_conv, batch, self._stacked_conv - ) - - def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: - if self.prev_num_accepted_tokens is None or self.cache_buf_idx is None: - return None - return ReplayStateUpdateMetadata( - prev_num_accepted_tokens=self.prev_num_accepted_tokens, - cache_buf_idx=self.cache_buf_idx, - replay_step_width=self.replay_step_width, - replay_history_size=self.replay_history_size, - ) - - def shutdown(self) -> None: - self._ssm_states = () - self._conv_states = () - self._stacked_conv = None - self.intermediate_conv = None - self.intermediate_indices = None - self.rand_seed = None - self.prev_num_accepted_tokens = None - self.cache_buf_idx = None - self.old_x = None - self.old_B = None - self.old_dt = None - self.old_dA_cumsum = None - - -__all__ = ["ReplayHistory", "ReplayLayerCache", "advance_replay_state"] +__all__ = ["ReplayLayerCache", "advance_replay_state"] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 07ddddc7a2bf..a0483901524f 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -22,7 +22,10 @@ KDAReplayState, KimiK3HybridCacheManagerV2, ) -from tensorrt_llm._torch.modules.mamba.cache_manager import NemotronHybridCacheManagerV2 +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + NemotronHybridCacheManagerV2, + ReplayHistory, +) from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( KvCacheCreator, @@ -58,7 +61,6 @@ IntermediateState, _stack_state_views, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayHistory from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index 4ab8e1fab1a8..08598440c565 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -18,6 +18,7 @@ from tensorrt_llm._torch.modules.mamba.cache_manager import ( Mamba2State, NemotronHybridCacheManagerV2, + ReplayHistory, ) from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( PLE_CONV_STATE, @@ -35,10 +36,7 @@ MambaRole, MambaStateLayout, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ( - ReplayHistory, - ReplayLayerCache, -) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayLayerCache from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, @@ -376,12 +374,12 @@ def test_common_manager_consumes_model_state_contract(state): assert payload.temporal is manager.all_ssm_states[0] if isinstance(state, KDAReplayState): assert payload.kda_qkg_cache is not None - elif isinstance(state, ReplayHistory): + elif isinstance(state, (ReplayHistory, GDNReplayState)): assert payload.old_x is not None else: assert payload.intermediate_ssm is not None assert (manager.get_replay_state_update_metadata() is not None) == isinstance( - state, ReplayHistory + state, (ReplayHistory, GDNReplayState) ) manager._shutdown_model_state() assert manager.intermediate_state_indices is None @@ -755,8 +753,10 @@ def test_mamba2_seed_lifecycle_does_not_require_speculative_decoding(): def test_replay_and_intermediate_are_independent_algorithms(): assert ReplayHistory.__bases__ == (object,) - assert GDNReplayState.__bases__ == (ReplayHistory,) - assert ReplayHistory.__module__.endswith("mamba_cache_manager.replay") + assert GDNReplayState.__bases__ == (object,) + assert not issubclass(GDNReplayState, ReplayHistory) + assert ReplayHistory.__module__.endswith("modules.mamba.cache_manager") + assert GDNReplayState.__module__.endswith("modules.fla.cache_manager") assert not any("modules.mamba" in cls.__module__ for cls in GDNReplayState.__mro__) assert Mamba2State.__bases__ == (IntermediateState,) assert not issubclass(ReplayHistory, IntermediateState) @@ -863,13 +863,13 @@ def test_algorithm_borrows_only_tensor_views_and_releases_them(state): state.bind(_layout(), views.all_ssm_states, views.all_conv_states) assert state._conv_states[0] is views.all_conv_states[0] assert not hasattr(state, "manager") - if isinstance(state, (IntermediateState, ReplayHistory)): + if isinstance(state, (IntermediateState, ReplayHistory, GDNReplayState)): assert state._ssm_states[0] is views.all_ssm_states[0] assert not hasattr(state, "_publish_compatibility_views") state.shutdown() state.shutdown() assert state._conv_states == () - if isinstance(state, (IntermediateState, ReplayHistory)): + if isinstance(state, (IntermediateState, ReplayHistory, GDNReplayState)): assert state._ssm_states == () assert not any(isinstance(value, torch.Tensor) for value in vars(state).values()) # Releasing borrowed references must not mutate the persistent pool. From de86e498ccb1fa0388a6d01bd455720c2eefeb35 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:56:56 +0000 Subject: [PATCH 10/25] [None][refactor] Flatten model-owned cache states and payloads (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/modules/fla/cache_manager.py | 241 +++++++++++++++-- .../_torch/modules/kimi_kda/cache_manager.py | 170 ++++++++++-- .../_torch/modules/mamba/cache_manager.py | 246 +++++++++++++++--- .../_torch/modules/qwen4_exp/cache_manager.py | 10 +- .../kv_cache/mamba_cache_manager/__init__.py | 22 +- .../kv_cache/mamba_cache_manager/common.py | 161 +----------- .../kv_cache/mamba_cache_manager/legacy.py | 20 +- .../mamba_cache_manager_v2.py | 12 - .../kv_cache/mamba_cache_manager/replay.py | 65 ----- .../kv_cache/mamba_cache_manager/seeds.py | 66 ----- .../kv_cache/test_mamba_cache_manager.py | 18 +- .../kv_cache/test_mamba_model_managers.py | 109 +++++--- .../modules/mamba/test_mamba2_metadata.py | 8 +- .../modules/mamba/test_mamba_ssm_rand_seed.py | 14 +- 14 files changed, 697 insertions(+), 465 deletions(-) delete mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py delete mode 100644 tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index 6367ea92bee8..71046e203684 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -6,6 +6,8 @@ from __future__ import annotations import sys +from dataclasses import dataclass, field +from typing import NamedTuple import torch @@ -15,12 +17,9 @@ from typing_extensions import override from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( - MIN_REPLAY_HISTORY_SIZE, - IntermediateState, MambaAcceptanceBatch, MambaLayerCache, MambaStateLayout, - ReplayStateUpdateMetadata, _mamba_effective_tp_size, _promote_intermediate_states, _stack_state_views, @@ -28,13 +27,8 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayLayerCache -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.seeds import ( - _reset_mamba_seed_buffer, - allocate_mamba_seed_buffer, -) from tensorrt_llm._torch.utils import is_gdn_replay_enabled -from tensorrt_llm._utils import get_sm_version +from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.logger import logger from .cached_replay import ( @@ -42,6 +36,203 @@ commit_gdn_cached_replay_history_layers, ) +MIN_REPLAY_HISTORY_SIZE = 16 + + +class ReplayStateUpdateMetadata(NamedTuple): + """Shared tensors and fixed sizes for replay state updates.""" + + prev_num_accepted_tokens: torch.Tensor + cache_buf_idx: torch.Tensor + replay_step_width: int + replay_history_size: int + + +def _advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: torch.Tensor | None = None, +) -> None: + """Advance double-buffered replay bookkeeping after a verify step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to(replay_metadata.prev_num_accepted_tokens.dtype) + previous = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = ( + previous + replay_metadata.replay_step_width > replay_metadata.replay_history_size + ) + next_accepted = torch.where(wrote_checkpoint, accepted_tokens, previous + accepted_tokens) + cache_buffer = replay_metadata.cache_buf_idx[slots] + next_cache_buffer = torch.where(wrote_checkpoint, 1 - cache_buffer, cache_buffer) + if is_dummy_request is not None: + next_accepted = torch.where(is_dummy_request, previous, next_accepted) + next_cache_buffer = torch.where(is_dummy_request, cache_buffer, next_cache_buffer) + replay_metadata.prev_num_accepted_tokens[slots] = next_accepted + replay_metadata.cache_buf_idx[slots] = next_cache_buffer + + +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +def _splitmix64(value: int) -> int: + """Return the SplitMix64 finalizer for ``value``.""" + value = (value + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 30 + value = (value * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 27 + value = (value * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 31 + return value & _MAMBA_SSM_UINT64_MASK + + +def compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: + """Return a reproducible, nonzero int64 Philox seed.""" + folded = ( + _MAMBA_SSM_SEED_BASE + + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK + ) & _MAMBA_SSM_UINT64_MASK + seed = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + return seed or 1 + + +def allocate_mamba_seed_buffer( + cache_size: int, rank_offset: int, device: torch.device +) -> torch.Tensor: + """Allocate the graph-stable per-slot Philox seed buffer.""" + seeds = [compute_deterministic_mamba_seed(0, slot, rank_offset) for slot in range(cache_size)] + return torch.tensor(seeds, dtype=torch.int64, device=device) + + +def _reset_mamba_seed_buffer( + seed_buffer: torch.Tensor, + slots: torch.Tensor, + host_slots: list[int], + *, + counter: int, + rank_offset: int, +) -> None: + """Refresh selected slots in the graph-stable int64 Philox buffer.""" + seeds = [compute_deterministic_mamba_seed(counter, slot, rank_offset) for slot in host_slots] + seed_buffer[slots] = torch.tensor(seeds, dtype=torch.int64, pin_memory=prefer_pinned()).to( + seed_buffer.device, non_blocking=True + ) + + +@dataclass(frozen=True, kw_only=True) +class GDNLayerCache(MambaLayerCache): + """Per-layer GDN intermediate or replay views consumed by its mixer.""" + + intermediate_conv_window: torch.Tensor | None = None + intermediate_ssm: torch.Tensor | None = None + mamba_ssm_rand_seed: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + prev_num_accepted_tokens: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + cache_buf_idx: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + old_x: torch.Tensor | None = None + old_B: torch.Tensor | None = None + old_dt: torch.Tensor | None = None + old_dA_cumsum: torch.Tensor | None = None + + +class GDNIntermediateState: + """Own conventional speculative intermediate states and promotion.""" + + def __init__(self) -> None: + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self._stacked_ssm: torch.Tensor | None = None + self._stacked_conv: torch.Tensor | None = None + self.intermediate_ssm: torch.Tensor | None = None + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + """Borrow persistent per-layer views until shutdown; allocate owned scratch.""" + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + self._stacked_ssm = _stack_state_views(ssm_states) + self._stacked_conv = _stack_state_views(conv_states) + if not context.mamba_pp_layers: + return + + device = ssm_states[0].device + + if context.spec_config is not None: + tokens_per_step = context.spec_config.tokens_per_gen_step + common_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_ssm = torch.zeros( + common_shape + list(context.ssm_state_shape), + dtype=context.ssm_state_dtype, + device=device, + ) + self.intermediate_conv = torch.zeros( + common_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + fields: dict[str, torch.Tensor | None] = {} + if self.intermediate_conv is not None: + fields["intermediate_conv_window"] = self.intermediate_conv[layer_offset] + if self.intermediate_ssm is not None: + fields["intermediate_ssm"] = self.intermediate_ssm[layer_offset] + return fields + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + """Plain intermediate buffers have no request-persistent bookkeeping.""" + + def update(self, batch: MambaAcceptanceBatch) -> None: + if self.intermediate_indices is None: + return + _promote_intermediate_states( + self._ssm_states, self.intermediate_ssm, batch, self._stacked_ssm + ) + _promote_intermediate_states( + self._conv_states, self.intermediate_conv, batch, self._stacked_conv + ) + + def shutdown(self) -> None: + self._ssm_states = () + self._conv_states = () + self._stacked_ssm = None + self._stacked_conv = None + self.intermediate_ssm = None + self.intermediate_conv = None + self.intermediate_indices = None + + def make_layer_cache( + self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor + ) -> MambaLayerCache: + return GDNLayerCache(conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset)) + class GDNReplayState: """Own GDN replay buffers, bookkeeping, and checkpoint commits.""" @@ -93,7 +284,11 @@ def bind( self._ssm_states = tuple(ssm_states) self._conv_states = tuple(conv_states) self._stacked_conv = _stack_state_views(conv_states) - self._seed_rank_offset = context.seed_rank_offset + self._seed_rank_offset = ( + context.mapping.tp_rank * 1_000_003 + + context.mapping.pp_rank * 1_000_033 + + context.mapping.rank * 1_009 + ) if not context.mamba_pp_layers: return @@ -114,7 +309,7 @@ def bind( self.intermediate_indices = torch.arange( context.max_batch_size, dtype=torch.int32, device=device ) - self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) + self.rand_seed = allocate_mamba_seed_buffer(cache_size, self._seed_rank_offset, device) num_heads, head_dim, d_state = context.ssm_state_shape common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] @@ -209,7 +404,7 @@ def make_layer_cache( conv: torch.Tensor, temporal: torch.Tensor, ) -> MambaLayerCache: - return ReplayLayerCache( + return GDNLayerCache( conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset), @@ -239,10 +434,7 @@ def update(self, batch: MambaAcceptanceBatch) -> None: metadata = self.get_replay_metadata() if metadata is None: raise RuntimeError("GDN replay buffers are not bound") - # Import through the facade to retain the established monkeypatch point. - from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager - - mamba_cache_manager._advance_replay_state( + _advance_replay_state( metadata, batch.destination_state_indices, batch.num_accepted_tokens, @@ -361,10 +553,10 @@ def select_gdn_replay_state( def create_gdn_state( manager: MambaHybridCacheManagerV2, use_replay: bool | None -) -> IntermediateState | GDNReplayState: +) -> GDNIntermediateState | GDNReplayState: """Shared selection for the Qwen3.5 and Qwen4 cache managers.""" if use_replay is False: - return IntermediateState() + return GDNIntermediateState() if use_replay is True: if manager.spec_config is None: raise ValueError("GDN replay requires speculative decoding") @@ -375,12 +567,12 @@ def create_gdn_state( ssm_cache_dtype=manager.ssm_state_dtype, manager_cls=type(manager), ) - or IntermediateState() + or GDNIntermediateState() ) def validate_gdn_layout( - manager: MambaHybridCacheManagerV2, state: IntermediateState | GDNReplayState + manager: MambaHybridCacheManagerV2, state: GDNIntermediateState | GDNReplayState ) -> None: if isinstance(state, GDNReplayState): if manager.local_num_mamba_layers and manager._global_n_groups % _mamba_effective_tp_size( @@ -395,17 +587,22 @@ class Qwen35HybridCacheManagerV2(MambaHybridCacheManagerV2): @override def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) -> None: + kwargs.pop("mamba_ssm_stochastic_rounding", False) self._requested_replay = use_replay_state_update kwargs.setdefault("conv_state_layout", "q_k_v") super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> IntermediateState | GDNReplayState: + def _initialize_model_state(self) -> GDNIntermediateState | GDNReplayState: state = create_gdn_state(self, self._requested_replay) validate_gdn_layout(self, state) return state + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + state = self._speculative_state + return state.get_replay_metadata() if isinstance(state, GDNReplayState) else None + @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: state = self._speculative_state @@ -435,6 +632,8 @@ def get_gdn_cache_params(config, *, spec_config=None, quant_config=None): __all__ = [ + "GDNIntermediateState", + "GDNLayerCache", "GDNReplayState", "Qwen35HybridCacheManagerV2", "create_gdn_state", diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index 1677b936534d..df122dc85de2 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -16,21 +16,55 @@ from typing_extensions import override from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( - IntermediateState, MambaAcceptanceBatch, MambaLayerCache, MambaStateLayout, - ReplayStateUpdateMetadata, + _promote_intermediate_states, + _stack_state_views, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.seeds import ( - allocate_mamba_seed_buffer, - compute_deterministic_mamba_seed, -) from tensorrt_llm.logger import logger +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +def _splitmix64(value: int) -> int: + """Return the SplitMix64 finalizer for ``value``.""" + value = (value + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 30 + value = (value * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 27 + value = (value * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 31 + return value & _MAMBA_SSM_UINT64_MASK + + +def compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: + """Return a reproducible, nonzero int64 Philox seed.""" + folded = ( + _MAMBA_SSM_SEED_BASE + + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK + ) & _MAMBA_SSM_UINT64_MASK + seed = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + return seed or 1 + + +def allocate_mamba_seed_buffer( + cache_size: int, rank_offset: int, device: torch.device +) -> torch.Tensor: + """Allocate the graph-stable per-slot Philox seed buffer.""" + seeds = [compute_deterministic_mamba_seed(0, slot, rank_offset) for slot in range(cache_size)] + return torch.tensor(seeds, dtype=torch.int64, device=device) + @dataclass(frozen=True, kw_only=True) class KDAReplayLayerCache(MambaLayerCache): @@ -145,6 +179,101 @@ def allocate_conv_cache() -> torch.Tensor: return fields, scratch +@dataclass(frozen=True, kw_only=True) +class KDAIntermediateLayerCache(MambaLayerCache): + """Generic intermediate-state view used during speculative decoding.""" + + intermediate_conv_window: torch.Tensor | None = None + intermediate_ssm: torch.Tensor | None = None + + +class KDAIntermediateState: + """Own conventional speculative intermediate states and promotion.""" + + def __init__(self) -> None: + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self._stacked_ssm: torch.Tensor | None = None + self._stacked_conv: torch.Tensor | None = None + self.intermediate_ssm: torch.Tensor | None = None + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + """Borrow persistent per-layer views until shutdown; allocate owned scratch.""" + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + self._stacked_ssm = _stack_state_views(ssm_states) + self._stacked_conv = _stack_state_views(conv_states) + if not context.mamba_pp_layers: + return + + device = ssm_states[0].device + + if context.spec_config is not None: + tokens_per_step = context.spec_config.tokens_per_gen_step + common_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_ssm = torch.zeros( + common_shape + list(context.ssm_state_shape), + dtype=context.ssm_state_dtype, + device=device, + ) + self.intermediate_conv = torch.zeros( + common_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + + def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: + fields: dict[str, torch.Tensor | None] = {} + if self.intermediate_conv is not None: + fields["intermediate_conv_window"] = self.intermediate_conv[layer_offset] + if self.intermediate_ssm is not None: + fields["intermediate_ssm"] = self.intermediate_ssm[layer_offset] + return fields + + def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: + """Plain intermediate buffers have no request-persistent bookkeeping.""" + + def update(self, batch: MambaAcceptanceBatch) -> None: + if self.intermediate_indices is None: + return + _promote_intermediate_states( + self._ssm_states, self.intermediate_ssm, batch, self._stacked_ssm + ) + _promote_intermediate_states( + self._conv_states, self.intermediate_conv, batch, self._stacked_conv + ) + + def shutdown(self) -> None: + self._ssm_states = () + self._conv_states = () + self._stacked_ssm = None + self._stacked_conv = None + self.intermediate_ssm = None + self.intermediate_conv = None + self.intermediate_indices = None + + def make_layer_cache( + self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor + ) -> MambaLayerCache: + return KDAIntermediateLayerCache( + conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset) + ) + + class KDAReplayState: """Own all V2 state used by the KDA fused multi-token verifier.""" @@ -156,13 +285,13 @@ def intermediate_ssm(self) -> torch.Tensor | None: def intermediate_conv(self) -> torch.Tensor | None: return None - def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: - return None - def __init__( self, num_speculative_tokens: int, + *, + stochastic_rounding: bool = False, ) -> None: + self._stochastic_rounding = stochastic_rounding self.num_speculative_tokens = num_speculative_tokens self._conv_states: tuple[torch.Tensor, ...] = () self._committed_window = 0 @@ -214,7 +343,7 @@ def bytes_per_slot(self, context: MambaStateLayout, layer_id: int) -> int: shared_bytes = 0 if layer_id == context.mamba_pp_layers[0]: shared_bytes = 4 - if context.stochastic_rounding: + if self._stochastic_rounding: shared_bytes += 8 return float_elements * 4 + shared_bytes @@ -225,7 +354,11 @@ def bind( conv_states: list[torch.Tensor], ) -> None: """Borrow per-layer conv views [slot, conv_dim, window] and allocate scratch.""" - self._seed_rank_offset = context.seed_rank_offset + self._seed_rank_offset = ( + context.mapping.tp_rank * 1_000_003 + + context.mapping.pp_rank * 1_000_033 + + context.mapping.rank * 1_009 + ) self._conv_section_dims = context.conv_section_dims self._conv_states = tuple(conv_states) self._committed_window = context.conv_state_shape[1] if conv_states else 0 @@ -253,10 +386,10 @@ def bind( dtype=torch.int32, device=device, ) - if context.stochastic_rounding: + if self._stochastic_rounding: self.rand_seed = allocate_mamba_seed_buffer( cache_size, - context.seed_rank_offset, + self._seed_rank_offset, device, ) logger.info(f"Mamba Cache (kda-replay) is allocated for {cache_size} state slots") @@ -476,21 +609,24 @@ def __init__( ) -> None: if use_replay_state_update: raise ValueError("KDA replay and Mamba2 replay are mutually exclusive") + self._mamba_ssm_stochastic_rounding = kwargs.pop("mamba_ssm_stochastic_rounding", False) self._requested_num_spec = kda_replay_num_spec self._kda_replay: KDAReplayState | None = None kwargs.setdefault("conv_state_layout", "q_k_v") super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> IntermediateState | KDAReplayState: + def _initialize_model_state(self) -> KDAIntermediateState | KDAReplayState: num_spec = self._requested_num_spec if num_spec is None: num_spec = get_kda_replay_num_spec(self.spec_config, manager_supports_replay=True) if num_spec is not None: - self._kda_replay = KDAReplayState(num_spec) + self._kda_replay = KDAReplayState( + num_spec, stochastic_rounding=self._mamba_ssm_stochastic_rounding + ) self._kda_replay.validate(self._state_layout) return self._kda_replay - return IntermediateState() + return KDAIntermediateState() @property def use_kda_replay_update(self) -> bool: @@ -656,6 +792,8 @@ def get_kimi_cache_params(config, *, spec_config=None, quant_config=None): __all__ = [ + "KDAIntermediateLayerCache", + "KDAIntermediateState", "KDAReplayState", "KDAReplayLayerCache", "KimiK3HybridCacheManagerV2", diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py index 7a4553287033..ce9dbef95caf 100644 --- a/tensorrt_llm/_torch/modules/mamba/cache_manager.py +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -7,6 +7,8 @@ import os import sys +from dataclasses import dataclass, field +from typing import NamedTuple import torch @@ -16,61 +18,201 @@ from typing_extensions import override from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( - MIN_REPLAY_HISTORY_SIZE, - IntermediateState, MambaAcceptanceBatch, MambaLayerCache, MambaStateLayout, - ReplayStateUpdateMetadata, _promote_intermediate_states, _stack_state_views, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.mamba_cache_manager_v2 import ( MambaHybridCacheManagerV2, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayLayerCache -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.seeds import ( - _reset_mamba_seed_buffer, - allocate_mamba_seed_buffer, -) -from tensorrt_llm._utils import get_sm_version +from tensorrt_llm._utils import get_sm_version, prefer_pinned from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +MIN_REPLAY_HISTORY_SIZE = 16 + + +class ReplayStateUpdateMetadata(NamedTuple): + """Shared tensors and fixed sizes for replay state updates.""" + + prev_num_accepted_tokens: torch.Tensor + cache_buf_idx: torch.Tensor + replay_step_width: int + replay_history_size: int + + +def _advance_replay_state( + replay_metadata: ReplayStateUpdateMetadata, + state_indices: torch.Tensor, + accepted_tokens: torch.Tensor, + is_dummy_request: torch.Tensor | None = None, +) -> None: + """Advance double-buffered replay bookkeeping after a verify step.""" + slots = state_indices.long() + accepted_tokens = accepted_tokens.to(replay_metadata.prev_num_accepted_tokens.dtype) + previous = replay_metadata.prev_num_accepted_tokens[slots] + wrote_checkpoint = ( + previous + replay_metadata.replay_step_width > replay_metadata.replay_history_size + ) + next_accepted = torch.where(wrote_checkpoint, accepted_tokens, previous + accepted_tokens) + cache_buffer = replay_metadata.cache_buf_idx[slots] + next_cache_buffer = torch.where(wrote_checkpoint, 1 - cache_buffer, cache_buffer) + if is_dummy_request is not None: + next_accepted = torch.where(is_dummy_request, previous, next_accepted) + next_cache_buffer = torch.where(is_dummy_request, cache_buffer, next_cache_buffer) + replay_metadata.prev_num_accepted_tokens[slots] = next_accepted + replay_metadata.cache_buf_idx[slots] = next_cache_buffer + + +_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 +_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 +_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 +_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D +_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 +_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 + + +def _splitmix64(value: int) -> int: + """Return the SplitMix64 finalizer for ``value``.""" + value = (value + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 30 + value = (value * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 27 + value = (value * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK + value ^= value >> 31 + return value & _MAMBA_SSM_UINT64_MASK + + +def compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: + """Return a reproducible, nonzero int64 Philox seed.""" + folded = ( + _MAMBA_SSM_SEED_BASE + + counter * _MAMBA_SSM_SEED_MIX_COUNTER + + slot * _MAMBA_SSM_SEED_MIX_SLOT + + rank_offset * _MAMBA_SSM_SEED_MIX_RANK + ) & _MAMBA_SSM_UINT64_MASK + seed = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK + return seed or 1 + + +def allocate_mamba_seed_buffer( + cache_size: int, rank_offset: int, device: torch.device +) -> torch.Tensor: + """Allocate the graph-stable per-slot Philox seed buffer.""" + seeds = [compute_deterministic_mamba_seed(0, slot, rank_offset) for slot in range(cache_size)] + return torch.tensor(seeds, dtype=torch.int64, device=device) + + +def _reset_mamba_seed_buffer( + seed_buffer: torch.Tensor, + slots: torch.Tensor, + host_slots: list[int], + *, + counter: int, + rank_offset: int, +) -> None: + """Refresh selected slots in the graph-stable int64 Philox buffer.""" + seeds = [compute_deterministic_mamba_seed(counter, slot, rank_offset) for slot in host_slots] + seed_buffer[slots] = torch.tensor(seeds, dtype=torch.int64, pin_memory=prefer_pinned()).to( + seed_buffer.device, non_blocking=True + ) + + +@dataclass(frozen=True, kw_only=True) +class Mamba2LayerCache(MambaLayerCache): + """Per-layer Mamba2 intermediate or replay views consumed by its mixer.""" + + intermediate_conv_window: torch.Tensor | None = None + intermediate_ssm: torch.Tensor | None = None + mamba_ssm_rand_seed: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + prev_num_accepted_tokens: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + cache_buf_idx: torch.Tensor | None = field( + default=None, + metadata={"slot_shared": True}, + ) + old_x: torch.Tensor | None = None + old_B: torch.Tensor | None = None + old_dt: torch.Tensor | None = None + old_dA_cumsum: torch.Tensor | None = None + def mamba_seed_rank_offset(mapping: Mapping) -> int: """Return a distinct deterministic seed offset for this parallel rank.""" return mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + mapping.rank * 1_009 -class Mamba2State(IntermediateState): +class Mamba2State: """Intermediate promotion plus Mamba2 stochastic-rounding seeds.""" uses_replay = False - @override - def __init__(self) -> None: - super().__init__() + def __init__(self, *, stochastic_rounding: bool = False) -> None: + self._stochastic_rounding = stochastic_rounding + self._ssm_states: tuple[torch.Tensor, ...] = () + self._conv_states: tuple[torch.Tensor, ...] = () + self._stacked_ssm: torch.Tensor | None = None + self._stacked_conv: torch.Tensor | None = None + self.intermediate_ssm: torch.Tensor | None = None + self.intermediate_conv: torch.Tensor | None = None + self.intermediate_indices: torch.Tensor | None = None self.rand_seed: torch.Tensor | None = None self._seed_request_counter = 0 self._seed_rank_offset = 0 - @override def bind( self, context: MambaStateLayout, ssm_states: list[torch.Tensor], conv_states: list[torch.Tensor], ) -> None: - super().bind(context, ssm_states, conv_states) - self._seed_rank_offset = context.seed_rank_offset - if context.stochastic_rounding and ssm_states: + self._ssm_states = tuple(ssm_states) + self._conv_states = tuple(conv_states) + self._stacked_ssm = _stack_state_views(ssm_states) + self._stacked_conv = _stack_state_views(conv_states) + if not context.mamba_pp_layers: + return + + device = ssm_states[0].device + + if context.spec_config is not None: + tokens_per_step = context.spec_config.tokens_per_gen_step + common_shape = [ + len(context.mamba_pp_layers), + context.max_batch_size, + tokens_per_step, + ] + self.intermediate_ssm = torch.zeros( + common_shape + list(context.ssm_state_shape), + dtype=context.ssm_state_dtype, + device=device, + ) + self.intermediate_conv = torch.zeros( + common_shape + list(context.conv_state_shape), + dtype=context.conv_state_dtype, + device=device, + ) + self.intermediate_indices = torch.arange( + context.max_batch_size, dtype=torch.int32, device=device + ) + self._seed_rank_offset = ( + context.mapping.tp_rank * 1_000_003 + + context.mapping.pp_rank * 1_000_033 + + context.mapping.rank * 1_009 + ) + if self._stochastic_rounding and ssm_states: states = ssm_states[0] self.rand_seed = allocate_mamba_seed_buffer( - states.shape[0], context.seed_rank_offset, states.device + states.shape[0], self._seed_rank_offset, states.device ) - @override def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: if self.rand_seed is None: return @@ -83,18 +225,43 @@ def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: rank_offset=self._seed_rank_offset, ) - @override def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: - fields = super()._layer_cache_fields(layer_offset) + fields: dict[str, torch.Tensor | None] = {} + if self.intermediate_conv is not None: + fields["intermediate_conv_window"] = self.intermediate_conv[layer_offset] + if self.intermediate_ssm is not None: + fields["intermediate_ssm"] = self.intermediate_ssm[layer_offset] if self.rand_seed is not None: fields["mamba_ssm_rand_seed"] = self.rand_seed return fields - @override def shutdown(self) -> None: - super().shutdown() + self._ssm_states = () + self._conv_states = () + self._stacked_ssm = None + self._stacked_conv = None + self.intermediate_ssm = None + self.intermediate_conv = None + self.intermediate_indices = None self.rand_seed = None + def update(self, batch: MambaAcceptanceBatch) -> None: + if self.intermediate_indices is None: + return + _promote_intermediate_states( + self._ssm_states, self.intermediate_ssm, batch, self._stacked_ssm + ) + _promote_intermediate_states( + self._conv_states, self.intermediate_conv, batch, self._stacked_conv + ) + + def make_layer_cache( + self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor + ) -> MambaLayerCache: + return Mamba2LayerCache( + conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset) + ) + class ReplayHistory: """Own Mamba2 replay history, seeds, bookkeeping, and conv promotion.""" @@ -140,7 +307,11 @@ def bind( self._ssm_states = tuple(ssm_states) self._conv_states = tuple(conv_states) self._stacked_conv = _stack_state_views(conv_states) - self._seed_rank_offset = context.seed_rank_offset + self._seed_rank_offset = ( + context.mapping.tp_rank * 1_000_003 + + context.mapping.pp_rank * 1_000_033 + + context.mapping.rank * 1_009 + ) if not context.mamba_pp_layers: return @@ -161,7 +332,7 @@ def bind( self.intermediate_indices = torch.arange( context.max_batch_size, dtype=torch.int32, device=device ) - self.rand_seed = allocate_mamba_seed_buffer(cache_size, context.seed_rank_offset, device) + self.rand_seed = allocate_mamba_seed_buffer(cache_size, self._seed_rank_offset, device) num_heads, head_dim, d_state = context.ssm_state_shape common_replay_shape = [len(context.mamba_pp_layers), cache_size, 2] @@ -202,7 +373,7 @@ def make_layer_cache( conv: torch.Tensor, temporal: torch.Tensor, ) -> MambaLayerCache: - return ReplayLayerCache( + return Mamba2LayerCache( conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset), @@ -229,10 +400,7 @@ def update(self, batch: MambaAcceptanceBatch) -> None: metadata = self.get_replay_metadata() if metadata is None: raise RuntimeError("Mamba2 replay buffers are not bound") - # Import through the facade to retain the established monkeypatch point. - from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager - - mamba_cache_manager._advance_replay_state( + _advance_replay_state( metadata, batch.destination_state_indices, batch.num_accepted_tokens, @@ -268,14 +436,14 @@ def shutdown(self) -> None: def create_mamba2_state( - *, spec_config: object | None, use_replay: bool + *, spec_config: object | None, use_replay: bool, stochastic_rounding: bool = False ) -> Mamba2State | ReplayHistory: """Build Mamba2 intermediate state or compact replay history.""" if use_replay: if spec_config is None: raise ValueError("Mamba replay requires speculative decoding") return ReplayHistory(spec_config.tokens_per_gen_step) - return Mamba2State() + return Mamba2State(stochastic_rounding=stochastic_rounding) def select_mamba2_state( @@ -303,7 +471,9 @@ def select_mamba2_state( use_replay = False else: logger.info("Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1") - return create_mamba2_state(spec_config=spec_config, use_replay=use_replay) + return create_mamba2_state( + spec_config=spec_config, use_replay=use_replay, stochastic_rounding=stochastic_rounding + ) class NemotronHybridCacheManagerV2(MambaHybridCacheManagerV2): @@ -313,6 +483,7 @@ class NemotronHybridCacheManagerV2(MambaHybridCacheManagerV2): @override def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) -> None: + self._mamba_ssm_stochastic_rounding = kwargs.pop("mamba_ssm_stochastic_rounding", False) self._requested_replay = use_replay_state_update kwargs.setdefault("conv_state_layout", "x_b_c") super().__init__(*args, **kwargs) @@ -321,7 +492,9 @@ def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) def _initialize_model_state(self) -> Mamba2State | ReplayHistory: if self._requested_replay is not None: state = create_mamba2_state( - spec_config=self.spec_config, use_replay=self._requested_replay + spec_config=self.spec_config, + use_replay=self._requested_replay, + stochastic_rounding=self._mamba_ssm_stochastic_rounding, ) else: state = select_mamba2_state( @@ -342,6 +515,10 @@ def _initialize_model_state(self) -> Mamba2State | ReplayHistory: state.validate(self._state_layout) return state + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + state = self._speculative_state + return state.get_replay_metadata() if isinstance(state, ReplayHistory) else None + def get_mamba_ssm_rand_seed(self) -> torch.Tensor | None: return self._speculative_state.rand_seed @@ -366,6 +543,7 @@ def get_nemotron_cache_params(config, *, spec_config=None, quant_config=None): __all__ = [ + "Mamba2LayerCache", "Mamba2State", "ReplayHistory", "NemotronHybridCacheManagerV2", diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index 5649d732bbf2..7236c9f58765 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -25,7 +25,9 @@ from tensorrt_llm._torch.attention.backends.sparse.qsa.params import QSASparseParams from tensorrt_llm._torch.disaggregation.resource.page import MapperKind from tensorrt_llm._torch.modules.fla.cache_manager import ( + GDNIntermediateState, GDNReplayState, + ReplayStateUpdateMetadata, create_gdn_state, validate_gdn_layout, ) @@ -33,7 +35,6 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( - IntermediateState, _get_local_mamba_cache_layout, ) from tensorrt_llm._torch.pyexecutor.resource_manager import get_pp_layers @@ -193,6 +194,7 @@ def __init__( qwen4_exp_ple_cache_params: Qwen4ExpPLECacheParams | None = None, **kwargs, ) -> None: + kwargs.pop("mamba_ssm_stochastic_rounding", False) self._requested_replay = use_replay_state_update self._ple_params = qwen4_exp_ple_cache_params self._pretrained_config = pretrained_config @@ -502,7 +504,7 @@ def _ple_buffer_configs(self) -> list[BufferConfig]: ] @override - def _initialize_model_state(self) -> IntermediateState | GDNReplayState: + def _initialize_model_state(self) -> GDNIntermediateState | GDNReplayState: state = create_gdn_state(self, self._requested_replay) validate_gdn_layout(self, state) @@ -559,6 +561,10 @@ def _shutdown_model_state(self) -> None: self._ple_ngram_contexts.clear() super()._shutdown_model_state() + def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: + state = self._speculative_state + return state.get_replay_metadata() if isinstance(state, GDNReplayState) else None + @property def use_gdn_cached_replay_all_layer_commit(self) -> bool: state = self._speculative_state diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py index 2a6a83544775..cd1ebe8bb592 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py @@ -13,14 +13,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from .common import ( - BaseMambaCacheManager, - IntermediateLayerCache, - MambaHybridCacheManager, - MambaLayerCache, - MambaRole, - ReplayStateUpdateMetadata, - ) + from .common import BaseMambaCacheManager, MambaHybridCacheManager, MambaLayerCache, MambaRole from .legacy import ( CppMambaHybridCacheManager, MambaCacheManager, @@ -34,7 +27,6 @@ "_estimate_mamba_hybrid_cache_cost": ("common", "_estimate_mamba_hybrid_cache_cost"), "BaseMambaCacheManager": ("common", "BaseMambaCacheManager"), "CppMambaHybridCacheManager": ("legacy", "CppMambaHybridCacheManager"), - "MIN_REPLAY_HISTORY_SIZE": ("common", "MIN_REPLAY_HISTORY_SIZE"), "MambaCacheManager": ("legacy", "MambaCacheManager"), "MambaHybridCacheManager": ("common", "MambaHybridCacheManager"), "MambaHybridCacheManagerV2": ( @@ -45,14 +37,6 @@ "MambaRole": ("common", "MambaRole"), "MixedMambaHybridCacheManager": ("legacy", "MixedMambaHybridCacheManager"), "PythonMambaCacheManager": ("legacy", "PythonMambaCacheManager"), - "ReplayStateUpdateMetadata": ("common", "ReplayStateUpdateMetadata"), - "IntermediateLayerCache": ("common", "IntermediateLayerCache"), - "_advance_replay_state": ("common", "_advance_replay_state"), - "_allocate_mamba_seed_buffer": ("common", "_allocate_mamba_seed_buffer"), - "_compute_deterministic_mamba_seed": ( - "common", - "_compute_deterministic_mamba_seed", - ), "_get_local_mamba_cache_layout": ("common", "_get_local_mamba_cache_layout"), "_get_mamba_hybrid_pool_size": ("common", "_get_mamba_hybrid_pool_size"), "_get_num_cuda_graph_padding_dummy_slots": ( @@ -60,7 +44,6 @@ "_get_num_cuda_graph_padding_dummy_slots", ), "_mamba_effective_tp_size": ("common", "_mamba_effective_tp_size"), - "_mamba_rank_offset": ("common", "_mamba_rank_offset"), "_mamba_snapshot_rule_counts": ("common", "_mamba_snapshot_rule_counts"), "_promote_mamba_state_triton": ("common", "_promote_mamba_state_triton"), "use_py_mamba_cache_manager": ("common", "use_py_mamba_cache_manager"), @@ -84,7 +67,6 @@ def __dir__() -> list[str]: __all__ = [ "BaseMambaCacheManager", "CppMambaHybridCacheManager", - "MIN_REPLAY_HISTORY_SIZE", "MambaCacheManager", "MambaHybridCacheManager", "MambaHybridCacheManagerV2", @@ -92,7 +74,5 @@ def __dir__() -> list[str]: "MambaRole", "MixedMambaHybridCacheManager", "PythonMambaCacheManager", - "ReplayStateUpdateMetadata", - "IntermediateLayerCache", "use_py_mamba_cache_manager", ] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py index a6f79a091a0e..a521086d6635 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -20,7 +20,7 @@ import sys from abc import ABC, abstractmethod from dataclasses import dataclass, field, fields, replace -from typing import TYPE_CHECKING, List, Literal, NamedTuple, Optional, Protocol, Tuple +from typing import TYPE_CHECKING, List, Literal, Optional, Protocol, Tuple import torch @@ -49,15 +49,6 @@ GB = 1 << 30 -class ReplayStateUpdateMetadata(NamedTuple): - """Shared tensors and fixed sizes for replay state updates.""" - - prev_num_accepted_tokens: torch.Tensor - cache_buf_idx: torch.Tensor - replay_step_width: int - replay_history_size: int - - @dataclass(frozen=True, kw_only=True) class MambaLayerCache: """Persistent recurrent-state views for one local layer.""" @@ -78,18 +69,6 @@ def at_layer_idx(self, layer: int) -> "MambaLayerCache": return type(self)(**values) -@dataclass(frozen=True, kw_only=True) -class IntermediateLayerCache(MambaLayerCache): - """Generic intermediate-state view used during speculative decoding.""" - - intermediate_conv_window: torch.Tensor | None = None - intermediate_ssm: torch.Tensor | None = None - mamba_ssm_rand_seed: torch.Tensor | None = field( - default=None, - metadata={"slot_shared": True}, - ) - - @dataclass(frozen=True, kw_only=True) class MambaStateLayout: """Resolved recurrent geometry; slot capacity is filled after pool allocation.""" @@ -107,8 +86,6 @@ class MambaStateLayout: state_index_capacity: int slot_capacity: int | None spec_config: object | None - stochastic_rounding: bool = False - seed_rank_offset: int = 0 conv_section_dims: tuple[int, ...] = field(default_factory=tuple) conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c" @@ -160,8 +137,6 @@ def make_layer_cache( self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor ) -> MambaLayerCache: ... - def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: ... - class BaseMambaCacheManager(ABC): """Abstract interface for accessing Mamba/recurrent state caches.""" @@ -170,10 +145,6 @@ class BaseMambaCacheManager(ABC): def get_state_indices(self, *args, **kwargs) -> torch.Tensor: """Return slot indices of each request with shape ``[max_batch_size]``.""" - def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: - """Return replay metadata tensors and fixed replay sizes.""" - return None - @abstractmethod def get_conv_states(self, layer_idx: int) -> torch.Tensor: """Return conv states with shape ``[slot_size, conv_dim, d_conv - 1]``.""" @@ -194,42 +165,6 @@ def on_state_transfer_complete(self, request_ids: list[int]) -> None: """Notify a model manager after a disaggregated transfer.""" -MIN_REPLAY_HISTORY_SIZE = 16 - - -def _advance_replay_state( - replay_metadata: ReplayStateUpdateMetadata, - state_indices: torch.Tensor, - accepted_tokens: torch.Tensor, - is_dummy_request: torch.Tensor | None = None, -) -> None: - """Compatibility wrapper for shared replay bookkeeping.""" - from .replay import advance_replay_state - - advance_replay_state(replay_metadata, state_indices, accepted_tokens, is_dummy_request) - - -def _allocate_mamba_seed_buffer( - cache_size: int, rank_offset: int, device: torch.device -) -> torch.Tensor: - """Compatibility wrapper for shared per-slot seed allocation.""" - from .seeds import allocate_mamba_seed_buffer - - return allocate_mamba_seed_buffer(cache_size, rank_offset, device) - - -def _compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: - """Compatibility wrapper for shared deterministic seeding.""" - from .seeds import compute_deterministic_mamba_seed - - return compute_deterministic_mamba_seed(counter, slot, rank_offset) - - -def _mamba_rank_offset(mapping: Mapping) -> int: - """Stable identity for per-rank recurrent state initialization.""" - return mapping.tp_rank * 1_000_003 + mapping.pp_rank * 1_000_033 + mapping.rank * 1_009 - - class _PrefixReuseDiagnostics(Protocol): def _get_num_reusable_tokens_before_hybrid_pruning(self) -> int: ... @@ -717,13 +652,9 @@ def _estimate_mamba_hybrid_cache_cost( "BaseMambaCacheManager", "MambaHybridCacheManager", "MambaLayerCache", - "IntermediateLayerCache", "MambaStateLayout", "MambaAcceptanceBatch", - "IntermediateState", - "ReplayStateUpdateMetadata", "MambaRole", - "MIN_REPLAY_HISTORY_SIZE", "use_py_mamba_cache_manager", "get_tensor_size_bytes", ] @@ -798,93 +729,3 @@ def _promote_intermediate_states( batch.accepted_positions, batch.destination_state_indices, ) - - -class IntermediateState: - """Own conventional speculative intermediate states and promotion.""" - - def __init__(self) -> None: - self._ssm_states: tuple[torch.Tensor, ...] = () - self._conv_states: tuple[torch.Tensor, ...] = () - self._stacked_ssm: torch.Tensor | None = None - self._stacked_conv: torch.Tensor | None = None - self.intermediate_ssm: torch.Tensor | None = None - self.intermediate_conv: torch.Tensor | None = None - self.intermediate_indices: torch.Tensor | None = None - - def bind( - self, - context: MambaStateLayout, - ssm_states: list[torch.Tensor], - conv_states: list[torch.Tensor], - ) -> None: - """Borrow persistent per-layer views until shutdown; allocate owned scratch.""" - self._ssm_states = tuple(ssm_states) - self._conv_states = tuple(conv_states) - self._stacked_ssm = _stack_state_views(ssm_states) - self._stacked_conv = _stack_state_views(conv_states) - if not context.mamba_pp_layers: - return - - device = ssm_states[0].device - - if context.spec_config is not None: - tokens_per_step = context.spec_config.tokens_per_gen_step - common_shape = [ - len(context.mamba_pp_layers), - context.max_batch_size, - tokens_per_step, - ] - self.intermediate_ssm = torch.zeros( - common_shape + list(context.ssm_state_shape), - dtype=context.ssm_state_dtype, - device=device, - ) - self.intermediate_conv = torch.zeros( - common_shape + list(context.conv_state_shape), - dtype=context.conv_state_dtype, - device=device, - ) - self.intermediate_indices = torch.arange( - context.max_batch_size, dtype=torch.int32, device=device - ) - - def _layer_cache_fields(self, layer_offset: int) -> dict[str, torch.Tensor | None]: - fields: dict[str, torch.Tensor | None] = {} - if self.intermediate_conv is not None: - fields["intermediate_conv_window"] = self.intermediate_conv[layer_offset] - if self.intermediate_ssm is not None: - fields["intermediate_ssm"] = self.intermediate_ssm[layer_offset] - return fields - - def reset_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - """Plain intermediate buffers have no request-persistent bookkeeping.""" - - def get_replay_metadata(self) -> ReplayStateUpdateMetadata | None: - return None - - def update(self, batch: MambaAcceptanceBatch) -> None: - if self.intermediate_indices is None: - return - _promote_intermediate_states( - self._ssm_states, self.intermediate_ssm, batch, self._stacked_ssm - ) - _promote_intermediate_states( - self._conv_states, self.intermediate_conv, batch, self._stacked_conv - ) - - def shutdown(self) -> None: - self._ssm_states = () - self._conv_states = () - self._stacked_ssm = None - self._stacked_conv = None - self.intermediate_ssm = None - self.intermediate_conv = None - self.intermediate_indices = None - - def make_layer_cache( - self, layer_offset: int, conv: torch.Tensor, temporal: torch.Tensor - ) -> MambaLayerCache: - return IntermediateLayerCache( - conv=conv, temporal=temporal, **self._layer_cache_fields(layer_offset) - ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py index 31e340571297..51db405ed9a0 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py @@ -26,6 +26,20 @@ from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + MIN_REPLAY_HISTORY_SIZE, + ReplayStateUpdateMetadata, + _advance_replay_state, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + allocate_mamba_seed_buffer as _allocate_mamba_seed_buffer, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + compute_deterministic_mamba_seed as _compute_deterministic_mamba_seed, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + mamba_seed_rank_offset as _mamba_rank_offset, +) from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, @@ -43,17 +57,11 @@ from tensorrt_llm.mapping import Mapping from .common import ( - MIN_REPLAY_HISTORY_SIZE, BaseMambaCacheManager, MambaHybridCacheManager, - ReplayStateUpdateMetadata, - _advance_replay_state, - _allocate_mamba_seed_buffer, - _compute_deterministic_mamba_seed, _get_local_mamba_cache_layout, _get_mamba_hybrid_pool_size, _mamba_effective_tp_size, - _mamba_rank_offset, _promote_mamba_state_triton, get_tensor_size_bytes, ) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 788f99972d8c..99cbedfa0a5f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -79,11 +79,9 @@ MambaRole, MambaState, MambaStateLayout, - ReplayStateUpdateMetadata, _estimate_mamba_hybrid_cache_cost, _get_num_cuda_graph_padding_dummy_slots, _mamba_effective_tp_size, - _mamba_rank_offset, _mamba_regular_snapshot_interval, _mamba_snapshot_rule_counts, _PrefixReuseDiagnostics, @@ -138,7 +136,6 @@ def __init__( layer_mask: Optional[List[bool]] = None, is_estimating_kv_cache: bool = False, is_draft: bool = False, - mamba_ssm_stochastic_rounding: bool = False, conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", **kwargs, ) -> None: @@ -170,8 +167,6 @@ def __init__( self._mamba_layer_mask = list(mamba_layer_mask) self.spec_config = spec_config - self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding - self._seed_rank_offset = _mamba_rank_offset(mapping) self._recurrent_evicted_blocks_total = 0 self._recurrent_onboarded_blocks_total = 0 self._recurrent_dropped_blocks_total = 0 @@ -266,8 +261,6 @@ def __init__( state_index_capacity=(max_batch_size + self._num_reserved_dummy_slots), slot_capacity=None, spec_config=spec_config, - stochastic_rounding=mamba_ssm_stochastic_rounding, - seed_rank_offset=self._seed_rank_offset, conv_section_dims=tuple(self.conv_section_dims), conv_state_layout=conv_state_layout, ) @@ -741,11 +734,6 @@ def intermediate_conv_states(self) -> torch.Tensor | None: state = self._speculative_state return state.intermediate_conv if state is not None else None - @override - def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: - state = self._speculative_state - return state.get_replay_metadata() if state is not None else None - @override def mamba_layer_cache(self, layer_idx: int) -> MambaLayerCache: if self._speculative_state is not None and self.spec_config is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py deleted file mode 100644 index 047b2443171e..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/replay.py +++ /dev/null @@ -1,65 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Replay payload contract and double-buffer bookkeeping helpers. - -Buffer ownership, allocation, reset, and update policies belong to model modules. -""" - -from __future__ import annotations - -from dataclasses import dataclass, field - -import torch - -from .common import MambaLayerCache, ReplayStateUpdateMetadata - - -@dataclass(frozen=True, kw_only=True) -class ReplayLayerCache(MambaLayerCache): - """Shared Mamba2/GDN replay views consumed by their recurrent mixers.""" - - intermediate_conv_window: torch.Tensor | None = None - intermediate_ssm: torch.Tensor | None = None - mamba_ssm_rand_seed: torch.Tensor | None = field( - default=None, - metadata={"slot_shared": True}, - ) - prev_num_accepted_tokens: torch.Tensor | None = field( - default=None, - metadata={"slot_shared": True}, - ) - cache_buf_idx: torch.Tensor | None = field( - default=None, - metadata={"slot_shared": True}, - ) - old_x: torch.Tensor | None = None - old_B: torch.Tensor | None = None - old_dt: torch.Tensor | None = None - old_dA_cumsum: torch.Tensor | None = None - - -def advance_replay_state( - replay_metadata: ReplayStateUpdateMetadata, - state_indices: torch.Tensor, - accepted_tokens: torch.Tensor, - is_dummy_request: torch.Tensor | None = None, -) -> None: - """Advance double-buffered replay bookkeeping after a verify step.""" - slots = state_indices.long() - accepted_tokens = accepted_tokens.to(replay_metadata.prev_num_accepted_tokens.dtype) - previous = replay_metadata.prev_num_accepted_tokens[slots] - wrote_checkpoint = ( - previous + replay_metadata.replay_step_width > replay_metadata.replay_history_size - ) - next_accepted = torch.where(wrote_checkpoint, accepted_tokens, previous + accepted_tokens) - cache_buffer = replay_metadata.cache_buf_idx[slots] - next_cache_buffer = torch.where(wrote_checkpoint, 1 - cache_buffer, cache_buffer) - if is_dummy_request is not None: - next_accepted = torch.where(is_dummy_request, previous, next_accepted) - next_cache_buffer = torch.where(is_dummy_request, cache_buffer, next_cache_buffer) - replay_metadata.prev_num_accepted_tokens[slots] = next_accepted - replay_metadata.cache_buf_idx[slots] = next_cache_buffer - - -__all__ = ["ReplayLayerCache", "advance_replay_state"] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py deleted file mode 100644 index c92ddc5d747d..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/seeds.py +++ /dev/null @@ -1,66 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Deterministic per-slot seeds shared by recurrent cache algorithms.""" - -from __future__ import annotations - -import torch - -from tensorrt_llm._utils import prefer_pinned - -_MAMBA_SSM_SEED_MASK = (1 << 62) - 1 -_MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 -_MAMBA_SSM_SEED_BASE = 0x6A09E667F3BCC908 -_MAMBA_SSM_SEED_MIX_COUNTER = 0x2545F4914F6CDD1D -_MAMBA_SSM_SEED_MIX_SLOT = 0x1B873593CC9E2D51 -_MAMBA_SSM_SEED_MIX_RANK = 0x9E3779B97F4A7C15 - - -def _splitmix64(value: int) -> int: - """Return the SplitMix64 finalizer for ``value``.""" - value = (value + 0x9E3779B97F4A7C15) & _MAMBA_SSM_UINT64_MASK - value ^= value >> 30 - value = (value * 0xBF58476D1CE4E5B9) & _MAMBA_SSM_UINT64_MASK - value ^= value >> 27 - value = (value * 0x94D049BB133111EB) & _MAMBA_SSM_UINT64_MASK - value ^= value >> 31 - return value & _MAMBA_SSM_UINT64_MASK - - -def compute_deterministic_mamba_seed(counter: int, slot: int, rank_offset: int) -> int: - """Return a reproducible, nonzero int64 Philox seed.""" - folded = ( - _MAMBA_SSM_SEED_BASE - + counter * _MAMBA_SSM_SEED_MIX_COUNTER - + slot * _MAMBA_SSM_SEED_MIX_SLOT - + rank_offset * _MAMBA_SSM_SEED_MIX_RANK - ) & _MAMBA_SSM_UINT64_MASK - seed = _splitmix64(folded) & _MAMBA_SSM_SEED_MASK - return seed or 1 - - -def allocate_mamba_seed_buffer( - cache_size: int, rank_offset: int, device: torch.device -) -> torch.Tensor: - """Allocate the graph-stable per-slot Philox seed buffer.""" - seeds = [compute_deterministic_mamba_seed(0, slot, rank_offset) for slot in range(cache_size)] - return torch.tensor(seeds, dtype=torch.int64, device=device) - - -def _reset_mamba_seed_buffer( - seed_buffer: torch.Tensor, - slots: torch.Tensor, - host_slots: list[int], - *, - counter: int, - rank_offset: int, -) -> None: - """Refresh selected slots in the graph-stable int64 Philox buffer.""" - seeds = [compute_deterministic_mamba_seed(counter, slot, rank_offset) for slot in host_slots] - seed_buffer[slots] = torch.tensor(seeds, dtype=torch.int64, pin_memory=prefer_pinned()).to( - seed_buffer.device, non_blocking=True - ) - - -__all__ = ["allocate_mamba_seed_buffer", "compute_deterministic_mamba_seed"] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index a0483901524f..4225e8c5a0ec 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -23,8 +23,12 @@ KimiK3HybridCacheManagerV2, ) from tensorrt_llm._torch.modules.mamba.cache_manager import ( + MIN_REPLAY_HISTORY_SIZE, + Mamba2State, NemotronHybridCacheManagerV2, ReplayHistory, + ReplayStateUpdateMetadata, + _advance_replay_state, ) from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( @@ -42,25 +46,19 @@ KVCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - MIN_REPLAY_HISTORY_SIZE, CppMambaHybridCacheManager, MambaCacheManager, MambaHybridCacheManagerV2, MambaRole, MixedMambaHybridCacheManager, PythonMambaCacheManager, - ReplayStateUpdateMetadata, - _advance_replay_state, _get_local_mamba_cache_layout, _get_mamba_hybrid_pool_size, _get_num_cuda_graph_padding_dummy_slots, _mamba_snapshot_rule_counts, _promote_mamba_state_triton, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( - IntermediateState, - _stack_state_views, -) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import _stack_state_views from tensorrt_llm._torch.pyexecutor.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, @@ -3678,7 +3676,7 @@ def test_v2_kda_replay_policy_allocates_logical_slot_caches(): try: assert mgr._kda_replay is policy assert policy.intermediate_indices.shape == (4,) - assert mgr.get_replay_state_update_metadata() is None + assert not hasattr(mgr, "get_replay_state_update_metadata") layer_cache = mgr.mamba_layer_cache(0) assert not hasattr(layer_cache, "intermediate_ssm") @@ -4235,7 +4233,7 @@ def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): lambda **_kwargs: events.append("commit"), ) monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._advance_replay_state", + "tensorrt_llm._torch.modules.fla.cache_manager._advance_replay_state", lambda *_args, **_kwargs: events.append("advance"), ) monkeypatch.setattr( @@ -4318,7 +4316,7 @@ def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): def test_v2_hybrid_dynamic_tree_promotes_accepted_leaf_state(monkeypatch): mgr = object.__new__(NemotronHybridCacheManagerV2) mgr.local_num_mamba_layers = 1 - policy = IntermediateState() + policy = Mamba2State() policy.intermediate_indices = torch.arange(2, dtype=torch.int32) policy.intermediate_ssm = torch.empty((1, 2, 8)) policy.intermediate_conv = torch.empty((1, 2, 8)) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index 08598440c565..28b985154271 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -9,13 +9,21 @@ import pytest import torch -from tensorrt_llm._torch.modules.fla.cache_manager import GDNReplayState, Qwen35HybridCacheManagerV2 +from tensorrt_llm._torch.modules.fla.cache_manager import ( + GDNIntermediateState, + GDNLayerCache, + GDNReplayState, + Qwen35HybridCacheManagerV2, +) from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAIntermediateLayerCache, + KDAIntermediateState, KDAReplayLayerCache, KDAReplayState, KimiK3HybridCacheManagerV2, ) from tensorrt_llm._torch.modules.mamba.cache_manager import ( + Mamba2LayerCache, Mamba2State, NemotronHybridCacheManagerV2, ReplayHistory, @@ -29,14 +37,12 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( - IntermediateLayerCache, - IntermediateState, MambaAcceptanceBatch, MambaHybridCacheManager, + MambaLayerCache, MambaRole, MambaStateLayout, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.replay import ReplayLayerCache from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, @@ -173,6 +179,7 @@ def _state_views(layers=2, slots=5): def _prepare_model_layout(manager): """Supply constructor-resolved inputs when testing the pre-pool hook alone.""" + manager._mamba_ssm_stochastic_rounding = False manager._state_layout = replace(_layout(), spec_config=manager.spec_config) manager.local_num_mamba_layers = 2 manager._global_n_groups = 2 @@ -242,12 +249,12 @@ def test_k3_initializes_only_the_selected_state( else None ) _prepare_model_layout(manager) - intermediate_factory = Mock(wraps=IntermediateState) + intermediate_factory = Mock(wraps=KDAIntermediateState) replay_factory = Mock(wraps=KDAReplayState) selector = Mock(return_value=selected_num_spec) monkeypatch.setitem( KimiK3HybridCacheManagerV2._initialize_model_state.__globals__, - "IntermediateState", + "KDAIntermediateState", intermediate_factory, ) monkeypatch.setitem( @@ -265,7 +272,7 @@ def test_k3_initializes_only_the_selected_state( num_spec = requested_num_spec if requested_num_spec is not None else selected_num_spec if num_spec is not None: intermediate_factory.assert_not_called() - replay_factory.assert_called_once_with(num_spec) + replay_factory.assert_called_once_with(num_spec, stochastic_rounding=False) assert manager._speculative_state is manager._kda_replay state = manager._kda_replay else: @@ -315,7 +322,7 @@ def test_base_manager_without_model_state(): assert manager.intermediate_state_indices is None assert manager.intermediate_ssm_states is None assert manager.intermediate_conv_states is None - assert manager.get_replay_state_update_metadata() is None + assert not hasattr(manager, "get_replay_state_update_metadata") @pytest.mark.parametrize( @@ -336,7 +343,8 @@ def test_mamba2_seed_lifecycle_without_speculative_decoding(): manager = object.__new__(NemotronHybridCacheManagerV2) manager.spec_config = None manager._requested_replay = False - manager._state_layout = replace(_layout(), spec_config=None, stochastic_rounding=True) + manager._mamba_ssm_stochastic_rounding = True + manager._state_layout = replace(_layout(), spec_config=None) views = _state_views() manager.all_ssm_states = views.all_ssm_states manager.all_conv_states = views.all_conv_states @@ -354,7 +362,7 @@ def test_mamba2_seed_lifecycle_without_speculative_decoding(): @pytest.mark.parametrize( - "state", [IntermediateState(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)] + "state", [GDNIntermediateState(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)] ) def test_common_manager_consumes_model_state_contract(state): manager = object.__new__(MambaHybridCacheManagerV2) @@ -378,9 +386,7 @@ def test_common_manager_consumes_model_state_contract(state): assert payload.old_x is not None else: assert payload.intermediate_ssm is not None - assert (manager.get_replay_state_update_metadata() is not None) == isinstance( - state, (ReplayHistory, GDNReplayState) - ) + assert not hasattr(manager, "get_replay_state_update_metadata") manager._shutdown_model_state() assert manager.intermediate_state_indices is None assert manager.intermediate_ssm_states is None @@ -409,8 +415,8 @@ def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): _prepare_model_layout(manager) manager._speculative_state = manager._initialize_model_state() state = manager._speculative_state - assert isinstance(state, IntermediateState) == (mode != "replay") - assert isinstance(state, ReplayHistory) == (mode == "replay") + assert isinstance(state, (Mamba2State, GDNIntermediateState)) == (mode != "replay") + assert isinstance(state, (ReplayHistory, GDNReplayState)) == (mode == "replay") manager._state_layout = object() manager.local_num_mamba_layers = 0 @@ -543,7 +549,7 @@ def wrap_buffer(role, dtype, shape): monkeypatch.setattr(torch.cuda, "current_device", lambda: "cpu") monkeypatch.setattr(torch.cuda, "synchronize", lambda: None) manager._setup_states() - manager._speculative_state = GDNReplayState(3) if replay else IntermediateState() + manager._speculative_state = GDNReplayState(3) if replay else GDNIntermediateState() manager._setup_model_state() state = manager._speculative_state scratch_before = {} @@ -596,7 +602,7 @@ def test_gdn_bind_accepts_affine_and_indirect_views(affine): @pytest.mark.parametrize( - "state", [IntermediateState(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)] + "state", [GDNIntermediateState(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)] ) def test_no_local_recurrent_layers_allocate_no_algorithm_buffers(state): state.bind(_layout(layers=(), slot_capacity=0), [], []) @@ -657,30 +663,37 @@ def test_kda_scratch_cost_matches_allocated_slot_buffers(): @pytest.mark.parametrize("layer", [0, 1]) -@pytest.mark.parametrize("kind", ["intermediate", "mamba2", "replay", "gdn", "kda"]) +@pytest.mark.parametrize( + "kind", ["intermediate", "kda_intermediate", "mamba2", "replay", "gdn", "kda"] +) def test_layer_payload_preserves_storage_and_strides(kind, layer): states = { - "intermediate": IntermediateState, - "mamba2": Mamba2State, + "intermediate": GDNIntermediateState, + "kda_intermediate": KDAIntermediateState, + "mamba2": lambda: Mamba2State(stochastic_rounding=True), "replay": lambda: ReplayHistory(3), "gdn": lambda: GDNReplayState(3), - "kda": lambda: KDAReplayState(2), + "kda": lambda: KDAReplayState(2, stochastic_rounding=True), } state = states[kind]() - layout = replace(_layout(), stochastic_rounding=True) + layout = _layout() # Interleave persistent slots to exercise borrowed, non-contiguous pool views. conv = torch.zeros(2, 10, 24, 3)[:, ::2] ssm = torch.zeros(2, 10, 2, 4, 4)[:, ::2] state.bind(layout, list(ssm.unbind()), list(conv.unbind())) payload = state.make_layer_cache(layer, conv[layer], ssm[layer]) expected_type = { - "intermediate": IntermediateLayerCache, - "mamba2": IntermediateLayerCache, - "replay": ReplayLayerCache, - "gdn": ReplayLayerCache, + "intermediate": GDNLayerCache, + "kda_intermediate": KDAIntermediateLayerCache, + "mamba2": Mamba2LayerCache, + "replay": Mamba2LayerCache, + "gdn": GDNLayerCache, "kda": KDAReplayLayerCache, }[kind] assert type(payload) is expected_type + assert expected_type.__bases__ == (MambaLayerCache,) + if kind in ("replay", "gdn", "kda"): + assert expected_type.__module__ == type(state).__module__ pairs = [(payload.conv, conv[layer]), (payload.temporal, ssm[layer])] if kind == "kda": assert payload.has_kda_replay_caches @@ -709,6 +722,8 @@ def test_layer_payload_preserves_storage_and_strides(kind, layer): pairs.append((payload.intermediate_ssm, state.intermediate_ssm[layer])) if kind == "intermediate": assert payload.mamba_ssm_rand_seed is None + elif kind == "kda_intermediate": + assert not hasattr(payload, "mamba_ssm_rand_seed") else: assert payload.mamba_ssm_rand_seed is state.rand_seed for view, source in pairs: @@ -721,8 +736,8 @@ def test_layer_payload_preserves_storage_and_strides(kind, layer): def test_kda_seed_lifecycle_is_independent_of_layer_payload(): - state = KDAReplayState(2) - layout = replace(_layout(), stochastic_rounding=True, seed_rank_offset=17) + state = KDAReplayState(2, stochastic_rounding=True) + layout = replace(_layout(), mapping=Mapping(world_size=2, tp_size=2, rank=1)) views = _state_views() state.bind(layout, views.all_ssm_states, views.all_conv_states) original = state.rand_seed.clone() @@ -736,10 +751,8 @@ def test_kda_seed_lifecycle_is_independent_of_layer_payload(): def test_mamba2_seed_lifecycle_does_not_require_speculative_decoding(): - layout = replace( - _layout(), spec_config=None, stochastic_rounding=True, ssm_state_dtype=torch.float16 - ) - first, second = Mamba2State(), Mamba2State() + layout = replace(_layout(), spec_config=None, ssm_state_dtype=torch.float16) + first, second = Mamba2State(stochastic_rounding=True), Mamba2State(stochastic_rounding=True) first.bind(layout, _state_views().all_ssm_states, _state_views().all_conv_states) second.bind(layout, _state_views().all_ssm_states, _state_views().all_conv_states) torch.testing.assert_close(first.rand_seed, second.rand_seed) @@ -752,27 +765,30 @@ def test_mamba2_seed_lifecycle_does_not_require_speculative_decoding(): def test_replay_and_intermediate_are_independent_algorithms(): + assert KDAIntermediateState.__bases__ == (object,) + assert GDNIntermediateState.__bases__ == (object,) + assert KDAReplayState.__bases__ == (object,) assert ReplayHistory.__bases__ == (object,) assert GDNReplayState.__bases__ == (object,) assert not issubclass(GDNReplayState, ReplayHistory) assert ReplayHistory.__module__.endswith("modules.mamba.cache_manager") assert GDNReplayState.__module__.endswith("modules.fla.cache_manager") assert not any("modules.mamba" in cls.__module__ for cls in GDNReplayState.__mro__) - assert Mamba2State.__bases__ == (IntermediateState,) - assert not issubclass(ReplayHistory, IntermediateState) + assert Mamba2State.__bases__ == (object,) + assert not issubclass(ReplayHistory, GDNIntermediateState) assert not issubclass(GDNReplayState, Mamba2State) assert ReplayHistory(3).intermediate_ssm is None assert GDNReplayState(3).intermediate_ssm is None - state = IntermediateState() + state = GDNIntermediateState() assert not {"rand_seed", "_seed_request_counter", "_seed_rank_offset"}.intersection(vars(state)) assert not hasattr(state, "_promote_conv") @pytest.mark.parametrize("replay_cls", [ReplayHistory, GDNReplayState]) def test_mamba_seed_reset_matches_between_intermediate_and_replay(replay_cls): - layout = replace(_layout(), stochastic_rounding=True, seed_rank_offset=17) + layout = replace(_layout(), mapping=Mapping(world_size=2, tp_size=2, rank=1)) views = _state_views() - intermediate, replay = Mamba2State(), replay_cls(3) + intermediate, replay = Mamba2State(stochastic_rounding=True), replay_cls(3) for state in (intermediate, replay): state.bind(layout, views.all_ssm_states, views.all_conv_states) torch.testing.assert_close(intermediate.rand_seed, replay.rand_seed) @@ -800,7 +816,8 @@ def test_mamba_seed_reset_matches_between_intermediate_and_replay(replay_cls): @pytest.mark.parametrize( - "state_cls", [IntermediateState, Mamba2State, ReplayHistory, GDNReplayState] + "state_cls", + [GDNIntermediateState, KDAIntermediateState, Mamba2State, ReplayHistory, GDNReplayState], ) def test_accepted_state_promotion_preserves_algorithm_boundary(monkeypatch, state_cls): from tensorrt_llm._torch.pyexecutor.kv_cache import mamba_cache_manager @@ -856,20 +873,26 @@ def promote(destination, source, source_indices, positions, destination_indices) @pytest.mark.parametrize( "state", - [IntermediateState(), Mamba2State(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)], + [GDNIntermediateState(), Mamba2State(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)], ) def test_algorithm_borrows_only_tensor_views_and_releases_them(state): views = _state_views() state.bind(_layout(), views.all_ssm_states, views.all_conv_states) assert state._conv_states[0] is views.all_conv_states[0] assert not hasattr(state, "manager") - if isinstance(state, (IntermediateState, ReplayHistory, GDNReplayState)): + if isinstance( + state, + (GDNIntermediateState, KDAIntermediateState, Mamba2State, ReplayHistory, GDNReplayState), + ): assert state._ssm_states[0] is views.all_ssm_states[0] assert not hasattr(state, "_publish_compatibility_views") state.shutdown() state.shutdown() assert state._conv_states == () - if isinstance(state, (IntermediateState, ReplayHistory, GDNReplayState)): + if isinstance( + state, + (GDNIntermediateState, KDAIntermediateState, Mamba2State, ReplayHistory, GDNReplayState), + ): assert state._ssm_states == () assert not any(isinstance(value, torch.Tensor) for value in vars(state).values()) # Releasing borrowed references must not mutate the persistent pool. @@ -877,7 +900,7 @@ def test_algorithm_borrows_only_tensor_views_and_releases_them(state): @pytest.mark.parametrize( - "cls", [IntermediateState, Mamba2State, ReplayHistory, GDNReplayState, KDAReplayState] + "cls", [GDNIntermediateState, Mamba2State, ReplayHistory, GDNReplayState, KDAReplayState] ) def test_algorithm_methods_do_not_depend_on_manager(cls): for function in vars(cls).values(): @@ -904,7 +927,7 @@ def test_algorithm_methods_do_not_depend_on_manager(cls): ) def test_model_compatibility_properties_follow_owned_state(manager_cls): manager = object.__new__(manager_cls) - manager._speculative_state = IntermediateState() + manager._speculative_state = GDNIntermediateState() if manager_cls is KimiK3HybridCacheManagerV2: manager._kda_replay = None fields = { diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index abc7b55a119e..43f18a9e986a 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -21,6 +21,10 @@ import torch from tensorrt_llm._torch.modules.mamba import mamba2_metadata +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + MIN_REPLAY_HISTORY_SIZE, + ReplayStateUpdateMetadata, +) from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( REPLAY_WORK_CACHE_BUF_IDX, REPLAY_WORK_CACHE_SLOT, @@ -32,10 +36,6 @@ cu_seqlens_to_chunk_indices_offsets, cu_seqlens_to_chunk_indices_offsets_triton, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - MIN_REPLAY_HISTORY_SIZE, - ReplayStateUpdateMetadata, -) skip_no_cuda = pytest.mark.skipif( not torch.cuda.is_available(), diff --git a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py index 48ef92fb8d72..aab23884f0af 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py @@ -11,12 +11,16 @@ import pytest import torch -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - PythonMambaCacheManager, - _allocate_mamba_seed_buffer, - _compute_deterministic_mamba_seed, - _mamba_rank_offset, +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + allocate_mamba_seed_buffer as _allocate_mamba_seed_buffer, ) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + compute_deterministic_mamba_seed as _compute_deterministic_mamba_seed, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + mamba_seed_rank_offset as _mamba_rank_offset, +) +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import PythonMambaCacheManager from tensorrt_llm.mapping import Mapping From c418620832da034398c554a2499e383608a29dfb Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon, 21 Sep 2026 10:17:58 +0000 Subject: [PATCH 11/25] [None][refactor] Remove redundant model state forwarding hooks (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../_torch/modules/qwen4_exp/cache_manager.py | 4 +- .../mamba_cache_manager_v2.py | 26 ++--- .../kv_cache/test_mamba_model_managers.py | 101 +++++++++++++----- 3 files changed, 84 insertions(+), 47 deletions(-) diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index 7236c9f58765..11d46e867bca 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -556,10 +556,10 @@ def _extra_state_bytes_per_rank(cls, model_config, mapping, **kwargs) -> int: ) @override - def _shutdown_model_state(self) -> None: + def shutdown(self) -> None: self._ple_conv_states.clear() self._ple_ngram_contexts.clear() - super()._shutdown_model_state() + super().shutdown() def get_replay_state_update_metadata(self) -> ReplayStateUpdateMetadata | None: state = self._speculative_state diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 99cbedfa0a5f..8d57f3dd9a66 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -943,17 +943,12 @@ def _on_state_slots_relocated(self, old_slots: list[int], new_slots: list[int]) """Allow model-owned scratch state to follow the updated slot mapping.""" def _reset_context_mamba_slots(self, num_contexts: int) -> None: - if num_contexts: - self._reset_model_slots( + if num_contexts and self._speculative_state is not None: + self._speculative_state.reset_slots( self.cuda_state_indices[:num_contexts].long(), self._host_state_indices[:num_contexts].tolist(), ) - def _reset_model_slots(self, slots: torch.Tensor, host_slots: list[int]) -> None: - """Reset any model-owned request state for the supplied slots.""" - if self._speculative_state is not None: - self._speculative_state.reset_slots(slots, host_slots) - def _refresh_dummy_request_mask(self, is_dummy: list[bool]) -> None: count = len(is_dummy) if count > self._dummy_request_mask_host.shape[0]: @@ -1029,7 +1024,9 @@ def update_mamba_states( state_indices_d = state_indices[num_contexts : num_contexts + num_gens].to(torch.int32) src_state_indices = self._generation_state_indices[:num_gens] is_dummy_request = self._dummy_request_mask[num_contexts : num_contexts + num_gens] - self._update_speculative_state( + if self._speculative_state is None: + raise NotImplementedError("The model manager must provide speculative state") + self._speculative_state.update( MambaAcceptanceBatch( attention_metadata=attn_metadata, num_contexts=num_contexts, @@ -1042,11 +1039,6 @@ def update_mamba_states( ), ) - def _update_speculative_state(self, batch: MambaAcceptanceBatch) -> None: - if self._speculative_state is None: - raise NotImplementedError("The model manager must provide speculative state") - self._speculative_state.update(batch) - @override def try_commit_blocks(self, request: LlmRequest, kv_cache=None) -> None: should_block_reuse = ( @@ -1330,7 +1322,8 @@ def _iter_cache_buffers_for_invalid_check(self) -> Iterable[tuple[int, torch.Ten @override def shutdown(self): - self._shutdown_model_state() + if self._speculative_state is not None: + self._speculative_state.shutdown() self._recurrent_buffers.clear() self.all_ssm_states = [] self.all_conv_states = [] @@ -1340,10 +1333,5 @@ def shutdown(self): self._branch_snapshot_points.clear() super().shutdown() - def _shutdown_model_state(self) -> None: - """Release model-owned state before the common pool is destroyed.""" - if self._speculative_state is not None: - self._speculative_state.shutdown() - __all__ = ["MambaHybridCacheManagerV2"] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index 28b985154271..7105fb1149a6 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -34,7 +34,7 @@ Qwen4ExpHybridCacheManagerV2, Qwen4ExpPLECacheParams, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager.common import ( MambaAcceptanceBatch, @@ -177,6 +177,30 @@ def _state_views(layers=2, slots=5): ) +def _shutdown_manager(manager, monkeypatch): + """Exercise real shutdown while replacing only the native pool teardown.""" + manager._recurrent_buffers = {"buffer": object()} + manager._branch_snapshot_points = {"request": object()} + if isinstance(manager, Qwen4ExpHybridCacheManagerV2): + manager._ple_conv_states = {42: object()} + manager._ple_ngram_contexts = {42: object()} + + def check_before_pool_shutdown(): + assert manager.intermediate_state_indices is None + assert manager.intermediate_ssm_states is None + assert manager.intermediate_conv_states is None + assert not manager._recurrent_buffers + assert not manager._branch_snapshot_points + assert not manager.all_ssm_states and not manager.all_conv_states + if isinstance(manager, Qwen4ExpHybridCacheManagerV2): + assert not manager._ple_conv_states and not manager._ple_ngram_contexts + + pool_shutdown = Mock(side_effect=check_before_pool_shutdown) + monkeypatch.setattr(KVCacheManagerV2, "shutdown", pool_shutdown) + manager.shutdown() + pool_shutdown.assert_called_once_with() + + def _prepare_model_layout(manager): """Supply constructor-resolved inputs when testing the pre-pool hook alone.""" manager._mamba_ssm_stochastic_rounding = False @@ -287,38 +311,31 @@ def test_k3_initializes_only_the_selected_state( manager._state_layout = object() manager.all_ssm_states, manager.all_conv_states = [], [] - bind, update, reset = Mock(), Mock(), Mock() + bind = Mock() monkeypatch.setattr(state, "bind", bind) - monkeypatch.setattr(state, "update", update) - monkeypatch.setattr(state, "reset_slots", reset) manager._setup_model_state() bind.assert_called_once_with( manager._state_layout, manager.all_ssm_states, manager.all_conv_states ) - batch, slots, host_slots = object(), object(), [0] - manager._update_speculative_state(batch) - update.assert_called_once_with(batch) - manager._reset_model_slots(slots, host_slots) - reset.assert_called_once_with(slots, host_slots) shutdown = Mock(wraps=state.shutdown) monkeypatch.setattr(state, "shutdown", shutdown) - manager._shutdown_model_state() + _shutdown_manager(manager, monkeypatch) shutdown.assert_called_once_with() - manager._shutdown_model_state() + _shutdown_manager(manager, monkeypatch) assert shutdown.call_count == 2 if num_spec is not None: assert manager._speculative_state is manager._kda_replay -def test_base_manager_without_model_state(): +def test_base_manager_without_model_state(monkeypatch): manager = object.__new__(MambaHybridCacheManagerV2) manager.spec_config = None manager._speculative_state = manager._initialize_model_state() assert manager._speculative_state is None manager._setup_model_state() - manager._reset_model_slots(object(), [0]) - manager._shutdown_model_state() + manager._reset_context_mamba_slots(num_contexts=1) + _shutdown_manager(manager, monkeypatch) assert manager.intermediate_state_indices is None assert manager.intermediate_ssm_states is None assert manager.intermediate_conv_states is None @@ -339,7 +356,7 @@ def test_seed_accessor_is_not_a_common_manager_capability(manager_cls): assert not hasattr(manager_cls, "get_mamba_ssm_rand_seed") -def test_mamba2_seed_lifecycle_without_speculative_decoding(): +def test_mamba2_seed_lifecycle_without_speculative_decoding(monkeypatch): manager = object.__new__(NemotronHybridCacheManagerV2) manager.spec_config = None manager._requested_replay = False @@ -352,19 +369,21 @@ def test_mamba2_seed_lifecycle_without_speculative_decoding(): manager._setup_model_state() seeds = manager.get_mamba_ssm_rand_seed() before = seeds.clone() - manager._reset_model_slots(torch.tensor([1]), [1]) + manager.cuda_state_indices = torch.tensor([1], dtype=torch.int32) + manager._host_state_indices = manager.cuda_state_indices.clone() + manager._reset_context_mamba_slots(num_contexts=1) assert manager.get_mamba_ssm_rand_seed() is seeds assert seeds[1] != before[1] torch.testing.assert_close(seeds[[0, 2, 3, 4]], before[[0, 2, 3, 4]]) assert manager.intermediate_ssm_states is None - manager._shutdown_model_state() + _shutdown_manager(manager, monkeypatch) assert manager.get_mamba_ssm_rand_seed() is None @pytest.mark.parametrize( "state", [GDNIntermediateState(), ReplayHistory(3), GDNReplayState(3), KDAReplayState(2)] ) -def test_common_manager_consumes_model_state_contract(state): +def test_common_manager_consumes_model_state_contract(state, monkeypatch): manager = object.__new__(MambaHybridCacheManagerV2) manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) manager._speculative_state = state @@ -387,7 +406,7 @@ def test_common_manager_consumes_model_state_contract(state): else: assert payload.intermediate_ssm is not None assert not hasattr(manager, "get_replay_state_update_metadata") - manager._shutdown_model_state() + _shutdown_manager(manager, monkeypatch) assert manager.intermediate_state_indices is None assert manager.intermediate_ssm_states is None assert manager.intermediate_conv_states is None @@ -399,8 +418,15 @@ def test_base_manager_requires_model_for_speculative_decoding(): with pytest.raises(NotImplementedError, match="model-specific"): manager._initialize_model_state() manager._speculative_state = None + manager.local_num_mamba_layers = 1 + manager._generation_state_indices = torch.empty(0, dtype=torch.int32) + manager._dummy_request_mask = torch.empty(0, dtype=torch.bool) with pytest.raises(NotImplementedError, match="model manager"): - manager._update_speculative_state(object()) + manager.update_mamba_states( + SimpleNamespace(num_seqs=0, num_contexts=0), + torch.empty(0, dtype=torch.int32), + state_indices=torch.empty(0, dtype=torch.int32), + ) @pytest.mark.parametrize( @@ -448,15 +474,38 @@ def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): state.intermediate_conv = [conv] if mode != "plain" else None assert manager.intermediate_ssm_states == ([ssm] if mode == "intermediate" else None) assert manager.intermediate_conv_states == ([conv] if mode != "plain" else None) - batch, slots, host_slots = object(), object(), [0] - manager._reset_model_slots(slots, host_slots) - reset.assert_called_once_with(slots, host_slots) - manager._update_speculative_state(batch) - update.assert_called_once_with(batch) + manager.local_num_mamba_layers = 1 + manager.cuda_state_indices = torch.tensor([4, 2, 3], dtype=torch.int32) + manager._host_state_indices = manager.cuda_state_indices.clone() + manager._generation_state_indices = torch.arange(2, dtype=torch.int32) + manager._dummy_request_mask = torch.tensor([False, False, True]) + manager._reset_context_mamba_slots(num_contexts=1) + reset.assert_called_once() + slots, host_slots = reset.call_args.args + torch.testing.assert_close(slots, torch.tensor([4], dtype=torch.int64)) + assert host_slots == [4] + manager._reset_context_mamba_slots(num_contexts=0) + reset.assert_called_once() + + metadata = SimpleNamespace(num_seqs=3, num_contexts=1) + manager.update_mamba_states( + metadata, + torch.tensor([1, 3, 2], dtype=torch.int32), + state_indices=manager.cuda_state_indices, + ) + update.assert_called_once() + batch = update.call_args.args[0] + assert batch.attention_metadata is metadata + assert (batch.num_contexts, batch.num_generations) == (1, 2) + assert batch.num_accepted_tokens.tolist() == [3, 2] + assert batch.accepted_positions.tolist() == [2, 1] + assert batch.source_state_indices.tolist() == [0, 1] + assert batch.destination_state_indices.tolist() == [2, 3] + assert batch.is_dummy_request.tolist() == [False, True] shutdown = Mock(wraps=state.shutdown) monkeypatch.setattr(state, "shutdown", shutdown) - manager._shutdown_model_state() + _shutdown_manager(manager, monkeypatch) shutdown.assert_called_once_with() assert manager.intermediate_ssm_states is None and state.intermediate_conv is None if mode == "replay": From d10a78f00ed55978b0e95aafe72e8e8c76a5b03a Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 6 Oct 2026 05:06:22 +0000 Subject: [PATCH 12/25] [None][fix] Align cache manager regression fixtures after merge (By Agent) Use model-owned hybrid managers in transfer tests and exercise KDA replay seeding through the model manager's transfer-completion hook. Complete fixture fields required by the current factory and draft-cache costing paths, and annotate GLM cache manager overrides. Validation on umbriel-b200-027 via tekit: - Full SM100 native build with CUDA 13.4 and PyTorch 2.14 succeeded. - Selected cache, transfer, QSA, GLM, Qwen4 and Mamba regression suite: 1525 passed, 5 skipped, 2 subtests passed; pytest exit code 0. Skips require the unavailable Transformers GLM5-next implementation. - Full pre-commit checks passed. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../sparse/glm_kpool/cache_manager.py | 3 ++ .../kv_cache/test_kv_cache_budget_split.py | 5 ++ .../kv_cache/test_kv_cache_v2_vocab_size.py | 1 + .../kv_cache/test_mamba_cache_manager.py | 52 ++++++++++--------- .../disaggregated/test_mamba_transfer.py | 20 +++++-- 5 files changed, 53 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py index 202a3eabfc3f..da317c4d98a6 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py @@ -230,6 +230,7 @@ def _extra_buffers_per_layer(self, *, tokens_per_block: int) -> dict[int, list[B if layer_id in self.layer_offsets } + @override def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) -> int: index_bytes = ( self.index_state_dim * torch.bfloat16.itemsize @@ -241,6 +242,7 @@ def get_layer_bytes_per_token(self, local_layer_idx: int, data_role: DataRole) - cache_bytes = super().get_layer_bytes_per_token(local_layer_idx, data_role) return cache_bytes + index_bytes if data_role == Role.ALL else cache_bytes + @override def _attention_cache_bytes_per_token(self) -> int: return sum( self.get_layer_bytes_per_token(local_layer_idx, Role.ALL) @@ -248,6 +250,7 @@ def _attention_cache_bytes_per_token(self) -> int: ) @staticmethod + @override def get_cache_size_per_token( model_config: ModelConfig, mapping: Mapping, diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index de1a4ae95ddc..24eb5cf0de60 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -119,6 +119,7 @@ def test_native_full_target_cost_preserves_draft_capacity( c._max_batch_size = max_batch_size c._max_num_tokens = 8192 c._mapping = Mock(enable_attention_dp=False, tp_size=2) + c._mapping.has_cp_helix.return_value = False c._mapping.pp_layers.return_value = list(range(24)) c._mapping.is_last_pp_rank.return_value = True c._speculative_config = SimpleNamespace( @@ -226,6 +227,7 @@ def get_cache_size_per_token(model_config, *args, **kwargs): creator._kv_connector_manager = None creator._cache_transceiver_config = None creator._mapping = Mock(enable_attention_dp=False, tp_size=1) + creator._mapping.has_cp_helix.return_value = False creator._mapping.pp_layers.return_value = [0] creator._mapping.is_last_pp_rank.return_value = True creator._speculative_config = SimpleNamespace( @@ -1069,6 +1071,7 @@ class TestExternalDrafterKvDtype: def _creator(self, mocker, draft_quant_config): class DraftModelConfig: quant_config = draft_quant_config + sparse_attention_config = None pretrained_config = SimpleNamespace( num_hidden_layers=TestExternalDrafterKvDtype.DRAFT_LAYERS, hidden_size=32, @@ -1105,6 +1108,7 @@ def get_cache_size_per_token(model_config, *args, **kwargs): c._tokens_per_block = 64 c._max_seq_len = 16384 c._max_batch_size = 1 + c._max_beam_width = 1 # Read by _build_managers on the draft path (_util.py); the cost # assertions below are per-token, so the value only has to exist. c._max_num_tokens = 8192 @@ -1121,6 +1125,7 @@ def get_cache_size_per_token(model_config, *args, **kwargs): c._draft_config = draft_model_config c._kv_cache_manager_cls = ProbeKVCacheManager c._is_disagg = True + c._cache_transceiver_config = None c._is_encoder_decoder = Mock(return_value=False) c._should_create_separate_draft_kv_cache = Mock(return_value=True) c._get_num_draft_layers = Mock(return_value=self.DRAFT_LAYERS) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py index edc18dd8e58a..6b58b2303d28 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_vocab_size.py @@ -102,6 +102,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: SimpleNamespace( model=SimpleNamespace(model_config=model_config), is_multimodal=is_multimodal, + _max_cuda_graph_batch_size=4, ) if is_multimodal is not None else None diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 669ea7484ffe..6df848a8cc9e 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -163,6 +163,7 @@ def _hybrid_model_config(): pretrained_config=config, sparse_attention_config=None, get_num_mamba_layers=lambda: 1, + mapping=Mapping(), ) @@ -311,6 +312,7 @@ def _kimi_model_config() -> SimpleNamespace: quant_config=None, sparse_attention_config=None, get_num_mamba_layers=lambda: 2, + mapping=Mapping(), ) @@ -392,6 +394,7 @@ def test_kimi_explicit_v2_manager_enables_owned_kda_replay_policy( assert "speculative_state" not in kwargs mgr = object.__new__(KimiK3HybridCacheManagerV2) mgr.spec_config = spec_config + mgr._mamba_ssm_stochastic_rounding = False mgr._requested_num_spec = None mgr._state_layout = SimpleNamespace(spec_config=spec_config, mamba_pp_layers=()) mgr._speculative_state = mgr._initialize_model_state() @@ -4008,36 +4011,37 @@ def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): the seeded slots. """ torch.manual_seed(0) - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr.local_num_mamba_layers = 2 - mgr.conv_state_shape = [6, 4] - mgr.conv_section_dims = [2, 2, 2] + mgr = object.__new__(KimiK3HybridCacheManagerV2) + policy = KDAReplayState(num_speculative_tokens=2) + mgr._kda_replay = policy + policy._committed_window = 4 + policy._conv_section_dims = (2, 2, 2) mgr._request_id_to_state_index = {101: 1, 202: 3} - mgr.prev_num_accepted_tokens = torch.full((4,), 5, dtype=torch.int32) + policy.prev_num_accepted_tokens = torch.full((4,), 5, dtype=torch.int32) # The restored (transferred) conv pool, in the dtype the transceiver moves. mgr.all_conv_states = [(torch.randn(4, 6, 4) * 0.1).to(torch.bfloat16) for _ in range(2)] + policy._conv_states = tuple(mgr.all_conv_states) restored = [conv.clone() for conv in mgr.all_conv_states] def dim_contiguous_conv_cache() -> torch.Tensor: return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) - mgr.kda_conv_q = dim_contiguous_conv_cache() - mgr.kda_conv_k = dim_contiguous_conv_cache() - mgr.kda_conv_v = dim_contiguous_conv_cache() - mgr.kda_qkg_cache = torch.full((2, 4, 2, 3, 2), 7.0) - mgr.kda_v_cache = torch.full((2, 4, 2, 2), 7.0) - mgr.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0) + policy.kda_conv_q = dim_contiguous_conv_cache() + policy.kda_conv_k = dim_contiguous_conv_cache() + policy.kda_conv_v = dim_contiguous_conv_cache() + policy.kda_qkg_cache = torch.full((2, 4, 2, 3, 2), 7.0) + policy.kda_v_cache = torch.full((2, 4, 2, 2), 7.0) + policy.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0) - mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) + mgr.on_state_transfer_complete([101, 202]) for layer_offset, conv_state in enumerate(mgr.all_conv_states): assert conv_state.dtype is torch.bfloat16 for slot in (1, 3): for replay_buffer, start in ( - (mgr.kda_conv_q, 0), - (mgr.kda_conv_k, 2), - (mgr.kda_conv_v, 4), + (policy.kda_conv_q, 0), + (policy.kda_conv_k, 2), + (policy.kda_conv_v, 4), ): # fp32 replay cache, logical shape unchanged by the widening. assert replay_buffer.dtype is torch.float32 @@ -4052,19 +4056,19 @@ def dim_contiguous_conv_cache() -> torch.Tensor: atol=0, ) assert torch.count_nonzero(replay_buffer[layer_offset, slot, :, 4:]) == 0 - assert torch.count_nonzero(mgr.kda_qkg_cache[layer_offset, slot]) == 0 - assert torch.count_nonzero(mgr.kda_v_cache[layer_offset, slot]) == 0 - assert torch.count_nonzero(mgr.kda_beta_cache[layer_offset, slot]) == 0 - assert mgr.prev_num_accepted_tokens[slot] == 0 + assert torch.count_nonzero(policy.kda_qkg_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(policy.kda_v_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(policy.kda_beta_cache[layer_offset, slot]) == 0 + assert policy.prev_num_accepted_tokens[slot] == 0 # Slots no request restored keep their pre-seed contents. for slot in (0, 2): - assert (mgr.kda_conv_q[layer_offset, slot] == 7.0).all() - assert (mgr.kda_qkg_cache[layer_offset, slot] == 7.0).all() + assert (policy.kda_conv_q[layer_offset, slot] == 7.0).all() + assert (policy.kda_qkg_cache[layer_offset, slot] == 7.0).all() # The restored bf16 pool is read-only here: plain decode (speculation # off) reuses exactly the state the transfer wrote. torch.testing.assert_close(conv_state, restored[layer_offset], rtol=0, atol=0) - assert mgr.prev_num_accepted_tokens[0] == 5 - assert mgr.prev_num_accepted_tokens[2] == 5 + assert policy.prev_num_accepted_tokens[0] == 5 + assert policy.prev_num_accepted_tokens[2] == 5 @skip_no_cuda diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index f31a4ca7a17c..ba073c645058 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -37,7 +37,12 @@ from tensorrt_llm._torch.disaggregation.native.mixers.ssm import peer from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -from tensorrt_llm._torch.pyexecutor.config_utils import Qwen4ExpPLECacheParams +from tensorrt_llm._torch.modules.fla.cache_manager import Qwen35HybridCacheManagerV2 +from tensorrt_llm._torch.modules.mamba.cache_manager import NemotronHybridCacheManagerV2 +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + Qwen4ExpHybridCacheManagerV2, + Qwen4ExpPLECacheParams, +) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( MambaHybridCacheManagerV2, MixedMambaHybridCacheManager, @@ -224,7 +229,14 @@ def _create_managers( mapping = Mapping( world_size=tp, rank=rank, tp_size=tp, pp_size=1, enable_attention_dp=enable_attention_dp ) - manager_cls = MambaHybridCacheManagerV2 if use_v2 else MixedMambaHybridCacheManager + if not use_v2: + manager_cls = MixedMambaHybridCacheManager + elif with_ple: + manager_cls = Qwen4ExpHybridCacheManagerV2 + elif conv_state_layout == "q_k_v": + manager_cls = Qwen35HybridCacheManagerV2 + else: + manager_cls = NemotronHybridCacheManagerV2 manager_kwargs = ( { "is_disagg": True, @@ -292,7 +304,7 @@ def _zero_mamba_states(manager): for layer_idx in _mamba_layer_ids(manager): manager.get_conv_states(layer_idx).zero_() manager.get_ssm_states(layer_idx).zero_() - if isinstance(manager, MambaHybridCacheManagerV2): + if isinstance(manager, Qwen4ExpHybridCacheManagerV2): for state in manager._ple_conv_states.values(): state.zero_() for state in manager._ple_ngram_contexts.values(): @@ -854,7 +866,7 @@ def run_mamba_transfer_test( ), ) - # Only the V2 manager owns PLE state. + # PLE buffers are owned by the Qwen4 manager. if with_ple: for gen_rank, manager in enumerate(gen_mgrs): for req_idx, request_id in enumerate(gen_rids): From 640812f3ec6f01cb5fe9f99c1d77db74007efc79 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 6 Oct 2026 08:24:30 +0000 Subject: [PATCH 13/25] [None][refactor] Simplify model cache manager and spec state names (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../sparse/glm_kpool/cache_manager.py | 2 +- .../attention/backends/sparse/qsa/metadata.py | 6 +- .../attention/backends/sparse/registry.py | 4 +- .../_torch/modules/fla/cache_manager.py | 6 +- .../_torch/modules/kimi_kda/cache_manager.py | 6 +- .../_torch/modules/mamba/cache_manager.py | 6 +- .../_torch/modules/qwen4_exp/cache_manager.py | 6 +- tensorrt_llm/_torch/pyexecutor/_util.py | 16 +-- .../mamba_cache_manager_v2.py | 4 +- .../attention/sparse/qsa/test_qsa_sparse.py | 6 +- .../kv_cache/test_mamba_cache_manager.py | 50 +++---- .../kv_cache/test_mamba_model_managers.py | 122 +++++++++--------- .../modeling/test_glm5_next_contracts.py | 2 +- .../modeling/test_qsa_runtime_wiring.py | 12 +- .../_torch/modeling/test_qwen4_exp_support.py | 12 +- .../disaggregated/test_mamba_transfer.py | 14 +- 16 files changed, 137 insertions(+), 137 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py index da317c4d98a6..dee85fd4ba2d 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py @@ -101,7 +101,7 @@ def __init__( super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> KDAIntermediateState | KDAReplayState: + def _initialize_spec_state(self) -> KDAIntermediateState | KDAReplayState: num_spec = self._requested_num_spec if num_spec is None: num_spec = get_kda_replay_num_spec(self.spec_config, manager_supports_replay=True) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index 98b750518da1..0203af3e7331 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py @@ -44,7 +44,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: super().__init__(*args, **kwargs) def __post_init__(self) -> None: - from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager super().__post_init__() # Draft-loop state; set by the speculative driver, read by the indexer. @@ -54,8 +54,8 @@ def __post_init__(self) -> None: self.qsa_shared_topk_indices: Optional[torch.Tensor] = None self.qsa_shared_topk_visible_blocks: Optional[torch.Tensor] = None self.qsa_shared_topk_captured_slots: set[int] = set() - if not isinstance(self.kv_cache_manager, Qwen4ExpHybridCacheManagerV2): - raise TypeError("QSA sparse attention requires Qwen4ExpHybridCacheManagerV2") + if not isinstance(self.kv_cache_manager, Qwen4ExpHybridCacheManager): + raise TypeError("QSA sparse attention requires Qwen4ExpHybridCacheManager") if not isinstance(self.sparse_metadata_params, QSASparseMetadataParams): raise TypeError("QSA sparse attention requires QSASparseMetadataParams") self.qsa_has_local_layers = self.kv_cache_manager.qsa_position_layer_id is not None diff --git a/tensorrt_llm/_torch/attention/backends/sparse/registry.py b/tensorrt_llm/_torch/attention/backends/sparse/registry.py index 315af1b2b116..560714060409 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/registry.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/registry.py @@ -42,9 +42,9 @@ def get_sparse_attn_kv_cache_manager( elif sparse_attention_config.algorithm == "qsa": if not use_kv_cache_manager_v2: raise ValueError("QSA sparse attention requires KV cache manager V2") - from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 + from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager - return Qwen4ExpHybridCacheManagerV2 + return Qwen4ExpHybridCacheManager else: raise ValueError( f"Unsupported sparse attention algorithm: {sparse_attention_config.algorithm}" diff --git a/tensorrt_llm/_torch/modules/fla/cache_manager.py b/tensorrt_llm/_torch/modules/fla/cache_manager.py index 71046e203684..f962bbdc1deb 100644 --- a/tensorrt_llm/_torch/modules/fla/cache_manager.py +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -582,7 +582,7 @@ def validate_gdn_layout( state.validate(manager._state_layout) -class Qwen35HybridCacheManagerV2(MambaHybridCacheManagerV2): +class Qwen35HybridCacheManager(MambaHybridCacheManagerV2): """Qwen3.5/Qwen3Next recurrent state with shared GDN replay.""" @override @@ -593,7 +593,7 @@ def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> GDNIntermediateState | GDNReplayState: + def _initialize_spec_state(self) -> GDNIntermediateState | GDNReplayState: state = create_gdn_state(self, self._requested_replay) validate_gdn_layout(self, state) @@ -635,7 +635,7 @@ def get_gdn_cache_params(config, *, spec_config=None, quant_config=None): "GDNIntermediateState", "GDNLayerCache", "GDNReplayState", - "Qwen35HybridCacheManagerV2", + "Qwen35HybridCacheManager", "create_gdn_state", "select_gdn_replay_state", "validate_gdn_layout", diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index df122dc85de2..7a3c6c63bea3 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -596,7 +596,7 @@ def get_kda_replay_num_spec( return spec_config.tokens_per_gen_step - 1 -class KimiK3HybridCacheManagerV2(MambaHybridCacheManagerV2): +class KimiK3HybridCacheManager(MambaHybridCacheManagerV2): """KDA/MLA cache layout and K3-only replay lifecycle handling.""" @override @@ -616,7 +616,7 @@ def __init__( super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> KDAIntermediateState | KDAReplayState: + def _initialize_spec_state(self) -> KDAIntermediateState | KDAReplayState: num_spec = self._requested_num_spec if num_spec is None: num_spec = get_kda_replay_num_spec(self.spec_config, manager_supports_replay=True) @@ -796,7 +796,7 @@ def get_kimi_cache_params(config, *, spec_config=None, quant_config=None): "KDAIntermediateState", "KDAReplayState", "KDAReplayLayerCache", - "KimiK3HybridCacheManagerV2", + "KimiK3HybridCacheManager", "allocate_kda_replay_fields", "get_kda_replay_num_spec", "get_kimi_cache_params", diff --git a/tensorrt_llm/_torch/modules/mamba/cache_manager.py b/tensorrt_llm/_torch/modules/mamba/cache_manager.py index ce9dbef95caf..36e2a2b0d236 100644 --- a/tensorrt_llm/_torch/modules/mamba/cache_manager.py +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -476,7 +476,7 @@ def select_mamba2_state( ) -class NemotronHybridCacheManagerV2(MambaHybridCacheManagerV2): +class NemotronHybridCacheManager(MambaHybridCacheManagerV2): """Nemotron/Mamba2 geometry, replay selection and seed ownership.""" _speculative_state: Mamba2State | ReplayHistory @@ -489,7 +489,7 @@ def __init__(self, *args, use_replay_state_update: bool | None = None, **kwargs) super().__init__(*args, **kwargs) @override - def _initialize_model_state(self) -> Mamba2State | ReplayHistory: + def _initialize_spec_state(self) -> Mamba2State | ReplayHistory: if self._requested_replay is not None: state = create_mamba2_state( spec_config=self.spec_config, @@ -546,7 +546,7 @@ def get_nemotron_cache_params(config, *, spec_config=None, quant_config=None): "Mamba2LayerCache", "Mamba2State", "ReplayHistory", - "NemotronHybridCacheManagerV2", + "NemotronHybridCacheManager", "create_mamba2_state", "select_mamba2_state", "get_nemotron_cache_params", diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index 1e0b4f82c753..d26564a49f19 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -161,7 +161,7 @@ def _qwen4_exp_ple_state_bytes_per_rank( return local_ple_layers * ple_bytes_per_layer -class Qwen4ExpHybridCacheManagerV2(MambaHybridCacheManagerV2): +class Qwen4ExpHybridCacheManager(MambaHybridCacheManagerV2): """Hybrid GDN/KV manager with lifecycle-coupled sparse side buffers. QSA adds a per-layer index-K cache and request-wide position coordinates. @@ -519,7 +519,7 @@ def _ple_buffer_configs(self) -> list[BufferConfig]: ] @override - def _initialize_model_state(self) -> GDNIntermediateState | GDNReplayState: + def _initialize_spec_state(self) -> GDNIntermediateState | GDNReplayState: state = create_gdn_state(self, self._requested_replay) validate_gdn_layout(self, state) @@ -609,7 +609,7 @@ def get_qwen4_exp_ple_cache_params(config, *, total_layers: int, is_draft: bool) __all__ = [ - "Qwen4ExpHybridCacheManagerV2", + "Qwen4ExpHybridCacheManager", "Qwen4ExpPLECacheParams", "extract_qwen4_exp_ple_cache_params", "get_qwen4_exp_ple_layer_mask", diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index d1f896033daa..fe679ff6b340 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -406,17 +406,17 @@ def get_kv_cache_manager_cls( sparse_attn_config, use_kv_cache_manager_v2=True) if is_kimi_linear(config): from ..modules.kimi_kda.cache_manager import \ - KimiK3HybridCacheManagerV2 - return KimiK3HybridCacheManagerV2 + KimiK3HybridCacheManager + return KimiK3HybridCacheManager if is_qwen4_exp(config): from ..modules.qwen4_exp.cache_manager import \ - Qwen4ExpHybridCacheManagerV2 - return Qwen4ExpHybridCacheManagerV2 + Qwen4ExpHybridCacheManager + return Qwen4ExpHybridCacheManager if is_qwen3_hybrid(config): - from ..modules.fla.cache_manager import Qwen35HybridCacheManagerV2 - return Qwen35HybridCacheManagerV2 - from ..modules.mamba.cache_manager import NemotronHybridCacheManagerV2 - return NemotronHybridCacheManagerV2 + from ..modules.fla.cache_manager import Qwen35HybridCacheManager + return Qwen35HybridCacheManager + from ..modules.mamba.cache_manager import NemotronHybridCacheManager + return NemotronHybridCacheManager elif sparse_attn_config is not None: if sparse_attn_algorithm == "qsa": # The QSA manager extends the hybrid manager because it must retain diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 8d57f3dd9a66..e60d20ef85ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -264,7 +264,7 @@ def __init__( conv_section_dims=tuple(self.conv_section_dims), conv_state_layout=conv_state_layout, ) - self._speculative_state: MambaState | None = self._initialize_model_state() + self._speculative_state: MambaState | None = self._initialize_spec_state() if isinstance(num_kv_heads, int): per_layer_kv_heads = [num_kv_heads] * total_layers @@ -319,7 +319,7 @@ def _disagg_transfer_overwrites_whole_cached_prefix(self) -> bool: or super()._disagg_transfer_overwrites_whole_cached_prefix() ) - def _initialize_model_state(self) -> MambaState | None: + def _initialize_spec_state(self) -> MambaState | None: """Return validated model-owned state before constructing the pool.""" if self.spec_config is not None: raise NotImplementedError( diff --git a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py index de10bf11987a..750a406c84da 100644 --- a/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py +++ b/tests/unittest/_torch/attention/sparse/qsa/test_qsa_sparse.py @@ -36,7 +36,7 @@ select_qsa_paged_tokens, select_qsa_tokens, ) -from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode @@ -1451,7 +1451,7 @@ def test_fused_qsa_prefill_bounds_sparse_attention_to_visible_tokens() -> None: def test_qsa_side_buffers_use_exact_geometry_and_one_position_role() -> None: - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager._qsa_enabled = True manager.qsa_index_dim = 128 manager.qsa_index_kv_heads = 1 @@ -1477,7 +1477,7 @@ def get_page_stride(self, layer_idx, role): del layer_idx, role return 0 - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.qsa_position_layer_id = 7 manager.layer_offsets = {7: 2} manager.impl = _Impl() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 6df848a8cc9e..4df8c9436fa6 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -16,16 +16,16 @@ MambaLayerGroup, ) from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -from tensorrt_llm._torch.modules.fla.cache_manager import GDNReplayState, Qwen35HybridCacheManagerV2 +from tensorrt_llm._torch.modules.fla.cache_manager import GDNReplayState, Qwen35HybridCacheManager from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( _KDA_BETA_CACHE_ALIGNMENT_BYTES, KDAReplayState, - KimiK3HybridCacheManagerV2, + KimiK3HybridCacheManager, ) from tensorrt_llm._torch.modules.mamba.cache_manager import ( MIN_REPLAY_HISTORY_SIZE, Mamba2State, - NemotronHybridCacheManagerV2, + NemotronHybridCacheManager, ReplayHistory, ReplayStateUpdateMetadata, _advance_replay_state, @@ -324,7 +324,7 @@ def _capture_kimi_v2_manager_ctor( V2 manager and capture the constructor arguments.""" captured: dict[str, object] = {} - class RecordingV2Manager(KimiK3HybridCacheManagerV2): + class RecordingV2Manager(KimiK3HybridCacheManager): def __init__(self, *args: object, **kwargs: object) -> None: captured["args"] = args captured["kwargs"] = kwargs @@ -334,7 +334,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: monkeypatch.delenv("TRTLLM_USE_PY_MAMBA", raising=False) monkeypatch.delenv("TLLM_MAMBA_MANAGER_PREFERENCE", raising=False) - assert get_kv_cache_manager_cls(model_config, kv_cache_config) is KimiK3HybridCacheManagerV2 + assert get_kv_cache_manager_cls(model_config, kv_cache_config) is KimiK3HybridCacheManager _create_kv_cache_manager( model_engine=None, @@ -392,12 +392,12 @@ def test_kimi_explicit_v2_manager_enables_owned_kda_replay_policy( _, kwargs = _capture_kimi_v2_manager_ctor(monkeypatch, spec_config) assert "speculative_state" not in kwargs - mgr = object.__new__(KimiK3HybridCacheManagerV2) + mgr = object.__new__(KimiK3HybridCacheManager) mgr.spec_config = spec_config mgr._mamba_ssm_stochastic_rounding = False mgr._requested_num_spec = None mgr._state_layout = SimpleNamespace(spec_config=spec_config, mamba_pp_layers=()) - mgr._speculative_state = mgr._initialize_model_state() + mgr._speculative_state = mgr._initialize_spec_state() policy = mgr._kda_replay assert isinstance(policy, KDAReplayState) assert policy.num_speculative_tokens == spec_config.tokens_per_gen_step - 1 @@ -488,8 +488,8 @@ def test_v2_manager_rejects_model_type_kwarg() -> None: @pytest.mark.parametrize( ("use_v2", "enable_block_reuse", "expected"), [ - (True, False, Qwen35HybridCacheManagerV2), - (True, True, Qwen35HybridCacheManagerV2), + (True, False, Qwen35HybridCacheManager), + (True, True, Qwen35HybridCacheManager), (False, True, CppMambaHybridCacheManager), ("auto", True, CppMambaHybridCacheManager), ], @@ -606,7 +606,7 @@ def __init__(self, *args, **kwargs): assert captured_mixed["max_num_tokens"] == 256 assert captured_v2["use_replay_state_update"] is None assert "speculative_state" not in captured_v2 - mgr = object.__new__(Qwen35HybridCacheManagerV2) + mgr = object.__new__(Qwen35HybridCacheManager) mgr.spec_config = common_kwargs["spec_config"] mgr.ssm_state_dtype = torch.bfloat16 mgr._requested_replay = None @@ -614,7 +614,7 @@ def __init__(self, *args, **kwargs): mgr._state_layout = SimpleNamespace( spec_config=mgr.spec_config, mamba_pp_layers=(), n_groups_per_rank=1 ) - mgr._speculative_state = mgr._initialize_model_state() + mgr._speculative_state = mgr._initialize_spec_state() assert isinstance(mgr._speculative_state, GDNReplayState) assert captured_v2["max_num_tokens"] == 256 assert captured_v2["vocab_size"] == pretrained_config.vocab_size @@ -795,7 +795,7 @@ def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, bac backend=backend, transceiver_runtime="PYTHON" ), ) - is Qwen35HybridCacheManagerV2 + is Qwen35HybridCacheManager ) @@ -989,7 +989,7 @@ def test_qwen3_gdn_replay_uses_v2_preference( ) assert llm_args.kv_cache_config.use_kv_cache_manager_v2 is expected_v2 - expected_manager = Qwen35HybridCacheManagerV2 if expected_v2 else CppMambaHybridCacheManager + expected_manager = Qwen35HybridCacheManager if expected_v2 else CppMambaHybridCacheManager assert ( get_kv_cache_manager_cls( _hybrid_model_config(), @@ -2436,17 +2436,17 @@ def _build_v2_hybrid_with_mamba_layer( dtype=kv_cache_dtype, ) model_cls = ( - KimiK3HybridCacheManagerV2 + KimiK3HybridCacheManager if isinstance(speculative_state, KDAReplayState) - else Qwen35HybridCacheManagerV2 + else Qwen35HybridCacheManager if conv_state_layout == "q_k_v" - else NemotronHybridCacheManagerV2 + else NemotronHybridCacheManager ) class ModelStateManager(model_cls): - def _initialize_model_state(self): + def _initialize_spec_state(self): if speculative_state is None: - return super()._initialize_model_state() + return super()._initialize_spec_state() elif isinstance(speculative_state, KDAReplayState): self._kda_replay = speculative_state self._kda_replay.validate(self._state_layout) @@ -3778,7 +3778,7 @@ def test_v2_kda_replay_policy_supports_attention_only_pp_rank(): def test_v2_kda_replay_policy_records_acceptance_and_skips_dummy_rows(): - mgr = object.__new__(KimiK3HybridCacheManagerV2) + mgr = object.__new__(KimiK3HybridCacheManager) mgr.local_num_mamba_layers = 1 policy = KDAReplayState( num_speculative_tokens=2, @@ -3809,7 +3809,7 @@ def test_legacy_mamba_cache_manager_delegates_kda_replay_capability() -> None: def test_v2_kda_replay_policy_resets_context_slots(): - mgr = object.__new__(KimiK3HybridCacheManagerV2) + mgr = object.__new__(KimiK3HybridCacheManager) policy = KDAReplayState( num_speculative_tokens=2, ) @@ -3826,7 +3826,7 @@ def test_v2_kda_replay_policy_resets_context_slots(): def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): - mgr = object.__new__(KimiK3HybridCacheManagerV2) + mgr = object.__new__(KimiK3HybridCacheManager) mgr.local_num_mamba_layers = 1 mgr.spec_config = SimpleNamespace(decoding_type="NGram") mgr._request_id_to_state_index = {101: 1, 303: 2} @@ -4011,7 +4011,7 @@ def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): the seeded slots. """ torch.manual_seed(0) - mgr = object.__new__(KimiK3HybridCacheManagerV2) + mgr = object.__new__(KimiK3HybridCacheManager) policy = KDAReplayState(num_speculative_tokens=2) mgr._kda_replay = policy policy._committed_window = 4 @@ -4212,7 +4212,7 @@ def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): - mgr = object.__new__(Qwen35HybridCacheManagerV2) + mgr = object.__new__(Qwen35HybridCacheManager) batch_size = 16 mgr.local_num_mamba_layers = 1 policy = GDNReplayState(tokens_per_gen_step=4) @@ -4266,7 +4266,7 @@ def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): - mgr = object.__new__(NemotronHybridCacheManagerV2) + mgr = object.__new__(NemotronHybridCacheManager) mgr.local_num_mamba_layers = 1 mgr._request_id_to_state_index = { 100: 0, @@ -4318,7 +4318,7 @@ def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): def test_v2_hybrid_dynamic_tree_promotes_accepted_leaf_state(monkeypatch): - mgr = object.__new__(NemotronHybridCacheManagerV2) + mgr = object.__new__(NemotronHybridCacheManager) mgr.local_num_mamba_layers = 1 policy = Mamba2State() policy.intermediate_indices = torch.arange(2, dtype=torch.int32) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index 436bc037aad3..cd014b8014e6 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -16,25 +16,25 @@ GDNIntermediateState, GDNLayerCache, GDNReplayState, - Qwen35HybridCacheManagerV2, + Qwen35HybridCacheManager, ) from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( KDAIntermediateLayerCache, KDAIntermediateState, KDAReplayLayerCache, KDAReplayState, - KimiK3HybridCacheManagerV2, + KimiK3HybridCacheManager, ) from tensorrt_llm._torch.modules.mamba.cache_manager import ( Mamba2LayerCache, Mamba2State, - NemotronHybridCacheManagerV2, + NemotronHybridCacheManager, ReplayHistory, ) from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( PLE_CONV_STATE, PLE_NGRAM_CONTEXT, - Qwen4ExpHybridCacheManagerV2, + Qwen4ExpHybridCacheManager, Qwen4ExpPLECacheParams, ) from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2, Role @@ -59,10 +59,10 @@ @pytest.mark.parametrize( "manager_cls", [ - NemotronHybridCacheManagerV2, - Qwen35HybridCacheManagerV2, - Qwen4ExpHybridCacheManagerV2, - KimiK3HybridCacheManagerV2, + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, Glm5NextCacheManager, ], ) @@ -81,10 +81,10 @@ def test_model_managers_inherit_the_common_lifecycle(manager_cls): "manager_cls", [ MambaHybridCacheManagerV2, - NemotronHybridCacheManagerV2, - Qwen35HybridCacheManagerV2, - Qwen4ExpHybridCacheManagerV2, - KimiK3HybridCacheManagerV2, + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, Glm5NextCacheManager, ], ) @@ -128,10 +128,10 @@ def test_legacy_managers_keep_compatibility_accessors(): [ MambaHybridCacheManager, MambaHybridCacheManagerV2, - NemotronHybridCacheManagerV2, - Qwen35HybridCacheManagerV2, - Qwen4ExpHybridCacheManagerV2, - KimiK3HybridCacheManagerV2, + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, Glm5NextCacheManager, Mamba2State, ReplayHistory, @@ -187,7 +187,7 @@ def _shutdown_manager(manager, monkeypatch): """Exercise real shutdown while replacing only the native pool teardown.""" manager._recurrent_buffers = {"buffer": object()} manager._branch_snapshot_points = {"request": object()} - if isinstance(manager, Qwen4ExpHybridCacheManagerV2): + if isinstance(manager, Qwen4ExpHybridCacheManager): manager._ple_conv_states = {42: object()} manager._ple_ngram_contexts = {42: object()} @@ -198,7 +198,7 @@ def check_before_pool_shutdown(): assert not manager._recurrent_buffers assert not manager._branch_snapshot_points assert not manager.all_ssm_states and not manager.all_conv_states - if isinstance(manager, Qwen4ExpHybridCacheManagerV2): + if isinstance(manager, Qwen4ExpHybridCacheManager): assert not manager._ple_conv_states and not manager._ple_ngram_contexts pool_shutdown = Mock(side_effect=check_before_pool_shutdown) @@ -213,7 +213,7 @@ def _prepare_model_layout(manager): manager._state_layout = replace(_layout(), spec_config=manager.spec_config) manager.local_num_mamba_layers = 2 manager._global_n_groups = 2 - if isinstance(manager, Qwen4ExpHybridCacheManagerV2): + if isinstance(manager, Qwen4ExpHybridCacheManager): manager._ple_params = None manager._is_ple_draft = False manager._pretrained_config = None @@ -223,12 +223,12 @@ def _prepare_model_layout(manager): def test_qwen_managers_share_gdn_algorithm_not_model_inheritance(): states = [] - for cls in (Qwen35HybridCacheManagerV2, Qwen4ExpHybridCacheManagerV2): + for cls in (Qwen35HybridCacheManager, Qwen4ExpHybridCacheManager): manager = object.__new__(cls) manager._requested_replay = True manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) _prepare_model_layout(manager) - manager._speculative_state = manager._initialize_model_state() + manager._speculative_state = manager._initialize_spec_state() states.append(manager._speculative_state) assert all(type(state) is GDNReplayState for state in states) assert states[0] is not states[1] @@ -237,10 +237,10 @@ def test_qwen_managers_share_gdn_algorithm_not_model_inheritance(): @pytest.mark.parametrize( "manager_cls", [ - NemotronHybridCacheManagerV2, - Qwen35HybridCacheManagerV2, - Qwen4ExpHybridCacheManagerV2, - KimiK3HybridCacheManagerV2, + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, Glm5NextCacheManager, ], ) @@ -255,7 +255,7 @@ def test_model_initialization_validates_selected_replay(manager_cls): manager._state_layout, n_groups_per_rank=0, conv_state_layout="x_b_c" ) with pytest.raises(ValueError, match="requires"): - manager._initialize_model_state() + manager._initialize_spec_state() @pytest.mark.parametrize( @@ -268,7 +268,7 @@ def test_model_initialization_validates_selected_replay(manager_cls): (None, None, None), ], ) -@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManagerV2, Glm5NextCacheManager]) +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) def test_k3_initializes_only_the_selected_state( monkeypatch, requested_num_spec, selected_num_spec, decoding_type, manager_cls ): @@ -285,22 +285,22 @@ def test_k3_initializes_only_the_selected_state( replay_factory = Mock(wraps=KDAReplayState) selector = Mock(return_value=selected_num_spec) monkeypatch.setitem( - manager_cls._initialize_model_state.__globals__, + manager_cls._initialize_spec_state.__globals__, "KDAIntermediateState", intermediate_factory, ) monkeypatch.setitem( - manager_cls._initialize_model_state.__globals__, + manager_cls._initialize_spec_state.__globals__, "KDAReplayState", replay_factory, ) monkeypatch.setitem( - manager_cls._initialize_model_state.__globals__, + manager_cls._initialize_spec_state.__globals__, "get_kda_replay_num_spec", selector, ) - manager._speculative_state = manager._initialize_model_state() + manager._speculative_state = manager._initialize_spec_state() num_spec = requested_num_spec if requested_num_spec is not None else selected_num_spec if num_spec is not None: intermediate_factory.assert_not_called() @@ -339,7 +339,7 @@ def test_k3_initializes_only_the_selected_state( def test_base_manager_without_model_state(monkeypatch): manager = object.__new__(MambaHybridCacheManagerV2) manager.spec_config = None - manager._speculative_state = manager._initialize_model_state() + manager._speculative_state = manager._initialize_spec_state() assert manager._speculative_state is None manager._setup_model_state() manager._reset_context_mamba_slots(num_contexts=1) @@ -355,9 +355,9 @@ def test_base_manager_without_model_state(monkeypatch): [ MambaHybridCacheManager, MambaHybridCacheManagerV2, - Qwen35HybridCacheManagerV2, - Qwen4ExpHybridCacheManagerV2, - KimiK3HybridCacheManagerV2, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, Glm5NextCacheManager, ], ) @@ -366,7 +366,7 @@ def test_seed_accessor_is_not_a_common_manager_capability(manager_cls): def test_mamba2_seed_lifecycle_without_speculative_decoding(monkeypatch): - manager = object.__new__(NemotronHybridCacheManagerV2) + manager = object.__new__(NemotronHybridCacheManager) manager.spec_config = None manager._requested_replay = False manager._mamba_ssm_stochastic_rounding = True @@ -374,7 +374,7 @@ def test_mamba2_seed_lifecycle_without_speculative_decoding(monkeypatch): views = _state_views() manager.all_ssm_states = views.all_ssm_states manager.all_conv_states = views.all_conv_states - manager._speculative_state = manager._initialize_model_state() + manager._speculative_state = manager._initialize_spec_state() manager._setup_model_state() seeds = manager.get_mamba_ssm_rand_seed() before = seeds.clone() @@ -425,7 +425,7 @@ def test_base_manager_requires_model_for_speculative_decoding(): manager = object.__new__(MambaHybridCacheManagerV2) manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) with pytest.raises(NotImplementedError, match="model-specific"): - manager._initialize_model_state() + manager._initialize_spec_state() manager._speculative_state = None manager.local_num_mamba_layers = 1 manager._generation_state_indices = torch.empty(0, dtype=torch.int32) @@ -440,7 +440,7 @@ def test_base_manager_requires_model_for_speculative_decoding(): @pytest.mark.parametrize( "manager_cls", - [NemotronHybridCacheManagerV2, Qwen35HybridCacheManagerV2, Qwen4ExpHybridCacheManagerV2], + [NemotronHybridCacheManager, Qwen35HybridCacheManager, Qwen4ExpHybridCacheManager], ) @pytest.mark.parametrize("mode", ["plain", "intermediate", "replay"]) def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): @@ -448,14 +448,14 @@ def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): manager.spec_config = SimpleNamespace(tokens_per_gen_step=3) if mode != "plain" else None manager._requested_replay = mode == "replay" _prepare_model_layout(manager) - manager._speculative_state = manager._initialize_model_state() + manager._speculative_state = manager._initialize_spec_state() state = manager._speculative_state assert isinstance(state, (Mamba2State, GDNIntermediateState)) == (mode != "replay") assert isinstance(state, (ReplayHistory, GDNReplayState)) == (mode == "replay") manager._state_layout = object() manager.local_num_mamba_layers = 0 - if manager_cls is Qwen4ExpHybridCacheManagerV2: + if manager_cls is Qwen4ExpHybridCacheManager: manager._ple_conv_states = {} manager._ple_ngram_contexts = {} manager.mamba_layer_offsets = {42: 0} @@ -522,7 +522,7 @@ def test_model_owns_speculative_state_lifecycle(monkeypatch, manager_cls, mode): def test_ple_adds_buffers_without_overriding_snapshot_or_budget_solver(): - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager._mamba_layer_mask = [True, True] manager.mamba_pp_layers = [0, 1] manager.ssm_bytes, manager.conv_bytes = 128, 288 @@ -551,7 +551,7 @@ def test_warmup_cleanup_clears_registered_persistent_roles_only( mamba_cache_manager_v2 as module, ) - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager._mamba_layer_mask = [False, True] manager.pp_layers = manager.mamba_pp_layers = [1] manager.layer_offsets = manager.mamba_layer_offsets = {1: 0} @@ -649,7 +649,7 @@ def test_gdn_bind_accepts_affine_and_indirect_views(affine): assert (state._state_descriptors is not None) != affine assert state.old_x.shape == (3, 5, 2, 16, 2, 4) owners = [] - for cls in (Qwen35HybridCacheManagerV2, Qwen4ExpHybridCacheManagerV2): + for cls in (Qwen35HybridCacheManager, Qwen4ExpHybridCacheManager): owner = object.__new__(cls) owner._speculative_state = state owners.append(owner) @@ -977,17 +977,17 @@ def test_algorithm_methods_do_not_depend_on_manager(cls): @pytest.mark.parametrize( "manager_cls", [ - NemotronHybridCacheManagerV2, - Qwen35HybridCacheManagerV2, - Qwen4ExpHybridCacheManagerV2, - KimiK3HybridCacheManagerV2, + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, Glm5NextCacheManager, ], ) def test_model_compatibility_properties_follow_owned_state(manager_cls): manager = object.__new__(manager_cls) manager._speculative_state = GDNIntermediateState() - if manager_cls is KimiK3HybridCacheManagerV2: + if manager_cls is KimiK3HybridCacheManager: manager._kda_replay = None fields = { "intermediate_state_indices": "intermediate_indices", @@ -1006,7 +1006,7 @@ def test_model_compatibility_properties_follow_owned_state(manager_cls): def test_k3_resolves_transfer_requests_before_calling_algorithm(): - manager = object.__new__(KimiK3HybridCacheManagerV2) + manager = object.__new__(KimiK3HybridCacheManager) manager._request_id_to_state_index = {101: 3, 202: 1} replay = KDAReplayState(2) replay.prev_num_accepted_tokens = torch.zeros(5, dtype=torch.int32) @@ -1023,7 +1023,7 @@ def test_k3_resolves_transfer_requests_before_calling_algorithm(): @pytest.mark.parametrize("draft_lengths", [[], [0, 0], [1, 0], [1, 1]]) def test_k3_host_acceptance_handles_empty_mixed_and_completed_requests(draft_lengths): - manager = object.__new__(KimiK3HybridCacheManagerV2) + manager = object.__new__(KimiK3HybridCacheManager) replay = KDAReplayState(2) replay.prev_num_accepted_tokens = torch.tensor([5, 6], dtype=torch.int32) manager._kda_replay = replay @@ -1049,7 +1049,7 @@ def test_k3_host_acceptance_handles_empty_mixed_and_completed_requests(draft_len def _slot_view_manager( monkeypatch, *, scale=6, role_count=2, page_shape=(4, 2, 3), dtype=torch.float32 ): - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.layer_offsets = {42: 0} page_elements = 1 for dim in page_shape: @@ -1071,7 +1071,7 @@ def _slot_view_manager( get_page_index_converter=lambda layer, role: converter, get_page_index_upper_bound=lambda layer, role: 3 * scale - layer_offset, ) - globals_ = Qwen4ExpHybridCacheManagerV2._get_view_by_role_and_layer.__globals__ + globals_ = Qwen4ExpHybridCacheManager._get_view_by_role_and_layer.__globals__ monkeypatch.setitem(globals_, "TensorWrapper", lambda address, dtype, shape: (address, shape)) def convert(wrapper): @@ -1124,7 +1124,7 @@ def test_slot_role_view_rejects_incompatible_physical_layout(monkeypatch, invali def test_qwen4_position_uses_model_owned_slot_view(monkeypatch): - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.qsa_position_layer_id = 42 manager.tokens_per_block = 4 manager._get_view_by_role_and_layer = Mock( @@ -1140,8 +1140,8 @@ def test_qwen4_position_uses_model_owned_slot_view(monkeypatch): @pytest.mark.parametrize("kv_layout", ["NHD", "HND"]) def test_qwen4_main_kv_view_supplies_only_logical_geometry(monkeypatch, kv_layout): - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) - namespace = Qwen4ExpHybridCacheManagerV2.get_buffers.__globals__ + manager = object.__new__(Qwen4ExpHybridCacheManager) + namespace = Qwen4ExpHybridCacheManager.get_buffers.__globals__ manager._qsa_enabled = True manager.dtype = next(iter(namespace["QSA_SPARSE_KV_CACHE_DTYPES"])) manager.kv_cache_type = namespace["CacheTypeCpp"].SELF @@ -1167,7 +1167,7 @@ def test_qwen4_main_kv_view_supplies_only_logical_geometry(monkeypatch, kv_layou @pytest.mark.parametrize("invalid", [None, "pool", "scale", "expansion"]) def test_qwen4_attention_pool_layout_validates_backend_contract(invalid): - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.layer_offsets = {42: 0} manager.qsa_position_layer_id = 42 manager.qsa_sparse_layer_ids = [42] @@ -1189,7 +1189,7 @@ def test_qwen4_attention_pool_layout_validates_backend_contract(invalid): def test_qwen4_requires_one_mapping_and_skips_nonlocal_layers(): - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.qsa_position_layer_id = 42 manager.qsa_sparse_layer_ids = [42, 43, 99] manager.layer_offsets = {42: 0, 43: 1} @@ -1242,9 +1242,9 @@ def create_pool(self, *args, **kwargs): get_page_index_upper_bound=lambda layer, role: 5, ) - class ObservedQwen4(Qwen4ExpHybridCacheManagerV2): - def _initialize_model_state(self): - state = super()._initialize_model_state() + class ObservedQwen4(Qwen4ExpHybridCacheManager): + def _initialize_spec_state(self): + state = super()._initialize_spec_state() events.append("validate") return state diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index 63920b74dd38..746b44d07e47 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -521,7 +521,7 @@ def __init__(self, *args, **kwargs): "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", return_value=verify_kernel is True, ): - state = manager._initialize_model_state() + state = manager._initialize_spec_state() assert manager.use_kda_replay_update is (verify_kernel is True) if verify_kernel is True: assert state.num_speculative_tokens == 3 diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index 69f78472c639..8df5621ed9b0 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -16,7 +16,7 @@ from tensorrt_llm._torch.attention.backends.sparse.qsa.module import QSASparseHooks from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManagerV2 +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import MambaKVCacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 @@ -222,7 +222,7 @@ def test_qsa_empty_batch_keeps_the_regular_backend_path() -> None: def test_qsa_metadata_allows_a_pp_rank_without_local_sparse_layers(monkeypatch) -> None: monkeypatch.setattr(TrtllmAttentionMetadata, "__post_init__", lambda self: None) - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.qsa_position_layer_id = None metadata = object.__new__(QSAAttentionMetadata) metadata.kv_cache_manager = manager @@ -253,7 +253,7 @@ def test_qsa_hybrid_routes_to_sparse_v2_cache_manager(monkeypatch) -> None: manager_cls = get_kv_cache_manager_cls(model_config, kv_cache_config) - assert manager_cls is Qwen4ExpHybridCacheManagerV2 + assert manager_cls is Qwen4ExpHybridCacheManager def test_qsa_hybrid_rejects_kv_cache_manager_v1(monkeypatch) -> None: @@ -337,7 +337,7 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( manager = _create_kv_cache_manager( model_engine=None, - kv_cache_manager_cls=Qwen4ExpHybridCacheManagerV2, + kv_cache_manager_cls=Qwen4ExpHybridCacheManager, mapping=SimpleNamespace(enable_attention_dp=False), kv_cache_config=KvCacheConfig(use_kv_cache_manager_v2=True), tokens_per_block=128, @@ -362,7 +362,7 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( def test_qwen4_manager_also_supports_dense_attention(monkeypatch) -> None: monkeypatch.setattr(MambaHybridCacheManagerV2, "__init__", lambda self, *args, **kwargs: None) - manager = Qwen4ExpHybridCacheManagerV2(layer_mask=[True]) + manager = Qwen4ExpHybridCacheManager(layer_mask=[True]) assert not manager._qsa_enabled assert manager._extra_buffers_per_layer(tokens_per_block=32) == {} @@ -371,7 +371,7 @@ def test_qwen4_manager_also_supports_dense_attention(monkeypatch) -> None: def test_qsa_cache_manager_delegates_regular_kv_layout( monkeypatch: pytest.MonkeyPatch, dtype: DataType ) -> None: - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.dtype = dtype manager._qsa_enabled = True sentinel = torch.empty(0, dtype=torch.int8) diff --git a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py index 8f9b998c4168..7dac22c46df8 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -553,7 +553,7 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state( ) -> None: from tensorrt_llm._torch.configs import Qwen4ExpTextConfig from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( - Qwen4ExpHybridCacheManagerV2, + Qwen4ExpHybridCacheManager, extract_qwen4_exp_ple_cache_params, ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig @@ -573,10 +573,10 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state( ), ), } - with_ple = Qwen4ExpHybridCacheManagerV2.get_cache_size_per_token( + with_ple = Qwen4ExpHybridCacheManager.get_cache_size_per_token( SimpleNamespace(pretrained_config=config, quant_config=None), **common ) - without_ple = Qwen4ExpHybridCacheManagerV2.get_cache_size_per_token( + without_ple = Qwen4ExpHybridCacheManager.get_cache_size_per_token( SimpleNamespace(pretrained_config=no_ple_config, quant_config=None), **common ) @@ -593,7 +593,7 @@ def test_ple_states_use_v2_lifecycle_buffers(monkeypatch) -> None: from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( PLE_CONV_STATE, PLE_NGRAM_CONTEXT, - Qwen4ExpHybridCacheManagerV2, + Qwen4ExpHybridCacheManager, ) ngram_context = torch.full((12, 2), 11, dtype=torch.int64) @@ -615,8 +615,8 @@ def fake_get_state_buffer( return conv_state raise AssertionError(f"unexpected role {role}") - monkeypatch.setattr(Qwen4ExpHybridCacheManagerV2, "_get_state_buffer", fake_get_state_buffer) - manager = object.__new__(Qwen4ExpHybridCacheManagerV2) + monkeypatch.setattr(Qwen4ExpHybridCacheManager, "_get_state_buffer", fake_get_state_buffer) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager._ple_layer_ids = [1] manager._ple_ngram_context_shape = [2] manager._ple_conv_state_shape = [16, 6] diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index ba073c645058..b911ab6893ff 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -37,10 +37,10 @@ from tensorrt_llm._torch.disaggregation.native.mixers.ssm import peer from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 -from tensorrt_llm._torch.modules.fla.cache_manager import Qwen35HybridCacheManagerV2 -from tensorrt_llm._torch.modules.mamba.cache_manager import NemotronHybridCacheManagerV2 +from tensorrt_llm._torch.modules.fla.cache_manager import Qwen35HybridCacheManager +from tensorrt_llm._torch.modules.mamba.cache_manager import NemotronHybridCacheManager from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( - Qwen4ExpHybridCacheManagerV2, + Qwen4ExpHybridCacheManager, Qwen4ExpPLECacheParams, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( @@ -232,11 +232,11 @@ def _create_managers( if not use_v2: manager_cls = MixedMambaHybridCacheManager elif with_ple: - manager_cls = Qwen4ExpHybridCacheManagerV2 + manager_cls = Qwen4ExpHybridCacheManager elif conv_state_layout == "q_k_v": - manager_cls = Qwen35HybridCacheManagerV2 + manager_cls = Qwen35HybridCacheManager else: - manager_cls = NemotronHybridCacheManagerV2 + manager_cls = NemotronHybridCacheManager manager_kwargs = ( { "is_disagg": True, @@ -304,7 +304,7 @@ def _zero_mamba_states(manager): for layer_idx in _mamba_layer_ids(manager): manager.get_conv_states(layer_idx).zero_() manager.get_ssm_states(layer_idx).zero_() - if isinstance(manager, Qwen4ExpHybridCacheManagerV2): + if isinstance(manager, Qwen4ExpHybridCacheManager): for state in manager._ple_conv_states.values(): state.zero_() for state in manager._ple_ngram_contexts.values(): From 65e898bdbd7efe62d2c434d7b331d2d0e1cfca1d Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 6 Oct 2026 14:25:16 +0000 Subject: [PATCH 14/25] [None][fix] Use Qwen4 PLE roles in extractor tests (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../unittest/disaggregated/test_extractor.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index ec0a638524cb..129ba92f35ac 100644 --- a/tests/unittest/disaggregated/test_extractor.py +++ b/tests/unittest/disaggregated/test_extractor.py @@ -34,11 +34,12 @@ get_physical_pool, get_unique_layers, ) -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - MambaHybridCacheManagerV2, - MambaRole, +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + PLE_NGRAM_CONTEXT, + Qwen4ExpHybridCacheManager, ) +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaRole from tensorrt_llm._torch.pyexecutor.resource_manager import ( CacheTypeCpp, DataType, @@ -746,7 +747,7 @@ def test_mamba_layer_group_serialization(): def _make_fake_v2_mamba_manager(*, ple_on_layer=1): - """Duck-typed MambaHybridCacheManagerV2 whose descriptors hold two life cycles. + """Qwen4 manager stub whose descriptors hold two life cycles. Life cycle 0 is recurrent (layers 0 and 1): SSM and conv state coalesce into one 64-byte pool in ``[ssm, conv]`` order per layer, and PLE n-gram @@ -783,7 +784,7 @@ def buf(layer_id, role): ), SimpleNamespace( single_buffer_size=32, - buffer_ids=[buf(ple_on_layer, MambaRole.PLE_NGRAM_CONTEXT)], + buffer_ids=[buf(ple_on_layer, PLE_NGRAM_CONTEXT)], ), ], ) @@ -809,7 +810,7 @@ def buf(layer_id, role): ), ) - manager = object.__new__(MambaHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) manager.impl = SimpleNamespace( layer_grouping=((0, 1), (2,)), init_config=SimpleNamespace( @@ -863,7 +864,7 @@ def test_v2_builder_describes_recurrent_life_cycle_from_descriptors(): assert set(by_role) == { MAMBA_SSM_ROLE, MAMBA_CONV_ROLE, - frozenset({str(MambaRole.PLE_NGRAM_CONTEXT)}), + frozenset({str(PLE_NGRAM_CONTEXT)}), } ssm = by_role[MAMBA_SSM_ROLE] @@ -882,7 +883,7 @@ def test_v2_builder_describes_recurrent_life_cycle_from_descriptors(): assert conv.bytes_per_head is None assert [tuple(int(x) for x in e) for e in conv.buffer_entries] == [(0, 64, 64), (1, 192, 64)] - ple = by_role[frozenset({str(MambaRole.PLE_NGRAM_CONTEXT)})] + ple = by_role[frozenset({str(PLE_NGRAM_CONTEXT)})] assert ple.pool_idx == 1 assert ple.mapper_kind == MapperKind.REPLICATED assert ple.bytes_per_layer == 32 @@ -893,7 +894,7 @@ def test_v2_builder_describes_recurrent_life_cycle_from_descriptors(): extractor = KVRegionExtractorV1(page_table) view_index = {pv.pool_role: idx for idx, pv in enumerate(recurrent.pool_views)} conv_idx = view_index[MAMBA_CONV_ROLE] - ple_idx = view_index[frozenset({str(MambaRole.PLE_NGRAM_CONTEXT)})] + ple_idx = view_index[frozenset({str(PLE_NGRAM_CONTEXT)})] assert extractor.extract_slot(2, 0, conv_idx).memory.ptrs.tolist() == [ 0x10000 + 2 * 256 + 64, 0x10000 + 2 * 256 + 192, From 6e9aecdcfc2a35620bb7c2676f1729d63b91e6e8 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 14 Oct 2026 13:59:33 +0000 Subject: [PATCH 15/25] fix: reserve replay scratch and gate QSA cache allocation (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../sparse/glm_kpool/cache_manager.py | 10 +- .../_torch/modules/kimi_kda/cache_manager.py | 62 ++++++- .../_torch/modules/qwen4_exp/cache_manager.py | 4 +- .../kv_cache/kv_cache_manager_v2.py | 47 ++++- .../mamba_cache_manager_v2.py | 11 ++ .../kv_cache/test_mamba_cache_manager.py | 14 ++ .../kv_cache/test_mamba_model_managers.py | 166 ++++++++++++++++++ .../modeling/test_qsa_runtime_wiring.py | 23 ++- 8 files changed, 319 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py index dee85fd4ba2d..9b207a947c55 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py @@ -126,6 +126,11 @@ def _extra_scratch_bytes_per_slot(self) -> int: for layer_id in self.mamba_pp_layers ) + @classmethod + @override + def _extra_state_bytes_per_rank(cls, model_config, mapping: Mapping, **kwargs) -> int: + return KDAReplayState.estimate_bytes_per_slot(model_config, mapping, **kwargs) + @override def _on_state_slots_relocated(self, old_slots: list[int], new_slots: list[int]) -> None: if self._kda_replay is None: @@ -249,9 +254,10 @@ def _attention_cache_bytes_per_token(self) -> int: for local_layer_idx in range(self.num_local_layers) ) - @staticmethod + @classmethod @override def get_cache_size_per_token( + cls, model_config: ModelConfig, mapping: Mapping, *, @@ -261,7 +267,7 @@ def get_cache_size_per_token( max_seq_len: int | None = None, **kwargs, ) -> tuple[int, int]: - slope, fixed_cost = MambaHybridCacheManagerV2.get_cache_size_per_token( + slope, fixed_cost = super().get_cache_size_per_token( model_config, mapping, max_batch_size=max_batch_size, diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index 7a3c6c63bea3..3bb659b19d59 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -19,6 +19,8 @@ MambaAcceptanceBatch, MambaLayerCache, MambaStateLayout, + _get_local_mamba_cache_layout, + _mamba_effective_tp_size, _promote_intermediate_states, _stack_state_views, ) @@ -26,6 +28,7 @@ MambaHybridCacheManagerV2, ) from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping _MAMBA_SSM_SEED_MASK = (1 << 62) - 1 _MAMBA_SSM_UINT64_MASK = (1 << 64) - 1 @@ -327,25 +330,65 @@ def validate(self, context: MambaStateLayout) -> None: "KDA replay caches require equal [Q | K | V] convolution-state sections" ) - def bytes_per_slot(self, context: MambaStateLayout, layer_id: int) -> int: - if not context.mamba_pp_layers: - return 0 - section_dim = context.conv_section_dims[0] - committed_window = context.conv_state_shape[1] + def _layer_bytes_per_slot(self, section_dim: int, committed_window: int, num_heads: int) -> int: extended_window = committed_window + self.num_speculative_tokens - num_heads = context.ssm_state_shape[0] float_elements = ( 3 * section_dim * extended_window + self.num_speculative_tokens * 3 * section_dim + self.num_speculative_tokens * section_dim + self.num_speculative_tokens * _kda_beta_cache_padded_heads(num_heads) ) + return float_elements * 4 + + def bytes_per_slot(self, context: MambaStateLayout, layer_id: int) -> int: + if not context.mamba_pp_layers: + return 0 shared_bytes = 0 if layer_id == context.mamba_pp_layers[0]: shared_bytes = 4 if self._stochastic_rounding: shared_bytes += 8 - return float_elements * 4 + shared_bytes + return ( + self._layer_bytes_per_slot( + context.conv_section_dims[0], + context.conv_state_shape[1], + context.ssm_state_shape[0], + ) + + shared_bytes + ) + + @classmethod + def estimate_bytes_per_slot(cls, model_config, mapping: Mapping, **kwargs) -> int: + """Estimate rank-local replay storage using the runtime replay gate and geometry.""" + spec_config = kwargs.get("spec_config") + num_spec = kwargs.get("kda_replay_num_spec") + if num_spec is None: + num_spec = get_kda_replay_num_spec(spec_config, manager_supports_replay=True) + if num_spec is None: + return 0 + params, local_layers, _ = _get_local_mamba_cache_layout( + model_config, + mapping, + spec_config=spec_config, + is_draft=kwargs.get("is_draft", False), + use_separate_draft_kv_cache=kwargs.get("use_separate_draft_kv_cache", False), + ) + if not local_layers: + return 0 + tp_size = _mamba_effective_tp_size(mapping) + quant_config = model_config.quant_config + stochastic_rounding = kwargs.get( + "mamba_ssm_stochastic_rounding", + getattr(quant_config, "mamba_ssm_stochastic_rounding", False) + and params.mamba_ssm_cache_dtype == torch.float16, + ) + state = cls(num_spec) + layer_bytes = state._layer_bytes_per_slot( + params.n_groups * params.state_size // tp_size, + params.conv_kernel - 1, + params.num_heads // tp_size, + ) + return local_layers * layer_bytes + 4 + (8 if stochastic_rounding else 0) def bind( self, @@ -641,6 +684,11 @@ def _extra_scratch_bytes_per_slot(self) -> int: for layer_id in self.mamba_pp_layers ) + @classmethod + @override + def _extra_state_bytes_per_rank(cls, model_config, mapping: Mapping, **kwargs) -> int: + return KDAReplayState.estimate_bytes_per_slot(model_config, mapping, **kwargs) + @override def _on_state_slots_relocated(self, old_slots: list[int], new_slots: list[int]) -> None: if self._kda_replay is None: diff --git a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py index d26564a49f19..b996ad00c0ac 100644 --- a/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -201,7 +201,9 @@ def __init__( self.qsa_index_dim = self.qsa_index_kv_heads = 0 self.qsa_sparse_layer_ids = [] self.qsa_position_layer_id = None - self._qsa_enabled = sparse_attention_config is not None + self._qsa_enabled = ( + sparse_attention_config is not None and sparse_attention_config.algorithm == "qsa" + ) if self._qsa_enabled: if layer_mask is None: raise ValueError("QSA cache allocation requires the full-attention layer mask") diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 6d2bb516fb75..78f978bc0594 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1724,7 +1724,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: else self.event_manager ) if not has_host_cache_tier: - candidate = KVCacheManagerPy( + candidate = self._allocate_cache_pools( config, event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), @@ -1733,7 +1733,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: init_error: Optional[Exception] = None local_init_status = _KVCacheManagerInitStatus.KEEP_HOST try: - candidate = KVCacheManagerPy( + candidate = self._allocate_cache_pools( config, event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), @@ -1773,7 +1773,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: if not isinstance(tier, HostCacheTierConfig) ], ) - candidate = KVCacheManagerPy( + candidate = self._allocate_cache_pools( config, event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), @@ -1799,6 +1799,14 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: ) assert candidate is not None + if self._get_external_cache_bytes(candidate): + config = replace( + config, + cache_tiers=[ + replace(config.cache_tiers[0], quota=candidate.get_quota(GPU_LEVEL)), + *config.cache_tiers[1:], + ], + ) self.kv_cache_manager_py_config = config self.impl = candidate self.can_evict = len(config.cache_tiers) > 1 @@ -3067,6 +3075,39 @@ def _build_base_config( initial_pool_ratio=kv_cache_config.pool_ratio, ) + def _get_external_cache_bytes(self, cache: KVCacheManagerPy) -> int: + """Return GPU storage owned outside the native pools for this capacity.""" + return 0 + + def _allocate_cache_pools( + self, config: KVCacheManagerConfigPy, *, event_manager, cold_page_codec + ) -> KVCacheManagerPy: + """Reserve external storage before publishing pool addresses and capacities.""" + cache = KVCacheManagerPy( + config, event_manager=event_manager, cold_page_codec=cold_page_codec + ) + total_quota = config.cache_tiers[0].quota + previous_pool_quota = total_quota + external_bytes = self._get_external_cache_bytes(cache) + while external_bytes and cache.get_quota(GPU_LEVEL) + external_bytes > total_quota: + # Scale the combined footprint, then recheck the actual capacity: + # native allocation granularity and minimum slots affect rounding. + native_bytes = cache.get_quota(GPU_LEVEL) + pool_quota = native_bytes * total_quota // (native_bytes + external_bytes) + if ( + pool_quota <= 0 + or pool_quota >= previous_pool_quota + or not cache.resize(GPU_LEVEL, pool_quota) + ): + cache.shutdown() + raise ValueError( + "The GPU cache quota cannot accommodate native pools and " + f"external cache storage ({external_bytes} bytes)." + ) + previous_pool_quota = pool_quota + external_bytes = self._get_external_cache_bytes(cache) + return cache + def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """Customize the general cache config for a specialized cache manager.""" return config diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 0d22ddac3f2d..56a1e6ed8358 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -70,6 +70,7 @@ TokenIdExt, _KVCache, ) +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManager as KVCacheManagerPy from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy from .common import ( @@ -400,6 +401,16 @@ def get_disagg_role_layouts(self) -> dict[DataRole, RoleLayout]: def _extra_scratch_bytes_per_slot(self) -> int: return 0 + @override + def _get_external_cache_bytes(self, cache: KVCacheManagerPy) -> int: + bytes_per_slot = self._extra_scratch_bytes_per_slot() + if not bytes_per_slot or not self.mamba_pp_layers: + return 0 + layer_id = LayerId(self.layer_offsets[self.mamba_pp_layers[0]]) + scale = cache.get_page_index_scale(layer_id, MambaRole.SSM_STATE) + pages = cache.get_page_index_upper_bound(layer_id, MambaRole.SSM_STATE) + return ((pages + scale - 1) // scale) * bytes_per_slot + def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConfig]: """Declare all buffers sharing one recurrent layer's slot lifecycle.""" return [ diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 5256d75dba9e..4062242b322c 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Regression tests for Python, Cpp, and V2 Mamba cache managers.""" +import math import os from types import SimpleNamespace from unittest.mock import MagicMock @@ -88,6 +89,7 @@ ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + GPU_LEVEL, AttentionLayerConfig, BatchDesc, BufferConfig, @@ -3682,6 +3684,18 @@ def test_v2_kda_replay_policy_allocates_logical_slot_caches(): assert policy.intermediate_indices.shape == (4,) assert not hasattr(mgr, "get_replay_state_update_metadata") + total_quota = math.ceil( + mgr._get_quota_from_max_tokens(mgr.kv_cache_config.max_tokens) + / mgr.kv_cache_config.max_util_for_resume + ) + replay_bytes = sum( + buffer.untyped_storage().nbytes() + for buffer in (*policy._layer_replay_buffers(), policy.prev_num_accepted_tokens) + ) + native_bytes = mgr.impl.get_quota(GPU_LEVEL) + assert native_bytes + replay_bytes <= total_quota + assert mgr.kv_cache_manager_py_config.cache_tiers[0].quota == native_bytes + layer_cache = mgr.mamba_layer_cache(0) assert not hasattr(layer_cache, "intermediate_ssm") assert not hasattr(layer_cache, "intermediate_conv_window") diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index cd014b8014e6..ec5be1f5ea57 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -46,11 +46,13 @@ MambaRole, MambaStateLayout, ) +from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, BufferConfig, GpuCacheTierConfig, + HostCacheTierConfig, KVCacheManagerConfig, LayerId, ) @@ -720,6 +722,170 @@ def test_kda_scratch_cost_matches_allocated_slot_buffers(): assert actual == estimated * layout.slot_capacity +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("pool_share", [0.25, 0.75]) +@pytest.mark.parametrize("quota_source", ["max_tokens", "max_gpu_total_bytes"]) +def test_kda_pool_allocation_reserves_actual_replay_capacity( + monkeypatch, manager_cls, replay, pool_share, quota_source +): + from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as module + + layout = _layout() + manager = object.__new__(manager_cls) + manager._state_layout = layout + manager.mamba_pp_layers = list(layout.mamba_pp_layers) + manager.layer_offsets = {0: 0, 1: 1} + manager._kda_replay = KDAReplayState(2) if replay else None + manager.ssm_bytes, manager.conv_bytes = 128, 288 + manager.max_batch_size = layout.max_batch_size + manager.mapping = layout.mapping + manager._num_reserved_dummy_slots = 1 + manager.kv_cache_config = KvCacheConfig(enable_block_reuse=False) + persistent_bytes_per_slot = 2 * (128 + 288) + total_quota = 1 << 20 + if quota_source == "max_tokens": + monkeypatch.setattr( + KVCacheManagerV2, "_get_quota_from_max_tokens", lambda self, tokens: tokens * 64 + ) + total_quota = manager._get_quota_from_max_tokens(16384) + + class Pool: + def __init__(self, config, **kwargs): + self.quota = config.cache_tiers[0].quota + self.resize_calls = 0 + + @property + def slots(self): + return int(self.quota * pool_share) // persistent_bytes_per_slot + + def get_quota(self, level): + return self.quota + + def get_page_index_scale(self, layer, role): + return 2 + + def get_page_index_upper_bound(self, layer, role): + return self.slots * 2 - 1 + + def resize(self, level, quota): + self.resize_calls += 1 + self.quota = quota + return True + + monkeypatch.setattr(module, "KVCacheManagerPy", Pool) + config = KVCacheManagerConfig( + tokens_per_block=32, + layers=[], + cache_tiers=[GpuCacheTierConfig(quota=total_quota), HostCacheTierConfig(quota=12345)], + ) + pool = manager._allocate_cache_pools(config, event_manager=None, cold_page_codec=None) + # The native pool can have many more slots than the live batch needs. + assert pool.slots > layout.max_batch_size + if replay: + state = manager._kda_replay + views = _state_views(slots=pool.slots) + state.bind( + layout.with_slot_capacity(pool.slots), views.all_ssm_states, views.all_conv_states + ) + scratch_bytes = sum( + buffer.untyped_storage().nbytes() + for buffer in (*state._layer_replay_buffers(), state.prev_num_accepted_tokens) + ) + assert pool.resize_calls > 0 + assert pool.quota + scratch_bytes <= total_quota + assert manager._get_external_cache_bytes(pool) == scratch_bytes + else: + assert pool.resize_calls == 0 + assert pool.quota == total_quota + # Host fallback retries must retain the original total GPU budget. + assert config.cache_tiers[0].quota == total_quota + assert config.cache_tiers[1].quota == 12345 + + +@pytest.mark.parametrize("resize_succeeds", [False, True]) +def test_external_cache_reservation_rejects_unusable_quota(monkeypatch, resize_succeeds): + from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as module + + manager = object.__new__(KVCacheManagerV2) + manager._get_external_cache_bytes = lambda cache: 128 + cache = Mock() + cache.get_quota.return_value = 1024 + cache.resize.return_value = resize_succeeds + monkeypatch.setattr(module, "KVCacheManagerPy", lambda *args, **kwargs: cache) + config = KVCacheManagerConfig( + tokens_per_block=32, layers=[], cache_tiers=[GpuCacheTierConfig(quota=1024)] + ) + # A backend that reports success without releasing memory must fail too. + with pytest.raises(ValueError, match="external cache storage"): + manager._allocate_cache_pools(config, event_manager=None, cold_page_codec=None) + cache.shutdown.assert_called_once() + + +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) +@pytest.mark.parametrize("stochastic_rounding", [False, True]) +@pytest.mark.parametrize("replay", [False, True]) +@pytest.mark.parametrize("reuse", [False, True]) +def test_kda_affine_estimator_includes_replay_scratch( + monkeypatch, manager_cls, stochastic_rounding, replay, reuse +): + from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import cache_manager as glm_module + from tensorrt_llm._torch.modules.kimi_kda import cache_manager as kda_module + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import common + + layout = _layout() + params = SimpleNamespace( + n_groups=2, + state_size=4, + conv_kernel=4, + num_heads=2, + mamba_ssm_cache_dtype=torch.float16, + get_states_bytes_per_layer=lambda mapping: 416, + get_layer_masks=lambda **kwargs: ([True, True], [False, False]), + ) + + def local_layout(*args, **kwargs): + return params, 2, 0 + + monkeypatch.setattr(common, "_get_local_mamba_cache_layout", local_layout) + monkeypatch.setattr(kda_module, "_get_local_mamba_cache_layout", local_layout) + monkeypatch.setattr( + kda_module, "get_kda_replay_num_spec", lambda *args, **kwargs: 2 if replay else None + ) + monkeypatch.setattr(glm_module, "extract_mamba_kv_cache_params", lambda *args, **kwargs: params) + monkeypatch.setattr(glm_module, "get_pp_layers", lambda *args, **kwargs: ([0, 1], None)) + config = SimpleNamespace( + pretrained_config=SimpleNamespace(index_head_dim=128), + quant_config=SimpleNamespace(mamba_ssm_stochastic_rounding=stochastic_rounding), + ) + monkeypatch.setattr(glm_module, "unwrap_glm5_next_text_config", lambda config: config) + kv_config = KvCacheConfig(enable_block_reuse=reuse) + cost = manager_cls.get_cache_size_per_token( + config, + layout.mapping, + max_batch_size=layout.max_batch_size, + kv_cache_config=kv_config, + spec_config=layout.spec_config, + ) + scratch_per_slot = 0 + if replay: + state = KDAReplayState(2, stochastic_rounding=stochastic_rounding) + state.bind(layout, _state_views().all_ssm_states, _state_views().all_conv_states) + buffers = [*state._layer_replay_buffers(), state.prev_num_accepted_tokens] + if stochastic_rounding: + buffers.append(state.rand_seed) + scratch_per_slot = ( + sum(buffer.untyped_storage().nbytes() for buffer in buffers) // layout.slot_capacity + ) + dummy_slots = common._get_num_cuda_graph_padding_dummy_slots( + layout.spec_config, layout.max_batch_size + ) + state_bytes = 2 * 416 + scratch_per_slot + interval = kv_config.mamba_state_config.periodic_snapshot_interval + slope = (state_bytes + interval - 1) // interval if reuse else 0 + assert cost == (slope, (layout.max_batch_size + dummy_slots) * state_bytes) + + @pytest.mark.parametrize("layer", [0, 1]) @pytest.mark.parametrize( "kind", ["intermediate", "kda_intermediate", "mamba2", "replay", "gdn", "kda"] diff --git a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index 8df5621ed9b0..6639bc521359 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -21,7 +21,11 @@ from tensorrt_llm._torch.pyexecutor.config_utils import MambaKVCacheParams from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm.bindings import DataType -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, QSASparseAttentionConfig +from tensorrt_llm.llmapi.llm_args import ( + KvCacheConfig, + QSASparseAttentionConfig, + SkipSoftmaxAttentionConfig, +) def _sparse_params() -> QSASparseParams: @@ -293,8 +297,13 @@ def test_qsa_rejects_non_hybrid_models(monkeypatch) -> None: ) -def test_qsa_cache_manager_uses_resolved_index_geometry( +@pytest.mark.parametrize( + "sparse_config", + [None, QSASparseAttentionConfig(), SkipSoftmaxAttentionConfig(threshold_scale_factor=1.0)], +) +def test_qwen4_cache_manager_uses_only_qsa_index_geometry( monkeypatch, + sparse_config, ) -> None: from tensorrt_llm._torch.pyexecutor import _util @@ -344,7 +353,7 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( max_seq_len=4096, max_batch_size=8, spec_config=None, - sparse_attention_config=QSASparseAttentionConfig(), + sparse_attention_config=sparse_config, max_num_tokens=1024, max_beam_width=1, kv_connector_manager=None, @@ -356,8 +365,12 @@ def test_qsa_cache_manager_uses_resolved_index_geometry( is_draft=False, ) - assert manager.qsa_index_dim == 96 - assert manager.qsa_index_kv_heads == 1 + qsa_enabled = isinstance(sparse_config, QSASparseAttentionConfig) + assert manager._qsa_enabled == qsa_enabled + assert manager.qsa_index_dim == (96 if qsa_enabled else 0) + assert manager.qsa_index_kv_heads == (1 if qsa_enabled else 0) + if not qsa_enabled: + assert manager._extra_buffers_per_layer(tokens_per_block=128) == {} def test_qwen4_manager_also_supports_dense_attention(monkeypatch) -> None: From 1cf2831564f858ad01a5ae8c842a450702d0737c Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:26:19 +0000 Subject: [PATCH 16/25] [None][fix] Correct hybrid replay accounting and lifecycle (By Agent) Keep model scratch out of native token quotas and reserve it against the physical GPU byte limit using actual pool capacity. Handle native allocation granularity and minimum slots, and reconstruct native GPU tier configs without dataclasses.replace. Move KDA request orchestration into the Kimi and GLM managers, preserving state algorithms independent of managers. Cover disabled and periodic snapshots, rounded pool budgets, and both model lifecycles. Validation: 444 pytest cases passed on B200 using the current CI image and compiled artifacts, including the prior MoE failure. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../sparse/glm_kpool/cache_manager.py | 39 ++++++++- .../_torch/modules/kimi_kda/cache_manager.py | 79 ++++++++--------- .../kv_cache/kv_cache_manager_v2.py | 65 +++++++++----- .../mamba_cache_manager_v2.py | 12 ++- .../kv_cache/test_mamba_cache_manager.py | 41 +++++---- .../kv_cache/test_mamba_model_managers.py | 87 ++++++++++++++++--- 6 files changed, 215 insertions(+), 108 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py index 39d248d7065a..91e3598f070b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py @@ -133,15 +133,46 @@ def _build_cache_config(self, config): @override def suspend_request(self, request) -> None: - if self._kda_replay is not None: - self._kda_replay.prepare_suspend(self, request) + replay = self._kda_replay + request_id = request.py_request_id + cache = self.kv_cache_map.get(request_id) + slot = self._request_id_to_state_index.get(request_id, -1) + if replay is not None and slot >= 0 and cache is not None and cache.is_active: + if request_id in replay.acceptance_pending: + raise RuntimeError( + f"Request {request_id} is being suspended before update_resources " + "recorded the acceptance of its last verify, so its pending KDA drafts are unknown" + ) + if request_id in replay.unseeded_requests: + replay.unseeded_requests.remove(request_id) + else: + with replay.on_cache_stream(self._stream): + replay.commit_pending_drafts(slot) + del self._request_id_to_state_index[request_id] + replay.suspended_requests.add(request_id) super().suspend_request(request) @override def _setup_state_indices(self, requests, num_contexts: int = 0) -> None: super()._setup_state_indices(requests, num_contexts) - if self._kda_replay is not None and self.local_num_mamba_layers > 0: - self._kda_replay.seed_resumed_requests(self, requests, num_contexts) + replay = self._kda_replay + if replay is None or self.local_num_mamba_layers == 0: + return + replay.suspended_requests.difference_update( + request.py_request_id for request in requests[:num_contexts] + ) + slots = [] + for request in requests[num_contexts:]: + if request.py_request_id in replay.suspended_requests: + replay.suspended_requests.remove(request.py_request_id) + slots.append(self._request_id_to_state_index[request.py_request_id]) + if slots: + with replay.on_cache_stream(self._stream): + replay.seed_transferred_slots( + torch.tensor( + slots, dtype=torch.long, device=replay.prev_num_accepted_tokens.device + ) + ) @override def prepare_resources(self, scheduled_batch) -> None: diff --git a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py index 220070d07f38..d58536202538 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -592,26 +592,6 @@ def on_cache_stream(stream): yield current.wait_stream(stream) - def prepare_suspend(self, manager: MambaHybridCacheManagerV2, request) -> None: - """Commit pending drafts and conv history before releasing a request's slot.""" - request_id = request.py_request_id - cache = manager.kv_cache_map.get(request_id) - slot = manager._request_id_to_state_index.get(request_id, -1) - if slot < 0 or cache is None or not cache.is_active: - return - if request_id in self.acceptance_pending: - raise RuntimeError( - f"Request {request_id} is being suspended before update_resources " - "recorded the acceptance of its last verify, so its pending KDA drafts are unknown" - ) - if request_id in self.unseeded_requests: - self.unseeded_requests.remove(request_id) - else: - with self.on_cache_stream(manager._stream): - self.commit_pending_drafts(slot) - del manager._request_id_to_state_index[request_id] - self.suspended_requests.add(request_id) - def commit_pending_drafts(self, slot: int) -> None: """Fold accepted drafts and the replay conv window into the persistent states.""" from ._kda_kernels import commit_kda_pending_drafts @@ -638,26 +618,6 @@ def commit_pending_drafts(self, slot: int) -> None: ) self.prev_num_accepted_tokens[slot] = 0 - def seed_resumed_requests( - self, manager: MambaHybridCacheManagerV2, requests, num_contexts: int - ) -> None: - """Rebuild replay history from the V2 state restored for resumed generation.""" - self.suspended_requests.difference_update( - request.py_request_id for request in requests[:num_contexts] - ) - slots = [] - for request in requests[num_contexts:]: - if request.py_request_id in self.suspended_requests: - self.suspended_requests.remove(request.py_request_id) - slots.append(manager._request_id_to_state_index[request.py_request_id]) - if slots: - with self.on_cache_stream(manager._stream): - self.seed_transferred_slots( - torch.tensor( - slots, dtype=torch.long, device=self.prev_num_accepted_tokens.device - ) - ) - def _layer_replay_buffers(self) -> tuple[torch.Tensor, ...]: buffers = ( self.kda_conv_q, @@ -777,15 +737,46 @@ def _build_cache_config(self, config): @override def suspend_request(self, request) -> None: - if self._kda_replay is not None: - self._kda_replay.prepare_suspend(self, request) + replay = self._kda_replay + request_id = request.py_request_id + cache = self.kv_cache_map.get(request_id) + slot = self._request_id_to_state_index.get(request_id, -1) + if replay is not None and slot >= 0 and cache is not None and cache.is_active: + if request_id in replay.acceptance_pending: + raise RuntimeError( + f"Request {request_id} is being suspended before update_resources " + "recorded the acceptance of its last verify, so its pending KDA drafts are unknown" + ) + if request_id in replay.unseeded_requests: + replay.unseeded_requests.remove(request_id) + else: + with replay.on_cache_stream(self._stream): + replay.commit_pending_drafts(slot) + del self._request_id_to_state_index[request_id] + replay.suspended_requests.add(request_id) super().suspend_request(request) @override def _setup_state_indices(self, requests, num_contexts: int = 0) -> None: super()._setup_state_indices(requests, num_contexts) - if self._kda_replay is not None and self.local_num_mamba_layers > 0: - self._kda_replay.seed_resumed_requests(self, requests, num_contexts) + replay = self._kda_replay + if replay is None or self.local_num_mamba_layers == 0: + return + replay.suspended_requests.difference_update( + request.py_request_id for request in requests[:num_contexts] + ) + slots = [] + for request in requests[num_contexts:]: + if request.py_request_id in replay.suspended_requests: + replay.suspended_requests.remove(request.py_request_id) + slots.append(self._request_id_to_state_index[request.py_request_id]) + if slots: + with replay.on_cache_stream(self._stream): + replay.seed_transferred_slots( + torch.tensor( + slots, dtype=torch.long, device=replay.prev_num_accepted_tokens.device + ) + ) @override def prepare_resources(self, scheduled_batch) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index fa93692e8016..6517e258f97d 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -1728,6 +1728,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: config, event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), + max_gpu_total_bytes=kv_cache_config.max_gpu_total_bytes, ) else: init_error: Optional[Exception] = None @@ -1737,6 +1738,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: config, event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), + max_gpu_total_bytes=kv_cache_config.max_gpu_total_bytes, ) except Exception as error: if isinstance(error, (CuError, KVCacheOutOfMemoryError)): @@ -1777,6 +1779,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: config, event_manager=event_sink, cold_page_codec=create_cold_page_codec(config), + max_gpu_total_bytes=kv_cache_config.max_gpu_total_bytes, ) except Exception as error: fallback_error = error.with_traceback(None) @@ -1803,7 +1806,7 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: config = replace( config, cache_tiers=[ - replace(config.cache_tiers[0], quota=candidate.get_quota(GPU_LEVEL)), + GpuCacheTierConfig(quota=candidate.get_quota(GPU_LEVEL)), *config.cache_tiers[1:], ], ) @@ -3080,33 +3083,49 @@ def _get_external_cache_bytes(self, cache: KVCacheManagerPy) -> int: return 0 def _allocate_cache_pools( - self, config: KVCacheManagerConfigPy, *, event_manager, cold_page_codec + self, + config: KVCacheManagerConfigPy, + *, + event_manager, + cold_page_codec, + max_gpu_total_bytes: int | None = None, ) -> KVCacheManagerPy: - """Reserve external storage before publishing pool addresses and capacities.""" + """Fit native pools and external storage within the physical GPU budget.""" cache = KVCacheManagerPy( config, event_manager=event_manager, cold_page_codec=cold_page_codec ) - total_quota = config.cache_tiers[0].quota - previous_pool_quota = total_quota external_bytes = self._get_external_cache_bytes(cache) - while external_bytes and cache.get_quota(GPU_LEVEL) + external_bytes > total_quota: - # Scale the combined footprint, then recheck the actual capacity: - # native allocation granularity and minimum slots affect rounding. - native_bytes = cache.get_quota(GPU_LEVEL) - pool_quota = native_bytes * total_quota // (native_bytes + external_bytes) - if ( - pool_quota <= 0 - or pool_quota >= previous_pool_quota - or not cache.resize(GPU_LEVEL, pool_quota) - ): - cache.shutdown() - raise ValueError( - "The GPU cache quota cannot accommodate native pools and " - f"external cache storage ({external_bytes} bytes)." - ) - previous_pool_quota = pool_quota - external_bytes = self._get_external_cache_bytes(cache) - return cache + # A token limit sizes the native pools; their minimum slots and allocation + # granularity can exceed that logical size. Only a byte limit caps the + # combined physical allocation, including model-owned scratch. + if not external_bytes or not max_gpu_total_bytes: + return cache + native_bytes = cache.get_quota(GPU_LEVEL) + if native_bytes + external_bytes <= max_gpu_total_bytes: + return cache + + lower_quota = 0 + upper_quota = config.cache_tiers[0].quota + pool_quota = min( + upper_quota - 1, + native_bytes * max_gpu_total_bytes // (native_bytes + external_bytes), + ) + while lower_quota < pool_quota < upper_quota: + if not cache.resize(GPU_LEVEL, pool_quota): + # A quota below the native minimum cannot be resized. Search + # above it before concluding that the combined budget cannot fit. + lower_quota = pool_quota + else: + external_bytes = self._get_external_cache_bytes(cache) + if cache.get_quota(GPU_LEVEL) + external_bytes <= max_gpu_total_bytes: + return cache + upper_quota = pool_quota + pool_quota = (lower_quota + upper_quota) // 2 + cache.shutdown() + raise ValueError( + "The GPU cache quota cannot accommodate native pools and " + f"external cache storage ({external_bytes} bytes)." + ) def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """Customize the general cache config for a specialized cache manager.""" diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py index 56a1e6ed8358..b0a4431b6510 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -432,13 +432,11 @@ def _max_resident_sequences(self) -> int: return self.max_batch_size * self.mapping.pp_size def _mamba_state_bytes_per_slot(self) -> int: - return ( - sum( - buffer.size - for layer_id in self.mamba_pp_layers - for buffer in self._get_recurrent_buffer_configs(layer_id) - ) - + self._extra_scratch_bytes_per_slot() + """Return persistent bytes per slot allocated by the native pools.""" + return sum( + buffer.size + for layer_id in self.mamba_pp_layers + for buffer in self._get_recurrent_buffer_configs(layer_id) ) def _num_ssm_snapshots_for_capacity( diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 875fc73fc7d9..4da3ae8093fa 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3,7 +3,6 @@ """Regression tests for Python, Cpp, and V2 Mamba cache managers.""" import contextlib -import math import os from types import SimpleNamespace from unittest.mock import MagicMock @@ -11,6 +10,9 @@ import pytest import torch +from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.cache_manager import ( + Glm5NextCacheManager, +) from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager from tensorrt_llm._torch.disaggregation.resource.page import ( MAMBA_CONV_ROLE, @@ -2422,6 +2424,8 @@ def _build_v2_hybrid_with_mamba_layer( mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, enable_kv_pool_rebalance=False, + max_tokens=512, + max_gpu_total_bytes=0, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2435,7 +2439,8 @@ def _build_v2_hybrid_with_mamba_layer( enable_attention_dp=enable_attention_dp, ) kv_cache_config = KvCacheConfig( - max_tokens=512, + max_tokens=max_tokens, + max_gpu_total_bytes=max_gpu_total_bytes, enable_block_reuse=enable_block_reuse, enable_partial_reuse=enable_partial_reuse, block_reuse_config=BlockReuseConfig( @@ -3678,7 +3683,8 @@ def test_v2_hybrid_intermediate_states_size_by_tokens_per_gen_step(): @skip_no_cuda -def test_v2_kda_replay_policy_allocates_logical_slot_caches(): +@pytest.mark.parametrize("max_tokens,byte_limit", [(512, 0), (None, 16 << 20), (512, 16 << 20)]) +def test_v2_kda_replay_policy_allocates_logical_slot_caches(max_tokens, byte_limit): spec_config = MTPDecodingConfig(max_draft_len=2) policy = KDAReplayState( num_speculative_tokens=2, @@ -3693,22 +3699,22 @@ def test_v2_kda_replay_policy_allocates_logical_slot_caches(): mamba_n_groups=6, mamba_ssm_cache_dtype=torch.float32, speculative_state=policy, + max_tokens=max_tokens, + max_gpu_total_bytes=byte_limit, ) try: assert mgr._kda_replay is policy assert policy.intermediate_indices.shape == (4,) assert not hasattr(mgr, "get_replay_state_update_metadata") - total_quota = math.ceil( - mgr._get_quota_from_max_tokens(mgr.kv_cache_config.max_tokens) - / mgr.kv_cache_config.max_util_for_resume - ) replay_bytes = sum( buffer.untyped_storage().nbytes() for buffer in (*policy._layer_replay_buffers(), policy.prev_num_accepted_tokens) ) native_bytes = mgr.impl.get_quota(GPU_LEVEL) - assert native_bytes + replay_bytes <= total_quota + assert mgr._get_external_cache_bytes(mgr.impl) == replay_bytes + if byte_limit: + assert native_bytes + replay_bytes <= byte_limit assert mgr.kv_cache_manager_py_config.cache_tiers[0].quota == native_bytes layer_cache = mgr.mamba_layer_cache(0) @@ -3946,11 +3952,11 @@ def suspend(self) -> None: self.is_active = False -def _v2_kda_replay_stub(num_slots=4): +def _v2_kda_replay_stub(num_slots=4, manager_cls=KimiK3HybridCacheManager): """A V2 hybrid manager stub with CPU KDA replay buffers and one layer's V2 conv states ([Q | K | V] sections of 2 channels, a committed window of 3), for the state-index setup that prepare_resources runs. Committing pending drafts is recorded instead of run (it needs the GPU kernels).""" - mgr = object.__new__(KimiK3HybridCacheManager) + mgr = object.__new__(manager_cls) mgr.local_num_mamba_layers = 1 mgr._kda_replay = KDAReplayState(2) mgr._speculative_state = mgr._kda_replay @@ -3996,12 +4002,13 @@ def _kda_replay_slot(mgr, slot): @pytest.mark.parametrize("resume_slot", [1, 0], ids=["new_slot", "old_slot"]) -def test_v2_kda_replay_suspended_request_resumes_from_its_v2_states(resume_slot): +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) +def test_v2_kda_replay_suspended_request_resumes_from_its_v2_states(resume_slot, manager_cls): """A suspend commits the request's pending drafts and conv window into its slot's V2 states and forgets the slot, which can then go to another request. The request resumes in a new slot, or in its old one after that request has finished. Either way its replay caches are seeded from the V2 conv state it resumes with, nothing is pending, and the other request's replay state is left alone.""" - mgr = _v2_kda_replay_stub() + mgr = _v2_kda_replay_stub(manager_cls=manager_cls) first = SimpleNamespace(py_request_id=1, is_dummy=False) second = SimpleNamespace(py_request_id=2, is_dummy=False) mgr.kv_cache_map[1] = _SlotCache(0) @@ -4043,14 +4050,15 @@ def test_v2_kda_replay_suspended_request_resumes_from_its_v2_states(resume_slot) @pytest.mark.parametrize("recorded_by", ["update_resources", "revert_allocate_generation"]) +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) def test_v2_kda_replay_refuses_a_suspend_before_the_host_drafter_acceptance( - monkeypatch, recorded_by + monkeypatch, recorded_by, manager_cls ): """With a host drafter, update_resources records the acceptance of a batch's verify, and a suspend needs it to commit the request's pending drafts. A suspend between the batch's prepare_resources and its update_resources raises instead of committing a stale count. Once update_resources has recorded it, or the batch is reverted without running, the suspend commits.""" - mgr = _v2_kda_replay_stub() + mgr = _v2_kda_replay_stub(manager_cls=manager_cls) mgr.spec_config = SimpleNamespace(decoding_type="NGram") for method in ("prepare_resources", "update_resources", "revert_allocate_generation"): monkeypatch.setattr(KVCacheManagerV2, method, MagicMock()) @@ -4077,10 +4085,11 @@ def test_v2_kda_replay_refuses_a_suspend_before_the_host_drafter_acceptance( @pytest.mark.parametrize("seeded", [False, True], ids=["before_seeding", "after_seeding"]) +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) def test_v2_kda_replay_disagg_generation_request_suspends_its_transferred_state( - monkeypatch, seeded + monkeypatch, seeded, manager_cls ): - mgr = _v2_kda_replay_stub() + mgr = _v2_kda_replay_stub(manager_cls=manager_cls) monkeypatch.setattr(KVCacheManagerV2, "prepare_resources", MagicMock()) request = SimpleNamespace(py_request_id=1, is_dummy=False, is_disagg_generation_init_state=True) mgr.kv_cache_map[1] = _SlotCache(0) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py index 92b03a2a3e60..86c0c327a634 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -46,7 +46,7 @@ MambaRole, MambaStateLayout, ) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, @@ -724,9 +724,10 @@ def test_kda_scratch_cost_matches_allocated_slot_buffers(): @pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) @pytest.mark.parametrize("replay", [False, True]) @pytest.mark.parametrize("pool_share", [0.25, 0.75]) -@pytest.mark.parametrize("quota_source", ["max_tokens", "max_gpu_total_bytes"]) +@pytest.mark.parametrize("quota_source", ["max_tokens", "max_gpu_total_bytes", "both"]) +@pytest.mark.parametrize("granularity", [1, 4096]) def test_kda_pool_allocation_reserves_actual_replay_capacity( - monkeypatch, manager_cls, replay, pool_share, quota_source + monkeypatch, manager_cls, replay, pool_share, quota_source, granularity ): from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as module @@ -743,15 +744,18 @@ def test_kda_pool_allocation_reserves_actual_replay_capacity( manager.kv_cache_config = KvCacheConfig(enable_block_reuse=False) persistent_bytes_per_slot = 2 * (128 + 288) total_quota = 1 << 20 - if quota_source == "max_tokens": + if quota_source in ("max_tokens", "both"): monkeypatch.setattr( KVCacheManagerV2, "_get_quota_from_max_tokens", lambda self, tokens: tokens * 64 ) total_quota = manager._get_quota_from_max_tokens(16384) + assert total_quota == 16384 * 64 + 4 * persistent_bytes_per_slot class Pool: def __init__(self, config, **kwargs): - self.quota = config.cache_tiers[0].quota + self.quota = ( + (config.cache_tiers[0].quota + granularity - 1) // granularity * granularity + ) self.resize_calls = 0 @property @@ -769,7 +773,7 @@ def get_page_index_upper_bound(self, layer, role): def resize(self, level, quota): self.resize_calls += 1 - self.quota = quota + self.quota = (quota + granularity - 1) // granularity * granularity return True monkeypatch.setattr(module, "KVCacheManagerPy", Pool) @@ -778,7 +782,10 @@ def resize(self, level, quota): layers=[], cache_tiers=[GpuCacheTierConfig(quota=total_quota), HostCacheTierConfig(quota=12345)], ) - pool = manager._allocate_cache_pools(config, event_manager=None, cold_page_codec=None) + byte_limit = total_quota if quota_source != "max_tokens" else None + pool = manager._allocate_cache_pools( + config, event_manager=None, cold_page_codec=None, max_gpu_total_bytes=byte_limit + ) # The native pool can have many more slots than the live batch needs. assert pool.slots > layout.max_batch_size if replay: @@ -791,12 +798,15 @@ def resize(self, level, quota): buffer.untyped_storage().nbytes() for buffer in (*state._layer_replay_buffers(), state.prev_num_accepted_tokens) ) - assert pool.resize_calls > 0 - assert pool.quota + scratch_bytes <= total_quota + if byte_limit: + assert pool.resize_calls > 0 + assert pool.quota + scratch_bytes <= byte_limit + else: + assert pool.resize_calls == 0 assert manager._get_external_cache_bytes(pool) == scratch_bytes else: assert pool.resize_calls == 0 - assert pool.quota == total_quota + assert pool.quota == (total_quota + granularity - 1) // granularity * granularity # Host fallback retries must retain the original total GPU budget. assert config.cache_tiers[0].quota == total_quota assert config.cache_tiers[1].quota == 12345 @@ -817,16 +827,62 @@ def test_external_cache_reservation_rejects_unusable_quota(monkeypatch, resize_s ) # A backend that reports success without releasing memory must fail too. with pytest.raises(ValueError, match="external cache storage"): - manager._allocate_cache_pools(config, event_manager=None, cold_page_codec=None) + manager._allocate_cache_pools( + config, event_manager=None, cold_page_codec=None, max_gpu_total_bytes=1024 + ) cache.shutdown.assert_called_once() +@pytest.mark.parametrize("byte_limit", [1500, 1800]) +def test_external_cache_reservation_handles_native_minimum_and_rounding(monkeypatch, byte_limit): + from tensorrt_llm._torch.pyexecutor.kv_cache import kv_cache_manager_v2 as module + + manager = object.__new__(KVCacheManagerV2) + manager._get_external_cache_bytes = lambda cache: cache.quota + + class Pool: + quota = 1024 + shutdown_called = False + + def get_quota(self, level): + return self.quota + + def resize(self, level, quota): + rounded = (quota + 255) // 256 * 256 + if rounded < 768: + return False + self.quota = rounded + return True + + def shutdown(self): + self.shutdown_called = True + + pool = Pool() + monkeypatch.setattr(module, "KVCacheManagerPy", lambda *args, **kwargs: pool) + config = KVCacheManagerConfig( + tokens_per_block=32, layers=[], cache_tiers=[GpuCacheTierConfig(quota=1024)] + ) + if byte_limit < 2 * 768: + with pytest.raises(ValueError, match="external cache storage"): + manager._allocate_cache_pools( + config, event_manager=None, cold_page_codec=None, max_gpu_total_bytes=byte_limit + ) + assert pool.shutdown_called + else: + manager._allocate_cache_pools( + config, event_manager=None, cold_page_codec=None, max_gpu_total_bytes=byte_limit + ) + assert pool.quota == 768 + assert not pool.shutdown_called + + @pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) @pytest.mark.parametrize("stochastic_rounding", [False, True]) @pytest.mark.parametrize("replay", [False, True]) @pytest.mark.parametrize("reuse", [False, True]) +@pytest.mark.parametrize("interval", [0, 32]) def test_kda_affine_estimator_includes_replay_scratch( - monkeypatch, manager_cls, stochastic_rounding, replay, reuse + monkeypatch, manager_cls, stochastic_rounding, replay, reuse, interval ): from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import cache_manager as glm_module from tensorrt_llm._torch.modules.kimi_kda import cache_manager as kda_module @@ -858,7 +914,10 @@ def local_layout(*args, **kwargs): quant_config=SimpleNamespace(mamba_ssm_stochastic_rounding=stochastic_rounding), ) monkeypatch.setattr(glm_module, "unwrap_glm5_next_text_config", lambda config: config) - kv_config = KvCacheConfig(enable_block_reuse=reuse) + kv_config = KvCacheConfig( + enable_block_reuse=reuse, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=interval), + ) cost = manager_cls.get_cache_size_per_token( config, layout.mapping, @@ -881,7 +940,7 @@ def local_layout(*args, **kwargs): ) state_bytes = 2 * 416 + scratch_per_slot interval = kv_config.mamba_state_config.periodic_snapshot_interval - slope = (state_bytes + interval - 1) // interval if reuse else 0 + slope = (state_bytes + interval - 1) // interval if reuse and interval > 0 else 0 assert cost == (slope, (layout.max_batch_size + dummy_slots) * state_bytes) From bc6edc27bfbd9f644474998938f3e1000a2b05cd Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:12:07 +0000 Subject: [PATCH 17/25] [None][fix] Adapt Qwen4 modeling v2 to layer caches (By Agent) Bind speculative conv and recurrent buffers through mamba_layer_cache and check the current GDN replay flag. The modeling_v2 target introduced by main still called the removed getters, causing both Qwen4 MTP accuracy gates to fail during initialization. Add a regression test using the real Qwen4 manager and GDN state views, including nontrivial layer offsets, and include it in B300 CI. Validation: 37 pool-contract and QSA wiring tests passed; pre-commit checks passed. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../modeling.py | 11 +++--- .../test_lists/test-db/l0_b300.yml | 1 + ...est_modeling_v2_qwen4_exp_pool_contract.py | 34 +++++++++++++++++++ 3 files changed, 41 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py index f71b25463bdd..0eea250b9c5c 100644 --- a/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py +++ b/tensorrt_llm/_torch/_experimental/modeling_v2/models/qwen4_exp/qwen3_8_flash_next_nvfp4__sm_103__tp1/modeling.py @@ -1316,18 +1316,19 @@ def _bind_step_buffers(self, cm, g: _PoolGeometry, conv0, ssm0) -> dict: f"the KV cache manager {type(cm).__name__} reports is_speculative() False, " "so it allocated none of the per-step state buffers the verify path writes" ) - assert not getattr(cm, "use_replay_state_update", False), ( + assert not cm.use_gdn_cached_replay_all_layer_commit, ( "the GDN cached-replay verify path is on. Export TRTLLM_USE_GDN_REPLAY=0 " "before the ranks start; see this target's __init__ for why" ) _require( cm, - ("get_intermediate_conv_states", "get_intermediate_ssm_states"), + ("mamba_layer_cache",), "reads both per-step state buffers", ) - # `get_intermediate_*` builds a fresh view per call, so bind them here. - int_conv = {i: cm.get_intermediate_conv_states(i) for i in self.gdn_layers} - int_ssm = {i: cm.get_intermediate_ssm_states(i) for i in self.gdn_layers} + # Bind the per-layer views once for subsequent forwards. + layer_caches = {i: cm.mamba_layer_cache(i) for i in self.gdn_layers} + int_conv = {i: cache.intermediate_conv_window for i, cache in layer_caches.items()} + int_ssm = {i: cache.intermediate_ssm for i, cache in layer_caches.items()} int_idx = getattr(cm, "intermediate_state_indices", None) g0 = self.gdn_layers[0] check_step_buffers(g, int_conv[g0], int_ssm[g0], conv0, ssm0, int_idx) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 0a29a32feca5..131059029192 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -44,6 +44,7 @@ l0_b300: - unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py - unittest/_torch/modeling/test_qsa_runtime_wiring.py - unittest/_torch/modeling/test_qwen4_exp_support.py + - unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/visual_gen/test_attention_flashinfer.py diff --git a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py index d89a528d1417..c85b39f123a9 100644 --- a/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py +++ b/tests/unittest/_torch/modeling_v2/test_modeling_v2_qwen4_exp_pool_contract.py @@ -9,6 +9,8 @@ from __future__ import annotations +from types import SimpleNamespace + import pytest import torch @@ -17,6 +19,8 @@ from tensorrt_llm._torch._experimental.modeling_v2.models.qwen4_exp.qwen3_8_flash_next_nvfp4__sm_103__tp1 import ( # noqa: E501 modeling as M, ) +from tensorrt_llm._torch.modules.fla.cache_manager import GDNIntermediateState +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager G = M._PoolGeometry( tokens_per_block=32, @@ -71,6 +75,36 @@ def test_valid_pools_pass() -> None: ) +def test_bind_step_buffers_uses_manager_layer_views() -> None: + state = GDNIntermediateState() + state.intermediate_conv = _t(2, SLOTS, 2, G.conv_dim, G.conv_kernel - 1) + state.intermediate_ssm = _t( + 2, SLOTS, 2, G.num_v_heads, G.head_v_dim, G.head_k_dim, dtype=torch.float32 + ) + state.intermediate_indices = _t(SLOTS, dtype=torch.int32) + manager = Qwen4ExpHybridCacheManager.__new__(Qwen4ExpHybridCacheManager) + manager.spec_config = SimpleNamespace(tokens_per_gen_step=2) + manager._speculative_state = state + manager.mamba_layer_offsets = {3: 1, 7: 0} + manager.all_conv_states = [_conv(), _conv()] + manager.all_ssm_states = [_ssm(), _ssm()] + + bound = M.Qwen4ExpTp1ModelingV2._bind_step_buffers( + SimpleNamespace(gdn_layers=[3, 7]), manager, G, _conv(), _ssm() + ) + + assert bound["int_idx"] is state.intermediate_indices + for role, storage in ( + ("int_conv", state.intermediate_conv), + ("int_ssm", state.intermediate_ssm), + ): + for layer, offset in manager.mamba_layer_offsets.items(): + view = bound[role][layer] + assert view._base is storage + assert view.storage_offset() == offset * storage.stride(0) + assert view.shape == storage.shape[1:] + + @pytest.mark.parametrize( "plane, match", [ From debdbb07998de565f9ae2b2c897779bc9430bbfa Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:48:03 +0000 Subject: [PATCH 18/25] [None][test] Close DSA cache managers during teardown (By Agent) Retain test-owned DSA cache managers until fixture teardown and explicitly shut them down on the test thread. Leaving request caches to cyclic GC can run CUDA event cleanup on a log-upload worker without a CUDA context, poisoning KVCM2 and causing cascading failures in the attention suite. Validation: reproduced the invalid-context failure with worker-thread GC; explicit shutdown prevents it. The complete DSA indexer module passed with worker-thread GC between tests (157 passed, 4 skipped). Pre-commit passed. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../attention/sparse/dsa/test_dsa_indexer.py | 78 ++++++++++++------- 1 file changed, 52 insertions(+), 26 deletions(-) diff --git a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py index 23d90f830a49..4af91655dbd5 100644 --- a/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py +++ b/tests/unittest/_torch/attention/sparse/dsa/test_dsa_indexer.py @@ -849,6 +849,18 @@ def create_dsa_cache_manager( return cache_manager, sparse_attn_config +@pytest.fixture +def dsa_cache_manager_factory(request: pytest.FixtureRequest): + """Keep managers alive until teardown and close them on the test thread.""" + + def create(**kwargs): + manager, config = create_dsa_cache_manager(**kwargs) + request.addfinalizer(manager.shutdown) + return manager, config + + return create + + def test_transform_local_topk_uses_v2_page_mapping(): """Use the V2 pool scale and offset when flattening DSA indices.""" cache_manager = SimpleNamespace( @@ -2060,7 +2072,7 @@ def validate_topk_indices(topk_indices_0, topk_indices_1, total_tokens): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper -def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): +def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(dsa_cache_manager_factory): """Recompute slot mappings without re-running full Indexer.prepare().""" head_dim = 128 block_size = 64 @@ -2071,7 +2083,7 @@ def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): num_ctx_tokens = seq_lens[0].item() num_tokens = seq_lens.sum().item() - cache_manager, _ = create_dsa_cache_manager( + cache_manager, _ = dsa_cache_manager_factory( batch_size=len(request_ids), head_dim=head_dim, tokens_per_block=block_size, @@ -2111,7 +2123,7 @@ def test_recompute_slot_mappings_matches_prepare_with_cached_tokens(): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper -def test_indexer_k_cache_scatter_custom_op(): +def test_indexer_k_cache_scatter_custom_op(dsa_cache_manager_factory): """ Direct comparison: CUDA kernel vs Python reference for k_cache scatter. @@ -2132,7 +2144,7 @@ def test_indexer_k_cache_scatter_custom_op(): layer_idx_python = 2 # Create cache manager with multiple layers - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -2302,7 +2314,7 @@ def _unravel_indices(flat_indices, shape): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper -def test_fp8_k_cache_roundtrip(): +def test_fp8_k_cache_roundtrip(dsa_cache_manager_factory): """Verify FP8 quantization scales survive write/read cycle for multiple requests.""" torch.manual_seed(42) @@ -2312,7 +2324,7 @@ def test_fp8_k_cache_roundtrip(): max_seq_len = 512 num_tokens_per_req = [150, 100] # Request 0: 3 blocks, Request 1: 2 blocks - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -2405,7 +2417,9 @@ def test_fp8_k_cache_roundtrip(): ], ) @pytest.mark.parametrize("compress_ratio", [1, 4]) -def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, compress_ratio): +def test_indexer_decode_with_paged_kv_cache( + dsa_cache_manager_factory, batch_size, next_n, backend, compress_ratio +): """ Test FP8 paged KV cache with two-phase workflow and variable context lengths. @@ -2458,7 +2472,7 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, compres print(f" Compress ratio: {compress_ratio}") # Setup: Create cache manager and indexer - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -2691,7 +2705,9 @@ def test_indexer_decode_with_paged_kv_cache(batch_size, next_n, backend, compres ), ], ) -def test_indexer_decode_with_paged_kv_cache_fp4(batch_size, next_n, backend): +def test_indexer_decode_with_paged_kv_cache_fp4( + dsa_cache_manager_factory, batch_size, next_n, backend +): """FP4 indexer paged decode test (DeepGEMM + CuTe DSL backends). Mirrors test_indexer_decode_with_paged_kv_cache but exercises the FP4 @@ -2752,7 +2768,7 @@ def test_indexer_decode_with_paged_kv_cache_fp4(batch_size, next_n, backend): print(f" Final lengths: {final_lens.tolist()}") # Setup: FP4 cache manager + indexer - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -3475,7 +3491,9 @@ def test_split_prefill_chunks(max_chunk_size, seq_lens, start_idx, expected_spec (1000, [400, 600, 2200, 500], "two_level"), # Mixed: request 2 needs Q-blocks ], ) -def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compress_ratio): +def test_indexer_chunked_prefill( + dsa_cache_manager_factory, chunk_size, seq_lens_list, chunking_type, compress_ratio +): """ Tests for indexer chunked prefill: 1. Request-level chunking: Multiple small requests packed together @@ -3539,7 +3557,7 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr print(f" Request {req_idx}: {seq_len} tokens → {num_q_blocks} Q-blocks") # Create cache manager and indexer - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -3740,7 +3758,9 @@ def test_indexer_chunked_prefill(chunk_size, seq_lens_list, chunking_type, compr @pytest.mark.parametrize("next_n", [1, 2, 4]) @pytest.mark.parametrize("index_topk", [2048]) @pytest.mark.parametrize("seq_len_range", [(2048, 8192), (512, 1024)]) -def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_len_range): +def test_indexer_decode_custom_vs_fallback( + dsa_cache_manager_factory, batch_size, next_n, index_topk, seq_len_range +): """ Test that the production and Torch Top-K modules produce identical results in the decode phase of sparse_attn_indexer. @@ -3799,7 +3819,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l num_cached_tokens = kv_lens.tolist() # Create cache manager and indexer - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -3974,7 +3994,7 @@ def test_indexer_decode_custom_vs_fallback(batch_size, next_n, index_topk, seq_l @skip_pre_hopper @pytest.mark.parametrize("step0_mode", ["gen", "context"]) @pytest.mark.parametrize("batch_size", [1, 2]) -def test_indexer_decode_mtp_topk_reuse(step0_mode, batch_size): +def test_indexer_decode_mtp_topk_reuse(dsa_cache_manager_factory, step0_mode, batch_size): """Verify mtp_index_share: draft step 0 stashes each request's last-token Top-K and draft steps > 0 reuse it verbatim (bit-exact), across > 1 draft step. Covers both step-0 shapes: "gen" (steady state, next_n > 1) and @@ -3990,7 +4010,7 @@ def test_indexer_decode_mtp_topk_reuse(step0_mode, batch_size): # path, which assumes num_gen_tokens == num_generations * (1 + md). md = 1 - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -4133,7 +4153,9 @@ def make_inputs(n_tokens): @pytest.mark.parametrize("batch_size", [4, 16]) @pytest.mark.parametrize("index_topk", [2048]) @pytest.mark.parametrize("chunk_size", [1024, 2048]) -def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chunk_size): +def test_indexer_prefill_chunked_custom_vs_fallback( + dsa_cache_manager_factory, batch_size, index_topk, chunk_size +): """ Test chunked prefill: production Top-K vs Torch Top-K with metadata.indexer_prefill_chunks != None. @@ -4160,7 +4182,7 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun total_tokens = seq_lens.sum().item() # Create cache manager and indexer - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -4251,7 +4273,9 @@ def test_indexer_prefill_chunked_custom_vs_fallback(batch_size, index_topk, chun @pytest.mark.parametrize("batch_size", [4, 16]) @pytest.mark.parametrize("index_topk", [2048]) @pytest.mark.parametrize("seq_len_range", [(1, 512)]) -def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, seq_len_range): +def test_indexer_prefill_single_pass_custom_vs_fallback( + dsa_cache_manager_factory, batch_size, index_topk, seq_len_range +): """ Test single-pass prefill: production Top-K vs Torch Top-K with metadata.indexer_prefill_chunks == None (else branch). @@ -4270,7 +4294,7 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, total_tokens = seq_lens.sum().item() # Create cache manager and indexer - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -4414,7 +4438,9 @@ def test_indexer_prefill_single_pass_custom_vs_fallback(batch_size, index_topk, @skip_pre_hopper @pytest.mark.parametrize("enable_indexer_skip", [True, False]) -def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): +def test_indexer_topk_multi_request_with_different_cache( + dsa_cache_manager_factory, enable_indexer_skip +): """ Test that custom topk kernel handles multi-request batches with different cached amounts. """ @@ -4442,7 +4468,7 @@ def test_indexer_topk_multi_request_with_different_cache(enable_indexer_skip): print(f" Req1: {cached_tokens[1]} cached + {seq_lens[1]} new = {total_kv_lens[1]} total") # Create cache manager - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -4718,7 +4744,7 @@ def is_attn_metadata(obj, cls): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_blackwell -def test_cutedsl_mqa_logits_output_buffer_persistent(): +def test_cutedsl_mqa_logits_output_buffer_persistent(dsa_cache_manager_factory): """Regression: the CuteDSL paged-MQA-logits output must have a STABLE address across calls, not be a per-forward ``torch.empty``. @@ -4739,7 +4765,7 @@ def test_cutedsl_mqa_logits_output_buffer_persistent(): heads = 32 kv_len = 4096 - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, @@ -4803,7 +4829,7 @@ def _mqa(): @pytest.mark.skipif(not has_deep_gemm(), reason="DeepGEMM not available") @skip_pre_hopper -def test_topk_indices_buffer_cuda_graph(): +def test_topk_indices_buffer_cuda_graph(dsa_cache_manager_factory): from tensorrt_llm._torch.memory_buffer_utils import Buffers batch_size, next_n = 4, 3 @@ -4813,7 +4839,7 @@ def test_topk_indices_buffer_cuda_graph(): layer_idx = 0 num_tokens = batch_size * next_n - cache_manager, sparse_attn_config = create_dsa_cache_manager( + cache_manager, sparse_attn_config = dsa_cache_manager_factory( batch_size=batch_size, head_dim=head_dim, tokens_per_block=block_size, From 866313687c12fdc70f4b7ab73713452a0bd6f6e5 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Thu, 15 Oct 2026 21:33:26 +0000 Subject: [PATCH 19/25] [None][test] Split B200 sampler invocations to avoid timeout (By Agent) Run beam-search tests separately from the remaining sampler coverage. CI 64455 reached the outer timeout after 1,315 cases, with 3,528 seconds spent in completed tests and no assertion failures. Separate invocations preserve coverage while giving each group its own existing time limit. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 1b7d163cbe9e..6b3346e0fefd 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -245,11 +245,13 @@ l0_b200: # UUID KV-event tests hardcode Qwen3-VL-2B (no model_dir fixture id), so they # are excluded by the qwen3_30b_a3b_fp8 filter above; select them explicitly. - unittest/_torch/multimodal/test_mm_encoder_standalone.py -k "test_kv_event_mm_keys and uuid" + # Run beam-search end-to-end tests in a separate sampler invocation. + - unittest/_torch/sampler/test_beam_search.py # test_speculative_d2h_parity_real_predictor (both no_stop_token and # stop_token_15) is a HW-agnostic sampler D2H-parity invariant (plain CUDA # stream/event, no SM gating), covered on A30 pre-merge; drop that whole # function's leaves on Blackwell, keep the rest of the sampler dir. - - unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" + - unittest/_torch/sampler --ignore=unittest/_torch/sampler/test_beam_search.py -k "not test_speculative_d2h_parity_real_predictor" - unittest/_torch/speculative/test_eagle3.py - unittest/_torch/speculative/test_lilicorr_loading.py - unittest/_torch/speculative/test_dspark_cute_dsl_attention.py From bb91945f75738d8e07195510df46026c31b2e9e2 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:05:43 +0000 Subject: [PATCH 20/25] [None][test] Preserve FP32 activation in fused MoE reference (By Agent) Keep the Blackwell gather GEMM reference and activation in FP32, matching the fused kernel's accumulator path. Convert to BF16 only at the input to the standalone FP4 quantizer, which does not accept FP32 inputs. This avoids rounding the GEMM output before Relu2 and other activations, which can move reference scale factors across FP8 quantization boundaries. Retain the existing accuracy tolerances and all parameterized cases. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py index 706d55ae10c7..7a69f36a733a 100644 --- a/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py @@ -1307,6 +1307,7 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( hidden_size, ) + # Match the fused epilogue's FP32 GEMM and activation computation. c_ref = cute_dsl_nvfp4_grouped_gemm_ref( a_gathered, b, @@ -1316,10 +1317,11 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell( tile_idx_to_group_idx, num_non_exiting_tiles, tile_size=tile_size, - output_dtype=torch.bfloat16, + output_dtype=torch.float32, scaling_vector_size=sf_vec_size, ) - c_ref = apply_activation_ref(c_ref, activation_type, swiglu_limit) + # The standalone FP4 quantizer requires a half-precision input. + c_ref = apply_activation_ref(c_ref, activation_type, swiglu_limit).to(torch.bfloat16) global_sf = c_ref[:num_valid_permuted_tokens].abs().max().float() / (448 * 6) c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False) From 5f92fb2c9f053e09d330437ead2138611d77cdea Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 20 Oct 2026 08:41:12 +0000 Subject: [PATCH 21/25] [None][test] Fix KV event expectations and mixed-batch scheduling (By Agent) Include stable generated tokens in V2 cache event expectations after upstream generation KV reuse. Keep exact token, event ordering, tier transition, and eviction checks. Submit the Nemotron mixed-batch requests before waiting for Ray streaming responses, with the first prefill filling the token budget. Preserve the iteration-overlap and numerical parity assertions. Validation with CI 65202 binaries on B200: reproduced the old KV event assertion; both single-GPU and attention-DP event tests pass; the Nemotron Ray TP1 breakable prefill test passes. Pre-commit passed. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../modeling/test_modeling_nemotron_h.py | 22 +++++++------------ .../llmapi/test_llm_kv_cache_events.py | 10 +++++---- 2 files changed, 14 insertions(+), 18 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index 0bbe517168dd..5dcf0addc06d 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -252,9 +252,7 @@ def test_nemotron_h_chunked_prefill(): # The mixed-batch section adds the streaming decode request and the context # request admitted while it decodes. _BCG_MIXED_BATCH_REQUESTS = 2 -# Decode length of that streaming request: with EOS ignored this is a ~1 s -# generation window, orders of magnitude longer than the client-side gap -# before the second request is submitted. +# Decode length of the streaming request, with EOS ignored. _BCG_STREAM_MAX_TOKENS = 64 # One (token_ids, logprobs) entry per request from each backend run. _BCG_NUM_REQUESTS = (sum(len(batch) for batch in _BCG_CONTEXT_BATCHES) + @@ -329,11 +327,11 @@ def _run_nemotron_h_prefill_backend(backend: PrefillCudaGraphBackend, per_request.append((list(output.outputs[0].token_ids), _first_step_logprobs(output))) - # Mixed batch: admit a context request while another request is - # decoding, so BCG replays a batch that carries both a context chunk - # and decode tokens. The overlap is proved afterwards from the two - # requests' iteration metrics rather than assumed. - decoding = llm.generate_async([_BCG_PROMPT_TOKEN] * 128, + # The first prefill fills the token budget, so the second request + # must start in a later iteration alongside the first request's decode. + # Submit both before waiting for streamed logits: response transport + # can otherwise delay admission until the first request has finished. + decoding = llm.generate_async([_BCG_PROMPT_TOKEN] * _BCG_MAX_NUM_TOKENS, sampling_params=SamplingParams( max_tokens=_BCG_STREAM_MAX_TOKENS, temperature=0.0, @@ -341,10 +339,6 @@ def _run_nemotron_h_prefill_backend(backend: PrefillCudaGraphBackend, return_generation_logits=True, return_perf_metrics=True), streaming=True) - # In streaming mode every response carries only the newest step's - # logits, so the first-step distribution must be read from the first - # streamed response, before the stream moves on. - next(decoding) admitted = llm.generate_async([_BCG_PROMPT_TOKEN] * 65, sampling_params=SamplingParams( max_tokens=4, @@ -353,8 +347,8 @@ def _run_nemotron_h_prefill_backend(backend: PrefillCudaGraphBackend, return_generation_logits=True, return_perf_metrics=True), streaming=False) - # Submit before the CPU logits conversion so it cannot delay admission. - # The first response stays current until the next stream consumption. + # Each streamed response carries only its newest step's logits. + next(decoding) decoding_first_step_logprobs = _first_step_logprobs(decoding) decoding_out = decoding.result() admitted_out = admitted.result() diff --git a/tests/unittest/llmapi/test_llm_kv_cache_events.py b/tests/unittest/llmapi/test_llm_kv_cache_events.py index 53ae9569333a..13b084ffa22c 100644 --- a/tests/unittest/llmapi/test_llm_kv_cache_events.py +++ b/tests/unittest/llmapi/test_llm_kv_cache_events.py @@ -991,7 +991,8 @@ def events_for_rank(events: list[dict]) -> list[dict]: tokens_per_block = llm.get_kv_cache_capacity()["tokensPerBlock"] prompt_length = 127 - blocks_per_request = (prompt_length + tokens_per_block - + cached_length = prompt_length + sampling_params.max_tokens - 1 + blocks_per_request = (cached_length + tokens_per_block - 1) // tokens_per_block num_requests = 1 cached_blocks = {} @@ -1059,11 +1060,12 @@ def events_for_rank(events: list[dict]) -> list[dict]: del cached_blocks[block_hash] num_removed_blocks += 1 - # V2 publishes the prompt, including its partial tail, but no decode tokens. + # The final sampled token has no KV until another forward pass. + cached_tokens = prompt + list(output.outputs[0].token_ids[:-1]) assert [[token["token_id"] for token in block["tokens"]] for block in stored_blocks] == [ - prompt[start:start + tokens_per_block] - for start in range(0, prompt_length, tokens_per_block) + cached_tokens[start:start + tokens_per_block] + for start in range(0, cached_length, tokens_per_block) ] if request_index == 0: assert num_offloaded_blocks == 0 From 58ef5168f66b92c62e668395182584d38995b48d Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:54:59 +0000 Subject: [PATCH 22/25] [None][test] Assign global IDs to disaggregated worker requests (By Agent) The worker test driver sends requests directly to multiple context workers without a disagg_request_id. Their independent local request ID sequences can collide in the generation worker's receive-session map, losing transfer completion notifications and causing KV transfer timeouts. Use the production global ID generator for each logical request and carry the ID from the context response into the generation request. Preserve the existing request concurrency, rounds, timeouts, and routing assertions. Validation: the original test_workers_kv_cache_aware_router passed on B200 with the CI 65255 binaries (1 passed in 172.28s). The current session registration methods reproduce the local-ID collision and retain both sessions with distinct global IDs. All applicable pre-commit checks passed. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tests/integration/defs/disaggregated/test_workers.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/integration/defs/disaggregated/test_workers.py b/tests/integration/defs/disaggregated/test_workers.py index 8a71bf8d16b8..445141fb76c5 100644 --- a/tests/integration/defs/disaggregated/test_workers.py +++ b/tests/integration/defs/disaggregated/test_workers.py @@ -34,6 +34,7 @@ from transformers import AutoTokenizer from tensorrt_llm import logger +from tensorrt_llm.llmapi.disagg_utils import get_global_disagg_request_id from tensorrt_llm.serve.disagg_auth import build_internal_disagg_auth_headers from tensorrt_llm.serve.openai_client import OpenAIHttpClient from tensorrt_llm.serve.openai_protocol import (CompletionRequest, @@ -186,7 +187,11 @@ async def send_disagg_request(self, session: aiohttp.ClientSession, ctx_request = copy.deepcopy(request) gen_request = copy.deepcopy(request) - ctx_request["disaggregated_params"] = {"request_type": "context_only"} + # Context workers have independent local request ID sequences. + ctx_request["disaggregated_params"] = { + "request_type": "context_only", + "disagg_request_id": get_global_disagg_request_id(node_id=0), + } ctx_response = await self.send_request(session, ctx_url, ctx_request) assert len(ctx_response["choices"]) == 1 From c5ceb28e526bef215a7d86189e7036ca8e649d75 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:26:20 +0000 Subject: [PATCH 23/25] [None][test] Close KV caches in model engine callback tests (By Agent) Register cleanup for the four callback tests that leave active KV caches open. Delayed collection on a thread without a CUDA context otherwise records an invalid CUDA event during cache destruction and poisons KVCM2 for subsequent tests. Validated with a background-GC reproducer (fails before, passes after) and the model engine test file: 104 tests and 58 subtests passed. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tests/unittest/_torch/executor/test_pytorch_model_engine.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 3a7293bd7e9d..78c2cc66ca0b 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -2720,6 +2720,7 @@ def test_forward_pass_callable_on_cuda_graph_on(self): enable_padding=True, )) model_engine, kv_cache_manager = create_model_engine_and_kvcache( llm_args) + self.addCleanup(kv_cache_manager.shutdown) mock_callable = Mock() model_engine.register_forward_pass_callable(mock_callable) @@ -2739,6 +2740,7 @@ def test_forward_pass_callable_on_cuda_graph_on(self): def test_forward_pass_callable_on_cuda_graph_off(self): model_engine, kv_cache_manager = create_model_engine_and_kvcache() + self.addCleanup(kv_cache_manager.shutdown) mock_callable = Mock() model_engine.register_forward_pass_callable(mock_callable) @@ -2758,6 +2760,7 @@ def test_forward_pass_callable_on_cuda_graph_off(self): def test_foward_pass_callable_off(self): model_engine, kv_cache_manager = create_model_engine_and_kvcache() + self.addCleanup(kv_cache_manager.shutdown) self.assertTrue(model_engine._runner.forward_pass_callable is None, "forward_pass_callback should be None by default") @@ -2775,6 +2778,7 @@ def test_foward_pass_callable_off(self): def test_foward_pass_callable_backward_compat(self): model_engine, kv_cache_manager = create_model_engine_and_kvcache() + self.addCleanup(kv_cache_manager.shutdown) self.assertTrue(model_engine._runner.forward_pass_callable is None, "forward_pass_callback should be None by default") From 5d024af033930e72e9742eaa057b81373f55bcbe Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:28:45 +0000 Subject: [PATCH 24/25] [None][test] Allow B300 attention suite to finish cold JIT coverage (By Agent) CI 65280 reached 4347 attention cases without an assertion failure before the aggregate's 3600-second limit terminated the pytest subprocess at 93% completion. The completed cases already consumed 3538 seconds. Give the B300 attention aggregate 90 minutes while preserving its test selection, execution order, exclusions, and assertions. Validation: unchanged YAML configuration apart from the timeout; actual CI parametrization applies a 5400-second marker; pre-commit checks. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b300.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index ce9c0d42d4c9..e7ca0c9988a4 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -34,8 +34,8 @@ l0_b300: # rebind mpi4py's serializer at import time, which applies process-wide on # every stage that collects them, including where the tests themselves skip. # The sparse/dsa/test_cute_dsl_* files ignored below keep their own entries (pre-existing). - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py - # Keep cold-JIT MSA coverage out of the 60-minute attention aggregate. + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp4_paged_mqa_logits.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_tiers.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py TIMEOUT (90) + # Run cold-JIT MSA coverage separately from the attention aggregate. # Cold-JIT runs take about 25 minutes; allow a 45-minute budget. - unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py TIMEOUT (45) - unittest/_torch/attention/sparse/dsa/test_cute_dsl_fp8_paged_mqa_logits.py From 31442501bd73a0ac25b4c8137fa9818748bce30d Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:59:50 +0000 Subject: [PATCH 25/25] [None][fix] Refresh serving schemas for rc2 (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- tensorrt_llm/schemas/trtllm-serve-config.schema.json | 4 ++-- tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json | 4 ++-- .../schemas/trtllm-serve-visual-gen-config.schema.json | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index f5f9184c5557..e09a97ce8afb 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -6285,7 +6285,7 @@ "type": "string" } }, - "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc1/_static/schemas/trtllm-serve-config.schema.json", + "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc2/_static/schemas/trtllm-serve-config.schema.json", "$schema": "https://json-schema.org/draft/2020-12/schema", "additionalProperties": false, "description": "Static YAML checks only; runtime validation remains authoritative.", @@ -7757,6 +7757,6 @@ } }, "required": [], - "title": "TensorRT-LLM 1.4.0rc1 trtllm-serve --config", + "title": "TensorRT-LLM 1.4.0rc2 trtllm-serve --config", "type": "object" } diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index c524c9fe14d4..07d6c962cbd8 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -7846,7 +7846,7 @@ "type": "string" } }, - "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc1/_static/schemas/trtllm-serve-disagg-config.schema.json", + "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc2/_static/schemas/trtllm-serve-disagg-config.schema.json", "$schema": "https://json-schema.org/draft/2020-12/schema", "additionalProperties": false, "description": "Static YAML checks only; runtime validation remains authoritative.", @@ -9502,6 +9502,6 @@ } }, "required": [], - "title": "TensorRT-LLM 1.4.0rc1 trtllm-serve disaggregated --config", + "title": "TensorRT-LLM 1.4.0rc2 trtllm-serve disaggregated --config", "type": "object" } diff --git a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json index cffa1974bcdc..c4f2e4b5320c 100644 --- a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json @@ -980,7 +980,7 @@ "type": "object" } }, - "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc1/_static/schemas/trtllm-serve-visual-gen-config.schema.json", + "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc2/_static/schemas/trtllm-serve-visual-gen-config.schema.json", "$schema": "https://json-schema.org/draft/2020-12/schema", "additionalProperties": false, "description": "Static YAML checks only; runtime validation remains authoritative.", @@ -1104,6 +1104,6 @@ "status": "prototype" } }, - "title": "TensorRT-LLM 1.4.0rc1 trtllm-serve --visual_gen_args", + "title": "TensorRT-LLM 1.4.0rc2 trtllm-serve --visual_gen_args", "type": "object" }