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/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/cache_manager.py index 85d5596065f1..d788e5ea7bdc 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 @@ -16,11 +16,24 @@ from __future__ import annotations +import sys from collections.abc import Sequence import torch +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override + from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAIntermediateState, + KDAReplayState, + _is_disagg_context_server, + bind_kda_replay_pool, + get_kda_replay_num_spec, +) from tensorrt_llm._torch.pyexecutor.config_utils import ( extract_mamba_kv_cache_params, unwrap_glm5_next_text_config, @@ -29,6 +42,7 @@ from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import get_pp_layers 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 BufferConfig, DataRole @@ -44,11 +58,16 @@ class Glm5NextCacheManager(MambaHybridCacheManagerV2): release and disaggregated transfer with the base manager. """ + kda_context_only = False + + @override def __init__( self, *args, sparse_layer_ids: Sequence[int] | None = None, index_state_dim: int = 0, + kda_replay_num_spec: int | None = None, + use_replay_state_update: bool = False, **kwargs, ) -> None: """Build the hybrid manager with one indexer buffer per sparse layer. @@ -77,8 +96,259 @@ def __init__( "glm5_next sparse layers need a positive index_state_dim " f"(got {self.index_state_dim})" ) + if use_replay_state_update: + raise ValueError("GLM KDA does not support Mamba2 replay") + self._mamba_ssm_stochastic_rounding = kwargs.pop("mamba_ssm_stochastic_rounding", False) + self.kda_context_only = _is_disagg_context_server(kwargs.get("is_disagg", 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_spec_state(self) -> KDAIntermediateState | KDAReplayState | 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) + self.kda_context_only = self.kda_context_only and num_spec is not None + if self.kda_context_only: + return None + if num_spec is not None: + 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 KDAIntermediateState() + + @property + def use_kda_replay_update(self) -> bool: + return self._kda_replay is not None + + @override + def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConfig]: + buffers = super()._get_recurrent_buffer_configs(global_layer_id) + if self._kda_replay is not None: + buffers.extend( + self._kda_replay.pool_buffer_configs(self._state_layout, global_layer_id) + ) + return buffers + + @override + def get_disagg_ignored_roles(self) -> frozenset[DataRole]: + roles = super().get_disagg_ignored_roles() + if self._kda_replay is not None: + roles |= self._kda_replay.disagg_ignored_roles(self._state_layout) + return roles + + @override + def _setup_model_state(self) -> None: + if self._kda_replay is not None: + bind_kda_replay_pool(self._kda_replay, self) + super()._setup_model_state() + + @override + def _build_cache_config(self, config): + if ( + self._kda_replay is not None + and self.local_num_mamba_layers > 0 + and self.kv_cache_config.enable_kv_pool_rebalance + ): + raise ValueError( + "KDA replay caches are sized once from the SSM pool's slot count, " + "and a rebalance can grow that pool past them: set " + "kv_cache_config.enable_kv_pool_rebalance=False." + ) + return super()._build_cache_config(config) + + @override + def suspend_request(self, request) -> None: + 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) + 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: + super().prepare_resources(scheduled_batch) + replay = self._kda_replay + if replay is None or self.local_num_mamba_layers == 0: + return + replay.unseeded_requests.update( + request.py_request_id + for request in scheduled_batch.context_requests + if request.is_disagg_generation_init_state + ) + if getattr(self.spec_config, "decoding_type", None) == "NGram": + replay.acceptance_pending.update( + request.py_request_id + for request in scheduled_batch.generation_requests + if not request.is_dummy + ) + + @override + def revert_allocate_generation(self, request) -> None: + if self._kda_replay is not None: + self._kda_replay.acceptance_pending.discard(request.py_request_id) + super().revert_allocate_generation(request) + + @override + def free_resources(self, request, pin_on_release: bool = False) -> None: + if self._kda_replay is not None: + self._kda_replay.suspended_requests.discard(request.py_request_id) + self._kda_replay.acceptance_pending.discard(request.py_request_id) + self._kda_replay.unseeded_requests.discard(request.py_request_id) + super().free_resources(request, pin_on_release) + + @override + def _extra_scratch_bytes_per_slot(self) -> int: + if self._kda_replay is None or not self._kda_replay._stochastic_rounding: + return 0 + return torch.int64.itemsize + + @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: + 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_mamba_states( + self, attn_metadata, num_accepted_tokens, state_indices=None, accepted_leaf_positions=None + ) -> None: + if self.kda_context_only: + return + super().update_mamba_states( + attn_metadata, num_accepted_tokens, state_indices, accepted_leaf_positions + ) + + @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 + replay.acceptance_pending.difference_update( + request.py_request_id for request in 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 + replay.unseeded_requests.difference_update(request_ids) + 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 _extra_buffers_per_layer(self, *, tokens_per_block: int) -> dict[int, list[BufferConfig]]: """One ``Role.INDEX_KEY`` buffer per sparse layer, keyed by local id.""" return { @@ -95,6 +365,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 @@ -106,14 +377,17 @@ 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) 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, *, @@ -123,7 +397,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, @@ -247,3 +521,33 @@ def get_latent_state_buffer(self, layer_idx: int) -> torch.Tensor | None: stride=(stride * flat.stride(0), *flat.stride()[1:]), storage_offset=flat.storage_offset(), ) + + +def get_glm5_next_cache_params(config, *, spec_config=None, quant_config=None): + """Resolve GLM KDA geometry and its required FP32 recurrent state.""" + from tensorrt_llm._torch.pyexecutor.config_utils import ( + build_mamba_kv_cache_params, + get_glm5_next_layer_masks, + ) + + linear = unwrap_glm5_next_text_config(config).linear_attn_config + attention, recurrent = get_glm5_next_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"glm5_next KDA: overriding mamba_ssm_cache_dtype " + f"{params.mamba_ssm_cache_dtype} -> torch.float32" + ) + params.mamba_ssm_cache_dtype = torch.float32 + return params diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py index 89a6410e1cc7..bb11190acf00 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/qsa/__init__.py @@ -1,15 +1,25 @@ # 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", + "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 deleted file mode 100644 index a5ddcaa70614..000000000000 --- a/tensorrt_llm/_torch/attention/backends/sparse/qsa/cache_manager.py +++ /dev/null @@ -1,310 +0,0 @@ -# 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 functools import cached_property -from typing import TYPE_CHECKING, Dict, List, Optional - -import torch - -from tensorrt_llm._torch.disaggregation.resource.page import MapperKind -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, DataRole, PageIndexMode - -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 .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 get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: - """Index positions are request-wide coordinates, so every rank holds - the same bytes. They carry no head axis, and the INDEXED fallback would - divide their size by the rank's KV-head count, which only cancels out - when both peers shard heads the same way.""" - return { - **super().get_disagg_role_mapper_kinds(), - QSA_INDEX_POSITION: MapperKind.REPLICATED, - } - - @cached_property - def qsa_local_sparse_layer_ids(self) -> List[int]: - """Sparse layers this rank owns, in model order.""" - return [ - layer_id for layer_id in self.qsa_sparse_layer_ids if layer_id in self.layer_offsets - ] - - @cached_property - def qsa_shared_topk_slots(self) -> Dict[int, int]: - """Row each sparse layer captures into in the shared MTP Top-K buffers. - - Every sparse layer runs its own indexer projection over its own index - cache, so a captured selection is only valid for the layer that - produced it. - """ - return {layer_id: slot for slot, layer_id in enumerate(self.qsa_local_sparse_layer_ids)} - - 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 = self.qsa_local_sparse_layer_ids - # 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"] diff --git a/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/qsa/metadata.py index fe1a441e3584..0203af3e7331 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,6 +44,8 @@ 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 Qwen4ExpHybridCacheManager + super().__post_init__() # Draft-loop state; set by the speculative driver, read by the indexer. self.qsa_in_mtp_draft_loop = False @@ -53,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, QSAMambaHybridCacheManagerV2): - raise TypeError("QSA sparse attention requires QSAMambaHybridCacheManagerV2") + 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 581299c02cc7..7d7e1e2f8c38 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/registry.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/registry.py @@ -30,7 +30,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 == "csa2": @@ -50,7 +49,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 Qwen4ExpHybridCacheManager + + return Qwen4ExpHybridCacheManager else: raise ValueError( f"Unsupported sparse attention algorithm: {sparse_attention_config.algorithm}" diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index f244d1cf2763..31dcace88254 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -18,7 +18,7 @@ Caching ------- -KDA states use ``MambaHybridCacheManagerV2`` with the ``KvCacheConfig`` +KDA states use ``KimiK3HybridCacheManager`` with the ``KvCacheConfig`` default ``use_kv_cache_manager_v2=True`` (selected in ``pyexecutor/_util.py``). Setting this field to ``"auto"`` falls back to V1 for ``KimiLinearForCausalLM``, which declares no V2 model preference. diff --git a/tensorrt_llm/_torch/models/modeling_qwen4_exp.py b/tensorrt_llm/_torch/models/modeling_qwen4_exp.py index c5350231a38f..c2e54cba74ec 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 @@ -493,7 +494,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..f962bbdc1deb --- /dev/null +++ b/tensorrt_llm/_torch/modules/fla/cache_manager.py @@ -0,0 +1,643 @@ +# 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, field +from typing import NamedTuple + +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 ( + MambaAcceptanceBatch, + MambaLayerCache, + MambaStateLayout, + _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.utils import is_gdn_replay_enabled +from tensorrt_llm._utils import get_sm_version, prefer_pinned +from tensorrt_llm.logger import logger + +from .cached_replay import ( + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE, + 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.""" + + 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 + self._state_descriptors: torch.Tensor | None = None + self._state_strides: tuple[int, int, int] | None = None + + @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: + 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.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 + + 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, 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] + 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 + 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", + ) + + 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 GDNLayerCache( + 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: + num_decodes = batch.num_generations + if num_decodes >= CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: + self.commit_all_layers(batch.attention_metadata, num_decodes) + metadata = self.get_replay_metadata() + if metadata is None: + raise RuntimeError("GDN replay buffers are not bound") + _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.""" + 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, + ) + + 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 + 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, "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 +) -> GDNIntermediateState | GDNReplayState: + """Shared selection for the Qwen3.5 and Qwen4 cache managers.""" + if use_replay is False: + return GDNIntermediateState() + 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 GDNIntermediateState() + ) + + +def validate_gdn_layout( + 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( + manager._state_layout.mapping + ): + raise ValueError("GDN replay groups must be divisible by the effective TP size") + state.validate(manager._state_layout) + + +class Qwen35HybridCacheManager(MambaHybridCacheManagerV2): + """Qwen3.5/Qwen3Next recurrent state with shared GDN replay.""" + + @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_spec_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 + return isinstance(state, GDNReplayState) and state.has_bound_states + + +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__ = [ + "GDNIntermediateState", + "GDNLayerCache", + "GDNReplayState", + "Qwen35HybridCacheManager", + "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..a927aa4bafb5 --- /dev/null +++ b/tensorrt_llm/_torch/modules/kimi_kda/cache_manager.py @@ -0,0 +1,1260 @@ +# 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 math +import os +import sys +from contextlib import contextmanager +from dataclasses import dataclass, field + +import torch + +from tensorrt_llm._utils import prefer_pinned + +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 ( + MambaAcceptanceBatch, + MambaLayerCache, + MambaRole, + MambaStateLayout, + _get_local_mamba_cache_layout, + _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.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, DataRole + +_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): + """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_k_cache: torch.Tensor | None = None + kda_g_cache: torch.Tensor | None = None + kda_v_cache: torch.Tensor | None = None + kda_beta_cache: torch.Tensor | None = None + # [slots, 3, num_spec, heads, head_dim]: fp32 vn, beta * k, and decay. + kda_state_tok: torch.Tensor | None = None + + @property + def has_kda_replay_caches(self) -> bool: + return self.kda_conv_q 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, + ) + + +_KDA_BETA_CACHE_ALIGNMENT_BYTES = 16 +# KVCM copies 16-byte payloads; the replay kernel reads only the first int32. +_KDA_REPLAY_COUNT_BYTES = 16 +_KDA_REPLAY_COUNT_ROLE = DataRole("prev_num_accepted_tokens") +_KDA_REPLAY_CONV_ROLES = ("kda_conv_q", "kda_conv_k", "kda_conv_v") +_KDA_DRAFT_RECORDS_ROLE = "kda_state_tok" + + +def _is_disagg_context_server(is_disagg: bool) -> bool: + """Whether this worker only serves the context phase of disagg requests.""" + return is_disagg and os.getenv("TRTLLM_DISAGG_ROLE") == "context" + + +def _kda_replay_layout( + num_spec: int, num_heads: int, head_dim: int, conv_window: int, draft_records: bool = False +) -> dict[str, tuple[torch.dtype, list[int]]]: + """Per-slot physical shapes of one layer's KDA replay caches. + + Conv caches hold raw bf16 projection inputs (exact in bf16); k / v / beta + are bf16 and the gate stays fp32 because it feeds the cumulative-gate + exponents of the replayed chunk. Beta rows keep the 16-byte CuTe + alignment by padding the head dimension. + + With ``draft_records``, the records of every draft of the last verify + round (``kda_state_tok``: per draft, the row innovations vn, then beta * k + and the decay, each fp32 [H, K] with V = K = ``head_dim``) take the place + of the k / g / v / beta caches. + """ + dim = num_heads * head_dim + layout = { + "kda_conv_q": (torch.bfloat16, [conv_window + num_spec, dim]), + "kda_conv_k": (torch.bfloat16, [conv_window + num_spec, dim]), + "kda_conv_v": (torch.bfloat16, [conv_window + num_spec, dim]), + } + if draft_records: + layout[_KDA_DRAFT_RECORDS_ROLE] = (torch.float32, [3, num_spec, num_heads, head_dim]) + return layout + heads_per_alignment = _KDA_BETA_CACHE_ALIGNMENT_BYTES // torch.bfloat16.itemsize + aligned_heads = ( + (num_heads + heads_per_alignment - 1) // heads_per_alignment + ) * heads_per_alignment + layout.update( + { + "kda_k_cache": (torch.bfloat16, [num_spec, dim]), + "kda_g_cache": (torch.float32, [num_spec, dim]), + "kda_v_cache": (torch.bfloat16, [num_spec, dim]), + "kda_beta_cache": (torch.bfloat16, [num_spec, aligned_heads]), + } + ) + return layout + + +def _kda_replay_bytes_per_slot( + layout: dict[str, tuple[torch.dtype, list[int]]], num_local_layers: int +) -> int: + """Replay bytes one recurrent-state slot holds across local layers.""" + if not layout or num_local_layers == 0: + return 0 + layer_bytes = sum(dtype.itemsize * math.prod(shape) for dtype, shape in layout.values()) + return num_local_layers * layer_bytes + _KDA_REPLAY_COUNT_BYTES + + +def _host_slots_to_device(slots: list[int], device: torch.device) -> torch.Tensor: + """Copy host slot IDs without synchronizing queued GPU work. + + Each copy uses fresh pinned storage so later host writes cannot race it. + """ + return torch.tensor( + slots, dtype=torch.long, pin_memory=device.type == "cuda" and prefer_pinned() + ).to(device, non_blocking=True) + + +def _kda_beta_cache_padded_heads(num_heads: int) -> int: + heads_per_alignment = _KDA_BETA_CACHE_ALIGNMENT_BYTES // torch.bfloat16.itemsize + return (num_heads + heads_per_alignment - 1) // heads_per_alignment * heads_per_alignment + + +def _allocate_kda_beta_cache(shape: tuple[int, ...], device: torch.device | None) -> torch.Tensor: + """Keep logical head count while aligning each physical row for CuTe.""" + return torch.zeros( + *shape[:-1], + _kda_beta_cache_padded_heads(shape[-1]), + dtype=torch.bfloat16, + device=device, + )[..., : shape[-1]] + + +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, + draft_records: bool = False, +) -> 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.bfloat16, + 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(), + } + if draft_records: + fields["kda_state_tok"] = torch.zeros( + num_local_layers, + cache_size, + 3, + num_speculative_tokens, + num_heads, + section_dim // num_heads, + dtype=torch.float32, + device=device, + ) + return fields, [ + value for name, value in fields.items() if name != "prev_num_accepted_tokens" + ] + fields.update( + { + "kda_k_cache": torch.zeros( + num_local_layers, + cache_size, + num_speculative_tokens, + section_dim, + dtype=torch.bfloat16, + device=device, + ), + "kda_g_cache": torch.zeros( + num_local_layers, + cache_size, + num_speculative_tokens, + section_dim, + dtype=torch.float32, + device=device, + ), + "kda_v_cache": torch.zeros( + num_local_layers, + cache_size, + num_speculative_tokens, + section_dim, + dtype=torch.bfloat16, + device=device, + ), + "kda_beta_cache": _allocate_kda_beta_cache( + (num_local_layers, cache_size, num_speculative_tokens, num_heads), device + ), + } + ) + scratch = [ + fields["kda_conv_q"], + fields["kda_conv_k"], + fields["kda_conv_v"], + fields["kda_k_cache"], + fields["kda_g_cache"], + fields["kda_v_cache"], + fields["kda_beta_cache"], + ] + 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.""" + + @property + def intermediate_ssm(self) -> torch.Tensor | None: + return None + + @property + def intermediate_conv(self) -> torch.Tensor | None: + return None + + def __init__( + self, + num_speculative_tokens: int, + *, + stochastic_rounding: bool = False, + draft_records: bool = False, + ) -> None: + self._stochastic_rounding = stochastic_rounding + self.draft_records = draft_records + self.pool_layers: list[dict[str, torch.Tensor]] | None = None + self.num_speculative_tokens = num_speculative_tokens + self.suspended_requests: set[int] = set() + self.acceptance_pending: set[int] = set() + self.unseeded_requests: set[int] = set() + self._ssm_states: tuple[torch.Tensor, ...] = () + 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_k_cache: torch.Tensor | None = None + self.kda_g_cache: torch.Tensor | None = None + self.kda_v_cache: torch.Tensor | None = None + self.kda_beta_cache: torch.Tensor | None = None + self.kda_state_tok: 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 pool_layout(self, context: MambaStateLayout) -> dict[str, tuple[torch.dtype, list[int]]]: + if not context.mamba_pp_layers: + return {} + 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" + ) + if self.draft_records and context.ssm_state_shape[1] != context.ssm_state_shape[-1]: + raise ValueError("KDA draft records require equal value and key head dimensions") + return _kda_replay_layout( + self.num_speculative_tokens, + context.ssm_state_shape[0], + context.ssm_state_shape[1], + context.conv_state_shape[1], + draft_records=self.draft_records, + ) + + def pool_buffer_configs(self, context: MambaStateLayout, layer_id: int) -> list[BufferConfig]: + layout = self.pool_layout(context) + buffers = [ + BufferConfig(role=DataRole(name), size=dtype.itemsize * math.prod(shape)) + for name, (dtype, shape) in layout.items() + ] + if layout and layer_id == context.mamba_pp_layers[0]: + buffers.append(BufferConfig(role=_KDA_REPLAY_COUNT_ROLE, size=_KDA_REPLAY_COUNT_BYTES)) + return buffers + + def disagg_ignored_roles(self, context: MambaStateLayout) -> frozenset[DataRole]: + roles = {DataRole(name) for name in self.pool_layout(context)} + if roles: + roles.add(_KDA_REPLAY_COUNT_ROLE) + return frozenset(roles) + + def _layer_bytes_per_slot(self, section_dim: int, committed_window: int, num_heads: int) -> int: + extended_window = committed_window + self.num_speculative_tokens + if self.draft_records: + return ( + 3 * section_dim * extended_window * torch.bfloat16.itemsize + + 3 * self.num_speculative_tokens * section_dim * torch.float32.itemsize + ) + bf16_elements = ( + 3 * section_dim * extended_window + + self.num_speculative_tokens * 2 * section_dim + + self.num_speculative_tokens * _kda_beta_cache_padded_heads(num_heads) + ) + gate_elements = self.num_speculative_tokens * section_dim + return bf16_elements * torch.bfloat16.itemsize + gate_elements * torch.float32.itemsize + + 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 ( + 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.""" + if _is_disagg_context_server(kwargs.get("is_disagg", False)): + return 0 + 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, draft_records=kwargs.get("kda_draft_records", False)) + return _kda_replay_bytes_per_slot( + _kda_replay_layout( + num_spec, + params.num_heads // tp_size, + params.head_dim, + params.conv_kernel - 1, + draft_records=state.draft_records, + ), + local_layers, + ) + (8 if stochastic_rounding else 0) + + def bind( + self, + context: MambaStateLayout, + ssm_states: list[torch.Tensor], + conv_states: list[torch.Tensor], + ) -> None: + """Borrow persistent views and bind bookkeeping; standalone use allocates replay storage.""" + 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._ssm_states = tuple(ssm_states) + 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 + if self.pool_layers is None: + 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, + draft_records=self.draft_records, + ) + 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 self._stochastic_rounding: + self.rand_seed = allocate_mamba_seed_buffer( + cache_size, + self._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]: + if self.pool_layers is not None: + return { + **self.pool_layers[layer_offset], + "prev_num_accepted_tokens": self.prev_num_accepted_tokens, + } + return { + "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_k_cache": None if self.kda_k_cache is None else self.kda_k_cache[layer_offset], + "kda_g_cache": None if self.kda_g_cache is None else self.kda_g_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] + ), + "kda_state_tok": None + if self.kda_state_tok is None + else self.kda_state_tok[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 + slots = _host_slots_to_device(host_slots, self.prev_num_accepted_tokens.device) + self.prev_num_accepted_tokens.index_fill_(0, 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 + if self.pool_layers is not None and self.rand_seed 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 = _host_slots_to_device([old for old, _ in moves], device) + destination_slots = _host_slots_to_device([new for _, new in moves], device) + if self.pool_layers is None: + 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) + + self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0) + for layer_offset, conv_state in enumerate(self._conv_states): + replay = self._layer_cache_fields(layer_offset) + for name in ( + *_KDA_REPLAY_CONV_ROLES, + "kda_k_cache", + "kda_g_cache", + "kda_v_cache", + "kda_beta_cache", + ): + buffer = replay.get(name) + if buffer is not None: + buffer.index_fill_(0, state_indices, 0) + selected_conv = conv_state.index_select(0, state_indices) + for section_idx, name in enumerate(_KDA_REPLAY_CONV_ROLES): + buffer = replay[name] + section = selected_conv[ + :, section_offsets[section_idx] : section_offsets[section_idx + 1], : + ] + buffer[:, :, :committed_window].index_copy_( + 0, state_indices, section.to(buffer.dtype) + ) + + 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 + + @staticmethod + @contextmanager + def on_cache_stream(stream): + """Order replay commits and seeding against V2 page migration.""" + current = torch.cuda.current_stream() + stream.wait_stream(current) + with torch.cuda.stream(stream): + yield + current.wait_stream(stream) + + 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 + + assert self.prev_num_accepted_tokens is not None + slots = torch.full( + (1,), slot, dtype=torch.int32, device=self.prev_num_accepted_tokens.device + ) + num_accepted = self.prev_num_accepted_tokens.index_select(0, slots) + for layer, (ssm_states, conv_states) in enumerate(zip(self._ssm_states, self._conv_states)): + replay = self._layer_cache_fields(layer) + commit_kda_pending_drafts( + ssm_states, + conv_states, + replay["kda_conv_q"], + replay["kda_conv_k"], + replay["kda_conv_v"], + replay.get("kda_k_cache"), + replay.get("kda_g_cache"), + replay.get("kda_v_cache"), + replay.get("kda_beta_cache"), + slots, + num_accepted, + state_tok=replay.get(_KDA_DRAFT_RECORDS_ROLE), + ) + self.prev_num_accepted_tokens[slot] = 0 + + def _layer_replay_buffers(self) -> tuple[torch.Tensor, ...]: + buffers = ( + self.kda_conv_q, + self.kda_conv_k, + self.kda_conv_v, + self.kda_k_cache, + self.kda_g_cache, + self.kda_v_cache, + self.kda_beta_cache, + self.kda_state_tok, + ) + if any(buffer is None for buffer in buffers[:3]): + 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.pool_layers = None + self._ssm_states = () + self._conv_states = () + self.suspended_requests.clear() + self.acceptance_pending.clear() + self.unseeded_requests.clear() + self.prev_num_accepted_tokens = None + self.kda_conv_q = None + self.kda_conv_k = None + self.kda_conv_v = None + self.kda_k_cache = None + self.kda_g_cache = None + self.kda_v_cache = None + self.kda_beta_cache = None + self.kda_state_tok = None + self.intermediate_indices = None + self.rand_seed = None + + +def bind_kda_replay_pool(state: KDAReplayState, manager: MambaHybridCacheManagerV2) -> None: + """Bind model-owned KDA replay roles to views of the manager's native pools.""" + layout = state.pool_layout(manager._state_layout) + if not layout: + return + required = { + (manager.layer_offsets[layer_id], str(role)) + for layer_id in manager.mamba_pp_layers + for role in (MambaRole.SSM_STATE, MambaRole.CONV_STATE, *layout) + } + required.add((manager.layer_offsets[manager.mamba_pp_layers[0]], str(_KDA_REPLAY_COUNT_ROLE))) + if not any( + required.issubset( + { + (int(buffer_id.layer_id), str(buffer_id.role)) + for coalesced in variant.coalesced_buffers + for buffer_id in coalesced.buffer_ids + } + ) + for group in manager.impl.pool_group_descs + for variant in group.slot_desc.variants + ): + raise RuntimeError( + "KDA replay and recurrent-state buffers must share one KVCM " + "pool group and lifecycle to reuse the same logical slot index" + ) + + state.pool_layers = [] + for layer_id in manager.mamba_pp_layers: + layer = {} + for name, (dtype, shape) in layout.items(): + view = manager._get_state_buffer( + manager.layer_offsets[layer_id], DataRole(name), dtype, shape + ) + if name in _KDA_REPLAY_CONV_ROLES: + view = view.transpose(-1, -2) + elif name == "kda_beta_cache": + view = view[..., : manager.ssm_state_shape[0]] + if view.shape[0] != manager._state_layout.slot_capacity: + raise RuntimeError(f"KDA replay buffer {name} has a different slot capacity") + layer[name] = view + state.pool_layers.append(layer) + state.prev_num_accepted_tokens = manager._get_state_buffer( + manager.layer_offsets[manager.mamba_pp_layers[0]], + _KDA_REPLAY_COUNT_ROLE, + torch.int32, + [_KDA_REPLAY_COUNT_BYTES // torch.int32.itemsize], + )[..., 0] + + +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 kda_mtp_replay_num_spec + + return kda_mtp_replay_num_spec(spec_config.tokens_per_gen_step - 1) + + +class KimiK3HybridCacheManager(MambaHybridCacheManagerV2): + """KDA/MLA cache layout and K3-only replay lifecycle handling.""" + + kda_context_only = False + + @override + def __init__( + self, + *args, + kda_replay_num_spec: int | None = None, + kda_draft_records: bool = False, + use_replay_state_update: bool = False, + **kwargs, + ) -> 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.kda_context_only = _is_disagg_context_server(kwargs.get("is_disagg", False)) + self._requested_num_spec = kda_replay_num_spec + self._requested_draft_records = kda_draft_records + self._kda_replay: KDAReplayState | None = None + kwargs.setdefault("conv_state_layout", "q_k_v") + super().__init__(*args, **kwargs) + + @override + def _initialize_spec_state(self) -> KDAIntermediateState | KDAReplayState | 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) + self.kda_context_only = self.kda_context_only and num_spec is not None + if self.kda_context_only: + return None + if num_spec is not None: + self._kda_replay = KDAReplayState( + num_spec, + stochastic_rounding=self._mamba_ssm_stochastic_rounding, + draft_records=self._requested_draft_records, + ) + self._kda_replay.validate(self._state_layout) + return self._kda_replay + return KDAIntermediateState() + + @property + def use_kda_replay_update(self) -> bool: + return self._kda_replay is not None + + @property + def keeps_kda_draft_records(self) -> bool: + return self.kda_state_tok is not None + + @property + def kda_state_tok(self) -> list[torch.Tensor] | None: + replay = self._kda_replay + if replay is None or not replay.pool_layers: + return None + if _KDA_DRAFT_RECORDS_ROLE not in replay.pool_layers[0]: + return None + return [layer[_KDA_DRAFT_RECORDS_ROLE] for layer in replay.pool_layers] + + def _kda_replay_layout(self) -> dict[str, tuple[torch.dtype, list[int]]]: + if self._kda_replay is None: + return {} + return self._kda_replay.pool_layout(self._state_layout) + + @override + def _get_recurrent_buffer_configs(self, global_layer_id: int) -> list[BufferConfig]: + buffers = super()._get_recurrent_buffer_configs(global_layer_id) + if self._kda_replay is not None: + buffers.extend( + self._kda_replay.pool_buffer_configs(self._state_layout, global_layer_id) + ) + return buffers + + @override + def get_disagg_ignored_roles(self) -> frozenset[DataRole]: + roles = super().get_disagg_ignored_roles() + if self._kda_replay is not None: + roles |= self._kda_replay.disagg_ignored_roles(self._state_layout) + return roles + + @override + def _setup_model_state(self) -> None: + if self._kda_replay is not None: + bind_kda_replay_pool(self._kda_replay, self) + super()._setup_model_state() + + @override + def _build_cache_config(self, config): + if ( + self._kda_replay is not None + and self.local_num_mamba_layers > 0 + and self.kv_cache_config.enable_kv_pool_rebalance + ): + raise ValueError( + "KDA replay views are bound once to the SSM pool's slot count, " + "and a rebalance can grow that pool past them: set " + "kv_cache_config.enable_kv_pool_rebalance=False." + ) + return super()._build_cache_config(config) + + @override + def suspend_request(self, request) -> None: + 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) + 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( + _host_slots_to_device(slots, replay.prev_num_accepted_tokens.device) + ) + + @override + def prepare_resources(self, scheduled_batch) -> None: + super().prepare_resources(scheduled_batch) + replay = self._kda_replay + if replay is None or self.local_num_mamba_layers == 0: + return + replay.unseeded_requests.update( + request.py_request_id + for request in scheduled_batch.context_requests + if request.is_disagg_generation_init_state + ) + if getattr(self.spec_config, "decoding_type", None) == "NGram": + replay.acceptance_pending.update( + request.py_request_id + for request in scheduled_batch.generation_requests + if not request.is_dummy + ) + + @override + def revert_allocate_generation(self, request) -> None: + if self._kda_replay is not None: + self._kda_replay.acceptance_pending.discard(request.py_request_id) + super().revert_allocate_generation(request) + + @override + def free_resources(self, request, pin_on_release: bool = False) -> None: + if self._kda_replay is not None: + self._kda_replay.suspended_requests.discard(request.py_request_id) + self._kda_replay.acceptance_pending.discard(request.py_request_id) + self._kda_replay.unseeded_requests.discard(request.py_request_id) + super().free_resources(request, pin_on_release) + + @override + def _extra_scratch_bytes_per_slot(self) -> int: + if self._kda_replay is None or not self._kda_replay._stochastic_rounding: + return 0 + return torch.int64.itemsize + + @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: + 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 = _host_slots_to_device(fresh_slots, self.cuda_state_indices.device) + self._kda_replay.reset_slots(slots, fresh_slots) + + @override + def update_mamba_states( + self, attn_metadata, num_accepted_tokens, state_indices=None, accepted_leaf_positions=None + ) -> None: + if self.kda_context_only: + return + super().update_mamba_states( + attn_metadata, num_accepted_tokens, state_indices, accepted_leaf_positions + ) + + @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 + replay.acceptance_pending.difference_update( + request.py_request_id for request in 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 + replay.unseeded_requests.difference_update(request_ids) + 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( + _host_slots_to_device(slots, 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) + + +def resolve_kimi_ssm_cache_dtype(config) -> torch.dtype: + """Default KDA state to FP32; BF16 requires an explicit cache configuration. + + BF16 prefill, decode and sequential verify stage state through FP32 and + round the committed state back to BF16. Fused MTP verify reads and writes + BF16 state directly, rounding the committed state to nearest. + Checkpoint-declared state dtypes do not override this default. + """ + from tensorrt_llm._torch.pyexecutor.config_utils import resolve_ssm_cache_dtype + + declared = resolve_ssm_cache_dtype(config) + if declared is not None and declared != torch.float32: + logger.info( + "Kimi K3: the checkpoint declares " + f"mamba_ssm_cache_dtype={declared}; keeping the fp32 " + "recurrent-state pool (kv_cache_config.mamba_ssm_cache_dtype " + "opts in to bfloat16)." + ) + return torch.float32 + + +def validate_kimi_state_cache_dtype( + mamba_ssm_cache_dtype: torch.dtype, mamba_ssm_stochastic_rounding: bool = False +) -> None: + """Reject cache dtypes and rounding modes unsupported by the KDA kernels.""" + if mamba_ssm_cache_dtype not in (torch.float32, torch.bfloat16): + raise ValueError( + "Kimi K3 KDA recurrent-state cache supports float32 (default) or " + f"bfloat16; got mamba_ssm_cache_dtype={mamba_ssm_cache_dtype}." + ) + if mamba_ssm_stochastic_rounding and mamba_ssm_cache_dtype != torch.float32: + raise ValueError( + "Kimi K3 KDA kernels round the committed recurrent state to " + "nearest; mamba_ssm_stochastic_rounding is not supported with a " + f"{mamba_ssm_cache_dtype} state cache." + ) + + +def get_kimi_cache_params(config, *, spec_config=None, quant_config=None): + """KDA uses equal Q/K/V sections and defaults to FP32 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, + ) + return params + + +__all__ = [ + "KDAIntermediateLayerCache", + "KDAIntermediateState", + "KDAReplayState", + "KDAReplayLayerCache", + "KimiK3HybridCacheManager", + "allocate_kda_replay_fields", + "bind_kda_replay_pool", + "get_kda_replay_num_spec", + "get_kimi_cache_params", + "resolve_kimi_ssm_cache_dtype", + "select_kda_replay_state", + "validate_kimi_state_cache_dtype", +] 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 30231196346d..3ade8c8fb9f1 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -553,7 +553,7 @@ def _project_output(self, core: Union[torch.Tensor, _KdaPrequantizedCore]) -> to 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..36e2a2b0d236 --- /dev/null +++ b/tensorrt_llm/_torch/modules/mamba/cache_manager.py @@ -0,0 +1,554 @@ +# 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 +from typing import NamedTuple + +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 ( + MambaAcceptanceBatch, + MambaLayerCache, + MambaStateLayout, + _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.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: + """Intermediate promotion plus Mamba2 stochastic-rounding seeds.""" + + uses_replay = False + + 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 + + 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_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], self._seed_rank_offset, states.device + ) + + 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, + ) + + 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] + if self.rand_seed is not None: + fields["mamba_ssm_rand_seed"] = self.rand_seed + return fields + + 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 + 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.""" + + @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.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 + + 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, 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] + 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 Mamba2LayerCache( + 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") + _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, 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(stochastic_rounding=stochastic_rounding) + + +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, "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, stochastic_rounding=stochastic_rounding + ) + + +class NemotronHybridCacheManager(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._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) + + @override + def _initialize_spec_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, + stochastic_rounding=self._mamba_ssm_stochastic_rounding, + ) + else: + 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(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") + 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 + + +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__ = [ + "Mamba2LayerCache", + "Mamba2State", + "ReplayHistory", + "NemotronHybridCacheManager", + "create_mamba2_state", + "select_mamba2_state", + "get_nemotron_cache_params", + "mamba_seed_rank_offset", +] diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index df2dcd3ae34e..af9e8645aa48 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -1174,14 +1174,12 @@ def forward_core( 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 efead8e0d799..cb0c04b96674 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -401,7 +401,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 @@ -419,17 +423,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 6a001d8f1242..9c41b9227e1e 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -647,15 +647,11 @@ def forward_core( 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..b996ad00c0ac --- /dev/null +++ b/tensorrt_llm/_torch/modules/qwen4_exp/cache_manager.py @@ -0,0 +1,622 @@ +# 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 functools import cached_property +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.disaggregation.resource.page import MapperKind +from tensorrt_llm._torch.modules.fla.cache_manager import ( + GDNIntermediateState, + GDNReplayState, + ReplayStateUpdateMetadata, + 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 ( + _get_local_mamba_cache_layout, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import get_pp_layers +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, DataRole, PageIndexMode + +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 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. + 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 get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: + """PLE state and request-wide QSA positions are replicated across TP.""" + return { + **super().get_disagg_role_mapper_kinds(), + PLE_CONV_STATE: MapperKind.REPLICATED, + PLE_NGRAM_CONTEXT: MapperKind.REPLICATED, + QSA_INDEX_POSITION: MapperKind.REPLICATED, + } + + @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: + 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 + 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 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") + 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) + + @cached_property + def qsa_local_sparse_layer_ids(self) -> list[int]: + """Sparse layers this rank owns, in model order.""" + return [ + layer_id for layer_id in self.qsa_sparse_layer_ids if layer_id in self.layer_offsets + ] + + @cached_property + def qsa_shared_topk_slots(self) -> dict[int, int]: + """Row each sparse layer captures into in the shared MTP Top-K buffers. + + Every sparse layer runs its own indexer projection over its own index + cache, so a captured selection is only valid for the layer that + produced it. + """ + return {layer_id: slot for slot, layer_id in enumerate(self.qsa_local_sparse_layer_ids)} + + @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 = self.qsa_local_sparse_layer_ids + # 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, + ) + + 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, + 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_view_by_role_and_layer( + 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_view_by_role_and_layer( + 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") + 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 + 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: " + 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_spec_state(self) -> GDNIntermediateState | GDNReplayState: + state = create_gdn_state(self, self._requested_replay) + + validate_gdn_layout(self, 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 + ) + return state + + @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: + super()._setup_model_state() + if self.local_num_mamba_layers: + self._setup_ple_states(self._state_layout.slot_capacity) + + @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(self) -> None: + self._ple_conv_states.clear() + self._ple_ngram_contexts.clear() + super().shutdown() + + 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 + return isinstance(state, GDNReplayState) and state.has_bound_states + + +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__ = [ + "Qwen4ExpHybridCacheManager", + "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 72b0e12823ab..38be9b86837e 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -53,11 +53,9 @@ 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 . import config_utils from .config_utils import (MambaKVCacheParams, _is_sliding_attention_layer, extract_mamba_kv_cache_params, - extract_qwen4_exp_ple_cache_params, get_layer_attention_window, is_hybrid_linear, is_kimi_linear, is_mla, is_nemotron_hybrid, is_qwen3_hybrid, is_qwen4_exp, resolve_vocab_size, @@ -313,7 +311,7 @@ def get_kv_cache_manager_cls( "hybrid managers do not implement per-token verify-group " "bookkeeping.") - # Kimi uses MambaHybridCacheManagerV2 with or without reuse when + # Kimi uses KimiK3HybridCacheManager with or without reuse when # use_kv_cache_manager_v2=True, the KvCacheConfig default. V1 can # be selected explicitly or resolved from "auto". Aggregated V1 # uses Mixed without reuse and the C++ hybrid manager with reuse. @@ -431,7 +429,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 \ + KimiK3HybridCacheManager + return KimiK3HybridCacheManager + if is_qwen4_exp(config): + from ..modules.qwen4_exp.cache_manager import \ + Qwen4ExpHybridCacheManager + return Qwen4ExpHybridCacheManager + if is_qwen3_hybrid(config): + 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 @@ -3212,6 +3222,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", } @@ -3230,29 +3241,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, @@ -3501,40 +3489,18 @@ 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. - kda_extra_kwargs = {} - kda_replay_manager_types = (MixedMambaHybridCacheManager, - MambaHybridCacheManagerV2) + kda_extra_kwargs = _mamba_conv_layout_kwargs(kv_cache_manager_cls, + "kimi_linear") + backbone = getattr(getattr(model_engine, "model", None), "model", None) if (spec_config is not None - and issubclass(kv_cache_manager_cls, kda_replay_manager_types)): - from ..modules.kimi_kda._kda_kernels import kda_mtp_replay_num_spec - kda_replay_num_spec = kda_mtp_replay_num_spec( - spec_config.tokens_per_gen_step - 1) - if kda_replay_num_spec is not None: - kda_extra_kwargs["kda_replay_num_spec"] = kda_replay_num_spec - # A model whose KDA verify replays the accepted drafts from - # its own per-draft records (instead of the replay caches) - # asks for a buffer of them with `kda_draft_records`. - backbone = getattr(getattr(model_engine, "model", None), - "model", None) - if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) - and getattr(backbone, "kda_draft_records", False)): - kda_extra_kwargs["kda_draft_records"] = True + and issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) + and getattr(backbone, "kda_draft_records", False)): + kda_extra_kwargs["kda_draft_records"] = True if is_glm5_next: - # The manager places an indexer buffer on every attention layer. from ..attention.backends.sparse.glm_kpool import \ glm_kpool_cache_row_dim kda_extra_kwargs["index_state_dim"] = glm_kpool_cache_row_dim( int(text_config.index_head_dim)) - # KDA's conv state is a [Q | K | V] concatenation whose three sections - # have identical width, i.e. the qwen3_next section layout. - kda_extra_kwargs.update( - _mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next")) kv_cache_manager = kv_cache_manager_cls( # mamba (KDA) cache parameters mamba_params.state_size, @@ -3630,54 +3596,9 @@ 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 - - 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, '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") - # 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 @@ -3717,7 +3638,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, ) @@ -3744,75 +3665,14 @@ 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, '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")) - 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, @@ -3841,7 +3701,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/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index 24b094deb4d4..f1281b8a9123 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -326,34 +326,12 @@ def resolve_ssm_cache_dtype(config): return None -KIMI_KDA_STATE_DTYPES = (torch.float32, torch.bfloat16) - - def resolve_auto_ssm_cache_dtype(config, fallback): - """Resolve mamba_ssm_cache_dtype="auto" for ``config``. - - Kimi K3 defaults to fp32: the HF reference (fla chunk/fused_recurrent - KDA kernels) carries the delta-rule recurrent state in fp32. A bf16 - state pool is an explicit opt-in through - kv_cache_config.mamba_ssm_cache_dtype; prefill, decode and sequential - verify then stage the addressed rows through an fp32 copy and round the - committed state back to bf16. The fused MTP verify kernel - (``trtllm::kda_mtp_decode``) reads and writes a bf16 pool directly and - rounds the committed state to nearest. - A checkpoint-declared - mamba_ssm_cache_dtype is not applied to Kimi K3 (the released - checkpoints do not carry the field); it is logged when it would have - changed the dtype. - """ + """Resolve the model's automatic recurrent-state dtype.""" if is_kimi_linear(config): - declared = resolve_ssm_cache_dtype(config) - if declared is not None and declared != torch.float32: - logger.info( - "Kimi K3: the checkpoint declares " - f"mamba_ssm_cache_dtype={declared}; keeping the fp32 " - "recurrent-state pool (kv_cache_config.mamba_ssm_cache_dtype " - "opts in to bfloat16).") - return torch.float32 + from ..modules.kimi_kda.cache_manager import \ + resolve_kimi_ssm_cache_dtype + return resolve_kimi_ssm_cache_dtype(config) return (resolve_ssm_cache_dtype(config) or resolve_hf_torch_dtype(config) or fallback) @@ -361,23 +339,12 @@ def resolve_auto_ssm_cache_dtype(config, fallback): def validate_kimi_kda_state_dtype(config, mamba_ssm_cache_dtype, mamba_ssm_stochastic_rounding=False): - """Reject state-cache settings the Kimi K3 KDA kernels cannot honor. - - The kernels read fp32 or bf16 state and round the committed state to - nearest; stochastic rounding is a Mamba2 fp16-cache feature that would - otherwise be accepted and silently ignored here. - """ - if not is_kimi_linear(config): - return - if mamba_ssm_cache_dtype not in KIMI_KDA_STATE_DTYPES: - raise ValueError( - "Kimi K3 KDA recurrent-state cache supports float32 (default) or " - f"bfloat16; got mamba_ssm_cache_dtype={mamba_ssm_cache_dtype}.") - if mamba_ssm_stochastic_rounding and mamba_ssm_cache_dtype != torch.float32: - raise ValueError( - "Kimi K3 KDA kernels round the committed recurrent state to " - "nearest; mamba_ssm_stochastic_rounding is not supported with a " - f"{mamba_ssm_cache_dtype} state cache.") + """Delegate KDA state-cache validation to the model owner.""" + if is_kimi_linear(config): + from ..modules.kimi_kda.cache_manager import \ + validate_kimi_state_cache_dtype + validate_kimi_state_cache_dtype(mamba_ssm_cache_dtype, + mamba_ssm_stochastic_rounding) def resolve_vocab_size(config) -> Optional[int]: @@ -540,61 +507,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, - ) - - def mamba_effective_tp_size(mapping) -> int: """TP degree for sizing per-rank mamba/KDA state (budgeting AND allocation). @@ -685,74 +597,39 @@ 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 elif is_glm5_next(config): - # GLM-5.3-Flash KDA state uses the same mamba parametrization as - # Kimi K3: [q | k | v] short-conv sections of identical width and an - # [H, V, K] fp32 delta-rule recurrent state, with the KDA geometry - # coming from the text config's linear_attn_config. - lin = unwrap_glm5_next_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_glm5_next_layer_masks(config) + from ..attention.backends.sparse.glm_kpool.cache_manager import \ + get_glm5_next_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 @@ -768,10 +645,6 @@ def extract_mamba_kv_cache_params( mamba_ssm_cache_dtype = resolve_auto_ssm_cache_dtype( config, torch.bfloat16) validate_kimi_kda_state_dtype(config, mamba_ssm_cache_dtype) - if is_glm5_next(config) and mamba_ssm_cache_dtype != torch.float32: - logger.info(f"glm5_next KDA: overriding mamba_ssm_cache_dtype " - f"{mamba_ssm_cache_dtype} -> torch.float32") - 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 07a5e93717aa..2a7cd971d760 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 @@ -1758,19 +1758,21 @@ 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), + max_gpu_total_bytes=kv_cache_config.max_gpu_total_bytes, ) else: init_error: Optional[Exception] = None local_init_status = _KVCacheManagerInitStatus.OK try: - candidate = KVCacheManagerPy( + candidate = self._allocate_cache_pools( 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: init_error = error.with_traceback(None) @@ -1786,6 +1788,14 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: raise RuntimeError("KV cache manager initialization failed on another rank") assert candidate is not None + if self._get_external_cache_bytes(candidate): + config = replace( + config, + cache_tiers=[ + GpuCacheTierConfig(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 @@ -3085,6 +3095,55 @@ def _build_base_config( ), ) + 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, + max_gpu_total_bytes: int | None = None, + ) -> KVCacheManagerPy: + """Fit native pools and external storage within the physical GPU budget.""" + cache = KVCacheManagerPy( + config, event_manager=event_manager, cold_page_codec=cold_page_codec + ) + external_bytes = self._get_external_cache_bytes(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.""" return config 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 b763ee4066c5..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ /dev/null @@ -1,5143 +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 contextlib import contextmanager -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._device_copy import DeviceCopyBuffer -from tensorrt_llm._torch.disaggregation.resource.page import (MapperKind, - RoleLayout) -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( - _RESERVED_REQUEST_IDS, 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.executor import KvCacheRetentionConfig -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) - -# Shared with the KV budget estimator so allocator and budgeting can never -# diverge on the sharding rule (config_utils is import-cycle-free). -from ..config_utils import is_glm5_next, is_kimi_linear, mamba_effective_tp_size - -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 - -_KDA_BETA_CACHE_ALIGNMENT_BYTES = 16 -# KVCM copies 16-byte payloads; the replay kernel reads only the first int32. -_KDA_REPLAY_COUNT_BYTES = 16 -_KDA_REPLAY_COUNT_ROLE = DataRole("prev_num_accepted_tokens") -_KDA_REPLAY_CONV_ROLES = ("kda_conv_q", "kda_conv_k", "kda_conv_v") -_KDA_DRAFT_RECORDS_ROLE = "kda_state_tok" - - -def _is_disagg_context_server(is_disagg: bool) -> bool: - """Whether this worker only serves the context phase of disagg requests.""" - return is_disagg and os.getenv("TRTLLM_DISAGG_ROLE") == "context" - - -def _kda_replay_layout( - num_spec: int, - num_heads: int, - head_dim: int, - conv_window: int, - draft_records: bool = False -) -> Dict[str, Tuple[torch.dtype, List[int]]]: - """Per-slot physical shapes of one layer's KDA replay caches. - - Conv caches hold raw bf16 projection inputs (exact in bf16); k / v / beta - are bf16 and the gate stays fp32 because it feeds the cumulative-gate - exponents of the replayed chunk. Beta rows keep the 16-byte CuTe - alignment by padding the head dimension. - - With ``draft_records``, the records of every draft of the last verify - round (``kda_state_tok``: per draft, the row innovations vn, then beta * k - and the decay, each fp32 [H, K] with V = K = ``head_dim``) take the place - of the k / g / v / beta caches. - """ - dim = num_heads * head_dim - layout = { - "kda_conv_q": (torch.bfloat16, [conv_window + num_spec, dim]), - "kda_conv_k": (torch.bfloat16, [conv_window + num_spec, dim]), - "kda_conv_v": (torch.bfloat16, [conv_window + num_spec, dim]), - } - if draft_records: - layout[_KDA_DRAFT_RECORDS_ROLE] = (torch.float32, - [3, num_spec, num_heads, head_dim]) - return layout - heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // - torch.bfloat16.itemsize) - aligned_heads = ((num_heads + heads_per_alignment - 1) // - heads_per_alignment) * heads_per_alignment - layout.update({ - "kda_k_cache": (torch.bfloat16, [num_spec, dim]), - "kda_g_cache": (torch.float32, [num_spec, dim]), - "kda_v_cache": (torch.bfloat16, [num_spec, dim]), - "kda_beta_cache": (torch.bfloat16, [num_spec, aligned_heads]), - }) - return layout - - -def _kda_replay_bytes_per_slot(layout: Dict[str, Tuple[torch.dtype, List[int]]], - num_local_layers: int) -> int: - """Replay bytes one recurrent-state slot holds across local layers.""" - if not layout or num_local_layers == 0: - return 0 - layer_bytes = sum(dtype.itemsize * math.prod(shape) - for dtype, shape in layout.values()) - return num_local_layers * layer_bytes + _KDA_REPLAY_COUNT_BYTES - - -def _allocate_kda_beta_cache( - shape: Tuple[int, ...], - device: Optional[torch.device], -) -> torch.Tensor: - """Allocate a logically shaped bf16 beta cache with aligned physical rows.""" - logical_num_heads = shape[-1] - heads_per_alignment = (_KDA_BETA_CACHE_ALIGNMENT_BYTES // - torch.bfloat16.itemsize) - padded_num_heads = ((logical_num_heads + heads_per_alignment - 1) // - heads_per_alignment) * heads_per_alignment - # CuTe requires every nested cache view to be 16-byte aligned. Keep the - # logical head count while padding the physical stride (six bf16 heads are - # 12 bytes, so an unpadded row makes alternating views only 4-byte aligned). - return torch.zeros( - *shape[:-1], - padded_num_heads, - dtype=torch.bfloat16, - device=device, - )[..., :logical_num_heads] - - -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. ``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 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 _host_slots_to_device(slots: List[int], - device: torch.device) -> torch.Tensor: - """Copy host slot ids to ``device`` without a host sync. - - A pageable copy would wait for the queued step. Each call stages through a - fresh pinned tensor, so no later host write can race the copy. - """ - return torch.tensor(slots, - dtype=torch.long, - pin_memory=device.type == "cuda" - and prefer_pinned()).to(device, non_blocking=True) - - -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) - - @property - def keeps_kda_draft_records(self) -> bool: - """Whether the KDA replay caches come with the records of every draft - of the last verify round (``kda_state_tok``), from which a verify - kernel replays the accepted drafts onto the slot's state.""" - return 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); or, - for a verify kernel that keeps the drafts' records, kda_conv_* and - kda_state_tok in place of the 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] bf16 - # (raw projection inputs, so bf16 is exact), 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: k - # [layers, slots, num_spec, H*K] bf16, gate [.., num_spec, H*K] fp32 - # (feeds the cumulative-gate exponents), v [.., num_spec, H*V] bf16, - # beta [.., num_spec, H] bf16. q is recomputed from the conv caches. - kda_k_cache: torch.Tensor | None = None - kda_g_cache: torch.Tensor | None = None - kda_v_cache: torch.Tensor | None = None - kda_beta_cache: torch.Tensor | None = None - # Optional: the records of every draft of the last verify round, - # [slots, 3, num_spec, H, K] fp32: the row innovations vn [num_spec, - # H, V], then beta * k and the decay [num_spec, H, K] (V = K). The - # SSM pool holds the state after the round's first, non-draft token; - # when the round accepted n > 0 drafts (prev_num_accepted_tokens), a - # verify kernel starts the next round by replaying the first n - # records onto it, S = decay * S + vn (beta * k), instead of - # replaying them from the caches above, which are then not - # allocated. A slot reset for a new request has n = 0, so a previous - # owner's records are never read. - kda_state_tok: 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 on the KDA fused-verify replay path: the extended conv - caches are allocated, with the per-draft replay caches of the - fused ``trtllm::kda_mtp_decode`` verify or, in their place, the - drafts' records (``kda_state_tok``).""" - return self.kda_conv_q 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(): - # raw bf16 projection inputs: bf16 is exact and halves - # the per-draft replay traffic of the verify kernel - return torch.zeros(num_local_layers, - max_batch_size, - extended_s, - section_dim, - dtype=torch.bfloat16, - 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_k_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - section_dim, - dtype=torch.bfloat16, - device=device) - spec_kwargs['kda_g_cache'] = torch.zeros(num_local_layers, - max_batch_size, - M, - 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.bfloat16, - device=device) - spec_kwargs['kda_beta_cache'] = _allocate_kda_beta_cache( - (num_local_layers, max_batch_size, M, nheads), device) - ssm_spec_cache = [ - spec_kwargs['kda_conv_q'], spec_kwargs['kda_conv_k'], - spec_kwargs['kda_conv_v'], spec_kwargs['kda_k_cache'], - spec_kwargs['kda_g_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_k_cache, self.mamba_cache.kda_g_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)) - mask_staging = torch.empty_like(self._dummy_request_mask_host, - pin_memory=prefer_pinned()) - mask_staging.copy_(self._dummy_request_mask_host) - self._dummy_request_mask.copy_(mask_staging, 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_k_cache=_drop(self.mamba_cache.kda_k_cache), - kda_g_cache=_drop(self.mamba_cache.kda_g_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 - - @property - def cuda_state_indices(self) -> torch.Tensor: - return self._state_indices_buffer.tensor - - @property - def _dummy_request_mask(self) -> Optional[torch.Tensor]: - if self._dummy_request_mask_buffer is None: - return None - return self._dummy_request_mask_buffer.tensor - - 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_buffer = 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_buffer.update(self._dummy_request_mask_host) - - 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.index_fill_(0, context_slots, 0) - self.cache_buf_idx.index_fill_(0, context_slots, 0) - if self.old_x is not None: - self.old_x.index_fill_(1, context_slots, 0) - if self.old_B is not None: - self.old_B.index_fill_(1, context_slots, 0) - if self.old_dt is not None: - self.old_dt.index_fill_(1, context_slots, 0) - if self.old_dA_cumsum is not None: - self.old_dA_cumsum.index_fill_(1, 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, - pin_memory=prefer_pinned()).to(self.mamba_ssm_rand_seed.device, - non_blocking=True) - - 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. - """ - - # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the - # cache manager __init__ coerces a positive spec recompute tail to a full - # re-prefill (see _has_recurrent_state on the attention-only base classes). - _has_recurrent_state = True - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - # kv cache 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 _estimated_kda_replay_num_spec(model_config, spec_config, - is_disagg: bool) -> Optional[int]: - """Replay-cache width a V2 manager for this model would allocate. - - Mirrors the executor's construction-time choice: KDA models (Kimi K3 and - GLM-5 Next) under speculative decoding get replay caches when the fused - verify kernel can run here (``kda_mtp_replay_num_spec``), except on a - disaggregated context server, which never verifies drafts. - """ - if spec_config is None or _is_disagg_context_server(is_disagg): - return None - pretrained_config = getattr(model_config, "pretrained_config", None) - if pretrained_config is None or not (is_kimi_linear(pretrained_config) - or is_glm5_next(pretrained_config)): - return None - from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ - kda_mtp_replay_num_spec - - return kda_mtp_replay_num_spec(spec_config.tokens_per_gen_step - 1) - - -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, - kda_replay_num_spec: Optional[int] = None, - **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, - ) - if kda_replay_num_spec is not None: - # KDA replay caches live in the same KVCM slot as the recurrent state, - # so every live, dummy and snapshot slot carries them. - tp_size = mamba_effective_tp_size(mapping) - state_bytes_per_rank += _kda_replay_bytes_per_slot( - _kda_replay_layout(kda_replay_num_spec, params.num_heads // tp_size, - params.head_dim, params.conv_kernel - 1), - local_mamba_layers) - 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. - - """ - - # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the - # cache manager __init__ coerces a positive spec recompute tail to a full - # re-prefill (see _has_recurrent_state on the attention-only base classes). - _has_recurrent_state = True - - def __init__( - self, - # mamba cache parameters - mamba_d_state: int, - mamba_d_conv: int, - mamba_num_heads: int, - mamba_n_groups: int, - mamba_head_dim: int, - mamba_num_layers: int, - mamba_layer_mask: List[bool], - mamba_cache_dtype: torch.dtype, - mamba_ssm_cache_dtype: torch.dtype, - kv_cache_config: KvCacheConfig, - kv_cache_type: CacheTypeCpp, - *, - num_layers: int, - num_kv_heads: Union[int, List[Optional[int]]], - head_dim: int, - tokens_per_block: int, - # 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._state_indices_buffer = DeviceCopyBuffer( - torch.zeros([self.max_batch_size], dtype=torch.int32, - device="cuda"), - pin_memory=prefer_pinned()) - 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_buffer = 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, - 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, - 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._state_indices_buffer.update(self._host_state_indices, force=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_buffer = DeviceCopyBuffer( - torch.zeros(self.max_batch_size, dtype=torch.bool, device=device), - pin_memory=prefer_pinned()) - 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 - - # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the - # cache manager __init__ coerces a positive spec recompute tail to a full - # re-prefill (see _has_recurrent_state on the attention-only base classes). - _has_recurrent_state = True - - # 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] = [] - # Set per instance on a disaggregated context server of a KDA model. - kda_context_only = False - - def _disagg_transfer_overwrites_whole_cached_prefix(self) -> bool: - # The incoming recurrent state overwrites the whole local slot and summarizes the - # complete prefix, so none of those tokens count as a local cache hit. - return (self.local_num_mamba_layers > 0 - or super()._disagg_transfer_overwrites_whole_cached_prefix()) - - 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, - kda_draft_records: bool = False, - **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") - - # A disaggregated context server never verifies drafts, so it keeps - # neither KDA replay caches nor intermediate verification buffers. - self.kda_context_only = bool( - kda_replay_num_spec is not None - and _is_disagg_context_server(kwargs.get("is_disagg", False))) - if self.kda_context_only: - kda_replay_num_spec = None - self._kda_replay_layers: List[Dict[str, torch.Tensor]] = [] - self._kda_replay_num_spec = kda_replay_num_spec - self._use_kda_replay_update = kda_replay_num_spec is not None - # Only with the KDA replay caches: also keep the records of every - # draft of the last verify round (kda_state_tok). - self._kda_draft_records = kda_draft_records and self._use_kda_replay_update - 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 = {} - # Requests whose suspend committed their KDA replay state into the V2 - # states; their next generation step re-seeds the replay caches. - self._kda_replay_suspended: set[int] = set() - # With a host drafter: requests of a prepared batch whose acceptance - # update_resources has not recorded yet. - self._kda_replay_acceptance_pending: set[int] = set() - self._kda_replay_unseeded: set[int] = set() - - state_index_capacity = (self.max_batch_size + - self._num_reserved_dummy_slots) - self._state_indices_buffer = DeviceCopyBuffer( - torch.zeros([state_index_capacity], - dtype=torch.int32, - device="cuda"), - pin_memory=prefer_pinned()) - 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._stacked_ssm_states = None - self._stacked_conv_states = None - 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 - - def get_disagg_role_mapper_kinds(self) -> Dict[DataRole, MapperKind]: - """Recurrent roles and how they reshard across TP. - - Convolution state is SECTIONED: its flat per-layer buffer is a - concatenation of sections (Mamba2 ``[x | B | C]``, GDN ``[Q | K | V]``) - that are each TP-sharded independently. SSM state keeps the INDEXED - default (sharded by head). PLE state is computed from replicated - inputs, so every rank holds identical bytes and the transfer copies - whole per-layer regions. - """ - return { - **super().get_disagg_role_mapper_kinds(), - MambaRole.CONV_STATE: - MapperKind.SECTIONED, - MambaRole.PLE_CONV_STATE: - MapperKind.REPLICATED, - MambaRole.PLE_NGRAM_CONTEXT: - MapperKind.REPLICATED, - } - - def get_disagg_role_layouts(self) -> Dict[DataRole, RoleLayout]: - """Per-layer geometry for resharding conv and SSM state.""" - if self.local_num_mamba_layers == 0: - return {} - d_conv_m1 = int(self.conv_state_shape[1]) - conv_elem_size = self.conv_state_dtype.itemsize - _, head_dim, d_state = self.ssm_state_shape - return { - MambaRole.CONV_STATE: - RoleLayout(section_bytes=tuple( - int(dim) * d_conv_m1 * conv_elem_size - for dim in self.conv_section_dims)), - MambaRole.SSM_STATE: - RoleLayout(bytes_per_head=int(head_dim) * int(d_state) * - self.ssm_state_dtype.itemsize), - } - - def get_disagg_ignored_roles(self) -> frozenset[DataRole]: - """Keep the KDA replay buffers out of disaggregated transfers. - - They are generation-local: a context server allocates none, and a - generation server seeds them from the transferred conv state. - """ - replay_roles = {DataRole(name) for name in self._kda_replay_layout()} - if replay_roles: - replay_roles.add(_KDA_REPLAY_COUNT_ROLE) - return super().get_disagg_ignored_roles() | frozenset(replay_roles) - - @property - def use_gdn_cached_replay_all_layer_commit(self) -> bool: - return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) - - @property - def kda_state_tok(self) -> Optional[List[torch.Tensor]]: - """Per local KDA layer, the drafts' records of every SSM slot - ([slots, 3, num_spec, H, K] fp32 KVCM state-buffer views), or None - when the replay path keeps the per-draft replay caches instead.""" - replay_layers = getattr(self, "_kda_replay_layers", None) - if not replay_layers or _KDA_DRAFT_RECORDS_ROLE not in replay_layers[0]: - return None - return [layer[_KDA_DRAFT_RECORDS_ROLE] for layer in replay_layers] - - @property - def keeps_kda_draft_records(self) -> bool: - return self.kda_state_tok is not None - - def _setup_mtp_intermediate_states(self, spec_config, - max_batch_size: int) -> None: - if self.kda_context_only: - super()._setup_mtp_intermediate_states(None, max_batch_size) - return - 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 _kda_replay_layout(self) -> Dict[str, Tuple[torch.dtype, List[int]]]: - """Per-slot shapes of the KDA replay buffers of one local layer.""" - if not self.use_kda_replay_update or self.local_num_mamba_layers == 0: - return {} - 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") - assert self._kda_replay_num_spec is not None - draft_records = getattr(self, "_kda_draft_records", False) - layout = _kda_replay_layout(self._kda_replay_num_spec, - self.ssm_state_shape[0], - self.ssm_state_shape[1], - self.conv_state_shape[1], - draft_records=draft_records) - if draft_records and tuple(layout[_KDA_DRAFT_RECORDS_ROLE][1]) != ( - self._kda_draft_records_shape()): - raise ValueError("KDA draft records require equal value and key " - "head dimensions") - return layout - - def _kda_replay_buffer_configs(self) -> List[BufferConfig]: - return [ - BufferConfig(role=DataRole(name), - size=dtype.itemsize * math.prod(shape)) - for name, (dtype, shape) in self._kda_replay_layout().items() - ] - - def _validate_kda_replay_colocation(self, - replay_roles: Iterable[str]) -> None: - """Verify replay and recurrent buffers share one slot lifecycle. - - Every layer addresses its replay caches with the recurrent state's - logical slot index, so all of them must be allocated, moved and - offloaded together: they must sit in one variant of one KVCM pool - group, not merely have equal capacities. - """ - required_buffers = {(self.layer_offsets[layer_id], str(role)) - for layer_id in self.mamba_pp_layers - for role in (MambaRole.SSM_STATE, - MambaRole.CONV_STATE, *replay_roles)} - required_buffers.add((self.layer_offsets[self.mamba_pp_layers[0]], - str(_KDA_REPLAY_COUNT_ROLE))) - if not any( - required_buffers.issubset( - {(int(buffer_id.layer_id), str(buffer_id.role)) - for coalesced in variant.coalesced_buffers - for buffer_id in coalesced.buffer_ids}) - for pool_group in self.impl.pool_group_descs - for variant in pool_group.slot_desc.variants): - raise RuntimeError( - "KDA replay and recurrent-state buffers must share one KVCM " - "pool group and lifecycle to reuse the same logical slot index") - - 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_replay_layers = [] - layout = self._kda_replay_layout() - if not layout: - return allocated - self._validate_kda_replay_colocation(layout) - # The replay caches are KVCM state buffers next to the SSM and conv - # states, so they follow the slot through suspend, offload and - # snapshot reuse. Pool capacity can exceed resident memory, so rows - # are initialized only once a request holds its slot. - for layer_id in self.mamba_pp_layers: - local_layer_idx = self.layer_offsets[layer_id] - layer = {} - for name, (dtype, shape) in layout.items(): - view = self._get_state_buffer(local_layer_idx, DataRole(name), - dtype, shape) - if name in _KDA_REPLAY_CONV_ROLES: - # Dim-contiguous [slots, dim, W - 1 + num_spec]. - view = view.transpose(-1, -2) - elif name == "kda_beta_cache": - view = view[..., :self.ssm_state_shape[0]] - if view.shape[0] != cache_size: - raise RuntimeError( - f"KDA replay buffer {name} of layer {layer_id} has " - f"{view.shape[0]} slots, but the recurrent state has " - f"{cache_size}") - layer[name] = view - self._kda_replay_layers.append(layer) - # All SSM layers share a slot lifecycle, so the accepted-draft count is - # stored once, on the first local KDA layer. - self.prev_num_accepted_tokens = self._get_state_buffer( - self.layer_offsets[self.mamba_pp_layers[0]], - _KDA_REPLAY_COUNT_ROLE, - torch.int32, - [_KDA_REPLAY_COUNT_BYTES // torch.int32.itemsize], - )[..., 0] - records = (", with the drafts' records in place of the replay " - "caches" if self.keeps_kda_draft_records else "") - logger.info("Mamba Cache (kda-replay) is bound to KVCM state buffers " - f"for {cache_size} state slots{records}") - 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, - kv_cache_retention_config: Optional[KvCacheRetentionConfig] = 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, - kv_cache_retention_config=kv_cache_retention_config, - ) - 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, - kda_replay_num_spec=_estimated_kda_replay_num_spec( - model_config, spec_config, kwargs.get("is_disagg", False)), - **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 _kda_draft_records_shape(self) -> Tuple[int, int, int, int]: - """One SSM slot's shape in a layer of ``kda_state_tok``: per draft, - the row innovations vn, then beta * k and the decay, each [H, K] (the - replay caches' equal [Q | K | V] sections make V = K).""" - return (3, self._kda_replay_num_spec, self.ssm_state_shape[0], - self.ssm_state_shape[-1]) - - def _mamba_state_bytes_per_slot(self) -> int: - base_bytes = (self.local_num_mamba_layers * - (self.ssm_bytes + self.conv_bytes) + - _kda_replay_bytes_per_slot(self._kda_replay_layout(), - self.local_num_mamba_layers)) - 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 - if (getattr(self, "_use_kda_replay_update", False) - and self.local_num_mamba_layers > 0 - and kv_cache_config.enable_kv_pool_rebalance): - # The KDA replay views and prev_num_accepted_tokens are bound - # once to the SSM pool's slot count at construction; a rebalance - # can grow the SSM pool past them, and the verify kernels index - # them by slot without a bound. - raise ValueError( - "The KDA replay caches (speculative decoding on KDA layers) " - "are sized once from the SSM pool's slot count, and a KV pool " - "rebalance can grow that pool past them: set " - "kv_cache_config.enable_kv_pool_rebalance=False.") - 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) - kda_replay_buffers = self._kda_replay_buffer_configs() - first_mamba_layer = (self.mamba_pp_layers[0] - if kda_replay_buffers else 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 = [ - 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()) - if kda_replay_buffers: - buffers.extend(kda_replay_buffers) - if global_layer_idx == first_mamba_layer: - # All SSM layers share a slot lifecycle, so the - # request's accepted-draft count is stored once and - # moves with the replay histories. - buffers.append( - BufferConfig(role=_KDA_REPLAY_COUNT_ROLE, - size=_KDA_REPLAY_COUNT_BYTES)) - 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 - ] - self._stacked_ssm_states = self._stack_state_views(self.all_ssm_states) - self._stacked_conv_states = self._stack_state_views( - self.all_conv_states) - 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, - ) - - @staticmethod - def _stack_state_views( - states: List[torch.Tensor]) -> Optional[torch.Tensor]: - """Returns one ``[num_layers, num_slots, *state_shape]`` view of ``states``. - - A stacked view collapses per-layer promotion loop into a single launch of _promote_mamba_state_triton. - This requires per-layer pool views to be *affine*: identical dtype/device/shape/strides, and base - pointers spaced by a constant. returns ``None`` otherwise. - """ - if not states: - return None - ref = states[0] - if len(states) == 1: - return ref.unsqueeze(0) - for state in states[1:]: - if (state.dtype != ref.dtype or state.device != ref.device - or state.shape != ref.shape - or state.stride() != ref.stride()): - return None - itemsize = ref.element_size() - layer_stride_bytes = states[1].data_ptr() - states[0].data_ptr() - if layer_stride_bytes <= 0 or layer_stride_bytes % itemsize != 0: - return None - base = states[0].data_ptr() - if any(state.data_ptr() != base + layer * layer_stride_bytes - for layer, state in enumerate(states)): - return None - layer_stride = layer_stride_bytes // itemsize - # Highest element one layer's view can reach, so the flat wrapper below - # covers every byte the stacked view addresses and no more. - layer_extent = 1 + sum((size - 1) * stride - for size, stride in zip(ref.shape, ref.stride())) - total_elems = (len(states) - 1) * layer_stride + layer_extent - flat = convert_to_torch_tensor( - TensorWrapper(base, ref.dtype, [total_elems])) - return flat.as_strided([len(states)] + list(ref.shape), - [layer_stride] + list(ref.stride())) - - 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_buffer = DeviceCopyBuffer( - torch.zeros(mask_capacity, dtype=torch.bool, device=device), - pin_memory=prefer_pinned()) - 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: - # Mirror the base: the guard page, when enabled, is held by a - # permanent reservation made at construction (not by a request), so - # exclude reserved ids before asserting the manager is otherwise empty. - reserved = set(self.kv_cache_map) & _RESERVED_REQUEST_IDS - assert not set(self.kv_cache_map) - reserved, ( - "get_num_free_blocks is only used when the kv cache manager is empty" - ) - 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[Tuple[int, 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 global_layer_id, KVCacheManagerV2.get_buffers( - self, global_layer_id) - - # SSM / conv / PLE tensors are not attention layers and hold no guard - # page, so yield a sentinel layer id (-1). The base consumer looks the - # id up via self._guard_page_by_layer.get(layer_id), which returns None - # for an unknown id and correctly skips guard restoration for them. - for buffer in self.all_ssm_states: - yield -1, buffer - for buffer in self.all_conv_states: - yield -1, buffer - for buffer in self._ple_ngram_contexts.values(): - yield -1, buffer - for buffer in self._ple_conv_states.values(): - yield -1, buffer - - 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 = _host_slots_to_device( - self._host_state_indices[:num_contexts].tolist(), - self.prev_num_accepted_tokens.device) - self.prev_num_accepted_tokens.index_fill_(0, 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 - device = self.prev_num_accepted_tokens.device - if isinstance(slots, torch.Tensor) and slots.device.type != "cpu": - slots = slots.to(device=device, dtype=torch.long, non_blocking=True) - else: - if isinstance(slots, torch.Tensor): - slots = slots.tolist() - if not slots: - return - slots = _host_slots_to_device(slots, device) - if slots.numel() > 0: - self.prev_num_accepted_tokens.index_fill_(0, slots, 0) - - def suspend_request(self, req: LlmRequest) -> None: - """Suspend the request's KV cache with its whole KDA state in V2. - - The drafts the last verify accepted are pending in the slot's KDA - replay buffers, and so is the conv window, which the fused verify does - not write to the V2 conv state. The pending drafts and the conv window - are committed into the slot's SSM and conv states here, before the - pages are released, so the suspended state is self-contained in the - recurrent state; the request's next generation step re-seeds the - replay buffers of whichever slot it resumes in from them. - """ - request_id = req.py_request_id - kv_cache = self.kv_cache_map.get(request_id) - slot = self._request_id_to_state_index.get(request_id, -1) - if (self.use_kda_replay_update and slot >= 0 and kv_cache is not None - and kv_cache.is_active): - # A host drafter's acceptance is recorded by update_resources. - # Executors record it before they can suspend the request: the - # host drafter runs without the overlap scheduler, and the PP - # loop does not evict requests in flight. - if request_id in self._kda_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 self._kda_replay_unseeded: - self._kda_replay_unseeded.remove(request_id) - else: - self._commit_kda_pending_drafts(slot) - del self._request_id_to_state_index[request_id] - self._kda_replay_suspended.add(request_id) - super().suspend_request(req) - - @contextmanager - def _on_kv_cache_stream(self): - """Run the enclosed work on the KV caches' stream, after the current - stream's earlier work and before its later work. - - V2 orders page migration against that stream: a suspend records the - event later migrations of the pages wait for, and a resume makes the - stream wait for the pages it brings back. - """ - current = torch.cuda.current_stream() - self._stream.wait_stream(current) - with torch.cuda.stream(self._stream): - yield - current.wait_stream(self._stream) - - def _commit_kda_pending_drafts(self, slot: int) -> None: - """Commit the slot's pending accepted drafts and its conv window into - its V2 SSM and conv states, in place, and clear its draft count. - - Afterwards the V2 states hold what the next verify would have started - from: the conv window exactly, the SSM state to the verify kernel's - accuracy (``commit_kda_pending_drafts``). - """ - from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ - commit_kda_pending_drafts - - assert self.prev_num_accepted_tokens is not None - with self._on_kv_cache_stream(): - slots = torch.full((1, ), - slot, - dtype=torch.int32, - device=self.prev_num_accepted_tokens.device) - num_accepted = self.prev_num_accepted_tokens.index_select(0, slots) - for layer_offset, (ssm_states, conv_states) in enumerate( - zip(self.all_ssm_states, self.all_conv_states)): - replay = self._kda_replay_layers[layer_offset] - # With the verify kernels' draft records (kda_state_tok), the - # accepted drafts are replayed from those, and the k / g / v / - # beta replay caches are not allocated. - commit_kda_pending_drafts( - ssm_states, - conv_states, - replay["kda_conv_q"], - replay["kda_conv_k"], - replay["kda_conv_v"], - replay.get("kda_k_cache"), - replay.get("kda_g_cache"), - replay.get("kda_v_cache"), - replay.get("kda_beta_cache"), - slots, - num_accepted, - state_tok=replay.get(_KDA_DRAFT_RECORDS_ROLE), - ) - self.prev_num_accepted_tokens[slot] = 0 - - def add_dummy_requests( - self, - request_ids: List[int], - token_nums: Optional[List[int]] = None, - is_gen: bool = False, - 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, - 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._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) - self._kda_replay_suspended.discard(request.py_request_id) - self._kda_replay_acceptance_pending.discard(request.py_request_id) - self._kda_replay_unseeded.discard(request.py_request_id) - 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) - if self.use_kda_replay_update: - self._kda_replay_unseeded.update( - request.py_request_id - for request in scheduled_batch.context_requests - if request.is_disagg_generation_init_state) - if self._record_kda_replay_in_update_resources: - # Recorded by this batch's update_resources, or dropped by - # revert_allocate_generation when the batch is not run. - self._kda_replay_acceptance_pending.update( - request.py_request_id - for request in scheduled_batch.generation_requests - if not request.is_dummy) - - def revert_allocate_generation(self, req: LlmRequest) -> None: - self._kda_replay_acceptance_pending.discard(req.py_request_id) - super().revert_allocate_generation(req) - - 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._state_indices_buffer.update(self._host_state_indices) - 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: - # A suspended request that comes back as a context request starts - # over. One that comes back generating has its whole KDA state in - # the V2 states of the slot it resumes in (suspend_request). A - # live request whose slot moved carries its replay buffers along: - # they are KVCM state buffers of the same slot. - for request in requests[:num_contexts]: - self._kda_replay_suspended.discard(request.py_request_id) - new_state_values = [ - self._request_id_to_state_index[request.py_request_id] - for request in replay_requests - ] - new_slots = [] - resumed_slots = [] - for request, old, new in zip(replay_requests, old_state_values, - new_state_values): - if request.py_request_id in self._kda_replay_suspended: - self._kda_replay_suspended.remove(request.py_request_id) - resumed_slots.append(new) - elif old < 0: - new_slots.append(new) - self._reset_kda_replay_slots(new_slots) - if resumed_slots: - resumed_indices = torch.tensor(resumed_slots, - dtype=torch.long, - pin_memory=prefer_pinned()) - with self._on_kv_cache_stream(): - self._seed_kda_replay_slots( - resumed_indices.to(self.prev_num_accepted_tokens.device, - non_blocking=True)) - - 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 - ] - self._kda_replay_unseeded.difference_update(request_ids) - if not slots: - return - - assert self.prev_num_accepted_tokens is not None - state_indices = _host_slots_to_device( - sorted(set(slots)), self.prev_num_accepted_tokens.device) - self._seed_kda_replay_slots(state_indices) - - def _seed_kda_replay_slots(self, state_indices: torch.Tensor) -> None: - """Seed the replay caches of distinct slots from their V2 conv - states, with no drafts pending.""" - assert self.prev_num_accepted_tokens is not None - 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) - - self.prev_num_accepted_tokens.index_fill_(0, state_indices, 0) - for replay_layer, conv_state in zip(self._kda_replay_layers, - self.all_conv_states): - for name, replay_buffer in replay_layer.items(): - # The drafts' records need no reset: with no pending drafts, - # none of them is read. - if name != _KDA_DRAFT_RECORDS_ROLE: - replay_buffer.index_fill_(0, state_indices, 0) - selected_conv = conv_state.index_select(0, state_indices) - for section_idx, name in enumerate(_KDA_REPLAY_CONV_ROLES): - replay_buffer = replay_layer[name] - section = selected_conv[:, section_offsets[section_idx]: - section_offsets[section_idx + 1], :] - replay_buffer[:, :, :committed_window].index_copy_( - 0, state_indices, section.to(replay_buffer.dtype)) - - 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 - spec_kwargs = dict(self._kda_replay_layers[layer_offset]) - spec_kwargs["prev_num_accepted_tokens"] = self.prev_num_accepted_tokens - 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 - self._kda_replay_acceptance_pending.difference_update( - request.py_request_id for request in 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 or self.kda_context_only: - 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: - self._promote_state(self.all_ssm_states, self._stacked_ssm_states, - self.intermediate_ssm_states, src_state_indices, - accepted_positions, state_indices_d) - - self._promote_state(self.all_conv_states, self._stacked_conv_states, - self.intermediate_conv_states, src_state_indices, - accepted_positions, state_indices_d) - - @staticmethod - def _promote_state(per_layer: List[torch.Tensor], - stacked: Optional[torch.Tensor], - intermediate: torch.Tensor, - src_state_indices: torch.Tensor, - accepted_positions: torch.Tensor, - state_indices_d: torch.Tensor) -> None: - """Promote every layer's accepted state, in one launch where possible. - - The kernel's grid is ``(num_layers * num_gens, tiles)``, so handing it - the stacked view does all layers at once; the per-layer loop is only - the fallback for a non-affine pool layout. - """ - if stacked is not None: - _promote_mamba_state_triton(stacked, intermediate, - src_state_indices, accepted_positions, - state_indices_d) - return - for layer_offset, dst in enumerate(per_layer): - _promote_mamba_state_triton( - dst.unsqueeze(0), - intermediate[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._stacked_ssm_states = None - self._stacked_conv_states = None - 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_buffer = 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_replay_layers = [] - self._kda_replay_suspended.clear() - self._kda_replay_acceptance_pending.clear() - self._kda_replay_unseeded.clear() - 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..cd1ebe8bb592 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/__init__.py @@ -0,0 +1,78 @@ +# 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 + 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"), + "MambaCacheManager": ("legacy", "MambaCacheManager"), + "MambaHybridCacheManager": ("common", "MambaHybridCacheManager"), + "MambaHybridCacheManagerV2": ( + "mamba_cache_manager_v2", + "MambaHybridCacheManagerV2", + ), + "MambaLayerCache": ("common", "MambaLayerCache"), + "MambaRole": ("common", "MambaRole"), + "MixedMambaHybridCacheManager": ("legacy", "MixedMambaHybridCacheManager"), + "PythonMambaCacheManager": ("legacy", "PythonMambaCacheManager"), + "_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_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", + "MambaCacheManager", + "MambaHybridCacheManager", + "MambaHybridCacheManagerV2", + "MambaLayerCache", + "MambaRole", + "MixedMambaHybridCacheManager", + "PythonMambaCacheManager", + "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..b88c7f3cfd53 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/common.py @@ -0,0 +1,730 @@ +# 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, 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.config_utils import mamba_effective_tp_size +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._utils import TensorWrapper, convert_to_torch_tensor +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 +_mamba_effective_tp_size = mamba_effective_tp_size + + +@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 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 + 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 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: ... + + +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]``.""" + + @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.""" + + +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. ``compute_dynamic_draft_len_mapping`` is created + later and covers only configured CUDA-graph batch sizes, so it cannot size + this persistent ID set. + """ + if spec_config is None: + return 1 + + draft_len_schedule = getattr(spec_config, "draft_len_schedule", None) + spec_dec_mode = getattr(spec_config, "spec_dec_mode", None) + supports_dynamic_draft_len = ( + spec_dec_mode is not None + and hasattr(spec_dec_mode, "support_dynamic_draft_len") + and spec_dec_mode.support_dynamic_draft_len() + ) + if draft_len_schedule and supports_dynamic_draft_len: + runtime_draft_lengths = set() + first_uncovered_batch_size = 1 + for batch_size_threshold, draft_len in draft_len_schedule.items(): + if first_uncovered_batch_size > max_batch_size: + break + if batch_size_threshold >= first_uncovered_batch_size: + runtime_draft_lengths.add(draft_len) + first_uncovered_batch_size = batch_size_threshold + 1 + if first_uncovered_batch_size <= max_batch_size: + runtime_draft_lengths.add(0) + else: + max_draft_len = getattr(spec_config, "max_draft_len", 0) or 0 + max_total_draft_tokens = getattr(spec_config, "max_total_draft_tokens", 0) or 0 + is_linear_tree = getattr( + spec_config, + "is_linear_tree", + max_draft_len == max_total_draft_tokens, + ) + static_draft_len = max_draft_len if is_linear_tree else max_total_draft_tokens + runtime_draft_lengths = {static_draft_len or 0} + + if (getattr(spec_config, "acceptance_rate_window_size", 0) or 0) > 0 and ( + getattr(spec_config, "acceptance_rate_threshold", 0) or 0 + ) > 0: + runtime_draft_lengths.add(0) + return len(runtime_draft_lengths) + + +class MambaRole: + """V2 buffer roles owned only by the hybrid Mamba manager.""" + + SSM_STATE = DataRole("ssm_state") + CONV_STATE = DataRole("conv_state") + + +def get_tensor_size_bytes(tensor): + """Calculate tensor size in bytes.""" + if isinstance(tensor, torch.Tensor): + 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 + + @property + def cuda_state_indices(self) -> torch.Tensor: + return self._state_indices_buffer.tensor + + @property + def _dummy_request_mask(self) -> Optional[torch.Tensor]: + if self._dummy_request_mask_buffer is None: + return None + return self._dummy_request_mask_buffer.tensor + + 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_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", + "MambaStateLayout", + "MambaAcceptanceBatch", + "MambaRole", + "use_py_mamba_cache_manager", + "get_tensor_size_bytes", +] + + +def _stack_state_views(states: List[torch.Tensor]) -> Optional[torch.Tensor]: + """Returns one ``[num_layers, num_slots, *state_shape]`` view of ``states``. + + A stacked view collapses per-layer promotion loop into a single launch of _promote_mamba_state_triton. + This requires per-layer pool views to be *affine*: identical dtype/device/shape/strides, and base + pointers spaced by a constant. returns ``None`` otherwise. + """ + if not states: + return None + ref = states[0] + # TensorWrapper exposes CUDA pointers; CPU contract tests use per-layer views. + if ref.device.type != "cuda": + return None + if len(states) == 1: + return ref.unsqueeze(0) + for state in states[1:]: + if ( + state.dtype != ref.dtype + or state.device != ref.device + or state.shape != ref.shape + or state.stride() != ref.stride() + ): + return None + itemsize = ref.element_size() + layer_stride_bytes = states[1].data_ptr() - states[0].data_ptr() + if layer_stride_bytes <= 0 or layer_stride_bytes % itemsize != 0: + return None + base = states[0].data_ptr() + if any( + state.data_ptr() != base + layer * layer_stride_bytes for layer, state in enumerate(states) + ): + return None + layer_stride = layer_stride_bytes // itemsize + # Highest element one layer's view can reach, so the flat wrapper below + # covers every byte the stacked view addresses and no more. + layer_extent = 1 + sum((size - 1) * stride for size, stride in zip(ref.shape, ref.stride())) + total_elems = (len(states) - 1) * layer_stride + layer_extent + flat = convert_to_torch_tensor(TensorWrapper(base, ref.dtype, [total_elems])) + return flat.as_strided([len(states)] + list(ref.shape), [layer_stride] + list(ref.stride())) + + +def _promote_intermediate_states( + destination_states: tuple[torch.Tensor, ...], + intermediate_states: torch.Tensor, + batch: MambaAcceptanceBatch, + stacked: torch.Tensor | None = None, +) -> 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 + + if stacked is not None: + mamba_cache_manager._promote_mamba_state_triton( + stacked, + intermediate_states, + batch.source_state_indices, + batch.accepted_positions, + batch.destination_state_indices, + ) + return + 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, + ) 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..e3bfeec169bb --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/legacy.py @@ -0,0 +1,2365 @@ +# 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 + +from tensorrt_llm._device_copy import DeviceCopyBuffer + +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.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, + 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 ( + BaseMambaCacheManager, + MambaHybridCacheManager, + _get_local_mamba_cache_layout, + _get_mamba_hybrid_pool_size, + _mamba_effective_tp_size, + _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 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", + ) + + +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. + + 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] bf16, + # 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: + # k and gate [layers, slots, num_spec, H*K], v [.., num_spec, H*V], + # beta [.., num_spec, H]. Gate is fp32; k, v, beta are bf16. + kda_k_cache: torch.Tensor | None = None + kda_g_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_k_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. + from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + allocate_kda_replay_fields, + ) + + replay_fields, ssm_spec_cache = allocate_kda_replay_fields( + num_local_layers=num_local_layers, + cache_size=max_batch_size, + num_speculative_tokens=self._kda_replay_num_spec, + conv_dim=conv_dim, + conv_kernel_size=d_conv, + num_heads=nheads, + device=device, + ) + spec_kwargs.update(replay_fields) + 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_k_cache, + self.mamba_cache.kda_g_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)) + mask_staging = torch.empty_like(self._dummy_request_mask_host, pin_memory=prefer_pinned()) + mask_staging.copy_(self._dummy_request_mask_host) + self._dummy_request_mask.copy_(mask_staging, 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_k_cache=_drop(self.mamba_cache.kda_k_cache), + kda_g_cache=_drop(self.mamba_cache.kda_g_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_buffer = 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_buffer.update(self._dummy_request_mask_host) + + 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.index_fill_(0, context_slots, 0) + self.cache_buf_idx.index_fill_(0, context_slots, 0) + if self.old_x is not None: + self.old_x.index_fill_(1, context_slots, 0) + if self.old_B is not None: + self.old_B.index_fill_(1, context_slots, 0) + if self.old_dt is not None: + self.old_dt.index_fill_(1, context_slots, 0) + if self.old_dA_cumsum is not None: + self.old_dA_cumsum.index_fill_(1, 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, + pin_memory=prefer_pinned(), + ).to(self.mamba_ssm_rand_seed.device, non_blocking=True) + + 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. + """ + + # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the + # cache manager __init__ coerces a positive spec recompute tail to a full + # re-prefill (see _has_recurrent_state on the attention-only base classes). + _has_recurrent_state = True + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + # kv cache 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 | None = 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__) + 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" + ) + + # 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. + + """ + + # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the + # cache manager __init__ coerces a positive spec recompute tail to a full + # re-prefill (see _has_recurrent_state on the attention-only base classes). + _has_recurrent_state = True + + def __init__( + self, + # mamba cache parameters + mamba_d_state: int, + mamba_d_conv: int, + mamba_num_heads: int, + mamba_n_groups: int, + mamba_head_dim: int, + mamba_num_layers: int, + mamba_layer_mask: List[bool], + mamba_cache_dtype: torch.dtype, + mamba_ssm_cache_dtype: torch.dtype, + kv_cache_config: KvCacheConfig, + kv_cache_type: CacheTypeCpp, + *, + num_layers: int, + num_kv_heads: Union[int, List[Optional[int]]], + head_dim: int, + tokens_per_block: int, + # 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 | None = 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__) + 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 + 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._state_indices_buffer = DeviceCopyBuffer( + torch.zeros([self.max_batch_size], dtype=torch.int32, device="cuda"), + pin_memory=prefer_pinned(), + ) + 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_buffer = 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, + 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, + 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._state_indices_buffer.update(self._host_state_indices, force=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_buffer = DeviceCopyBuffer( + torch.zeros(self.max_batch_size, dtype=torch.bool, device=device), + pin_memory=prefer_pinned(), + ) + 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..5b56920dc1d0 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager/mamba_cache_manager_v2.py @@ -0,0 +1,1354 @@ +# 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 + +from tensorrt_llm._device_copy import DeviceCopyBuffer + +if sys.version_info[:2] >= (3, 12): + from typing import override +else: + from typing_extensions import override + +from tensorrt_llm._torch.disaggregation.resource.page import MapperKind, RoleLayout +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( + _RESERVED_REQUEST_IDS, + 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.bindings.executor import KvCacheRetentionConfig +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 KVCacheManager as KVCacheManagerPy +from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy + +from .common import ( + MambaAcceptanceBatch, + MambaHybridCacheManager, + MambaLayerCache, + MambaRole, + MambaState, + MambaStateLayout, + _estimate_mamba_hybrid_cache_cost, + _get_num_cuda_graph_padding_dummy_slots, + _mamba_effective_tp_size, + _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 + # Recurrent snapshots use their own commit/history endpoints without draft backoff. + _supports_reuse_match_backoff = False + + # Recurrent (conv/SSM) state summarizes the whole matched prefix, so the + # cache manager __init__ coerces a positive spec recompute tail to a full + # re-prefill (see _has_recurrent_state on the attention-only base classes). + _has_recurrent_state = 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, + 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._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, + conv_section_dims=tuple(self.conv_section_dims), + conv_state_layout=conv_state_layout, + ) + self._speculative_state: MambaState | None = self._initialize_spec_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() + + @override + def _disagg_transfer_overwrites_whole_cached_prefix(self) -> bool: + # A transferred recurrent state summarizes and overwrites the whole prefix. + return ( + self.local_num_mamba_layers > 0 + or super()._disagg_transfer_overwrites_whole_cached_prefix() + ) + + 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( + "Speculative decoding requires a model-specific Mamba cache manager" + ) + return None + + @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 + + @override + def get_disagg_role_mapper_kinds(self) -> dict[DataRole, MapperKind]: + """Convolution sections shard independently; SSM state shards by head.""" + return { + **super().get_disagg_role_mapper_kinds(), + MambaRole.CONV_STATE: MapperKind.SECTIONED, + } + + @override + def get_disagg_role_layouts(self) -> dict[DataRole, RoleLayout]: + """Describe per-layer recurrent geometry for resharding across TP.""" + if self.local_num_mamba_layers == 0: + return {} + d_conv_m1 = int(self.conv_state_shape[1]) + _, head_dim, d_state = self.ssm_state_shape + return { + MambaRole.CONV_STATE: RoleLayout( + section_bytes=tuple( + int(dim) * d_conv_m1 * self.conv_state_dtype.itemsize + for dim in self.conv_section_dims + ) + ), + MambaRole.SSM_STATE: RoleLayout( + bytes_per_head=int(head_dim) * int(d_state) * self.ssm_state_dtype.itemsize + ), + } + + 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 [ + 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 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( + 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._state_indices_buffer = DeviceCopyBuffer( + torch.zeros([state_index_capacity], dtype=torch.int32, device="cuda"), + pin_memory=prefer_pinned(), + ) + self._host_state_indices = torch.zeros( + [state_index_capacity], dtype=torch.int32, pin_memory=prefer_pinned() + ) + + self._dummy_request_mask_buffer = DeviceCopyBuffer( + torch.zeros(state_index_capacity, dtype=torch.bool, device="cuda"), + pin_memory=prefer_pinned(), + ) + 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.""" + 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 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) + ) + + 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: + reserved = set(self.kv_cache_map) & _RESERVED_REQUEST_IDS + assert not set(self.kv_cache_map) - reserved, ( + "get_num_free_blocks is only used when the kv cache manager is empty" + ) + 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, + kv_cache_retention_config: Optional[KvCacheRetentionConfig] = 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, + kv_cache_retention_config=kv_cache_retention_config, + ) + 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, + 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, + 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._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._state_indices_buffer.update(self._host_state_indices) + 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 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 _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_buffer.update(self._dummy_request_mask_host) + + @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] + 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, + 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, + ), + ) + + @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[tuple[int, 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 global_layer_id, 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}" + ) + # Recurrent roles have no attention guard page. + yield -1, buffer + + @override + def shutdown(self): + if self._speculative_state is not None: + self._speculative_state.shutdown() + self._recurrent_buffers.clear() + self.all_ssm_states = [] + self.all_conv_states = [] + self._dummy_request_mask_buffer = None + self._dummy_request_mask_host = None + self._generation_state_indices = None + self._branch_snapshot_points.clear() + super().shutdown() + + +__all__ = ["MambaHybridCacheManagerV2"] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index bf552ac11533..3837bceb2110 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -8074,21 +8074,11 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): self.resource_manager.resource_managers[ 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 ]) # Cache state that the generation worker derives from the imported KV diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 341d16689a59..813903b95aa3 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4132,7 +4132,7 @@ }, "role": { "default": "both", - "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without moving any KV, so they pin no staging buffers and pay no per-request cost.", + "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without registering any KV cache, so it needs no GPUDirect RDMA.", "enum": [ "both", "producer", @@ -4172,8 +4172,8 @@ "title": "Segment Size" }, "stage_through_host": { - "default": true, - "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. The only transfer path supported today, so 'false' is ignored with a warning. Costs a copy each way and needs no GPUDirect RDMA.", + "default": false, + "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'.", "telemetry": { "exclude": true }, @@ -6948,6 +6948,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index b57ddfb733f8..6f9079a2da86 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -2314,6 +2314,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { @@ -5682,7 +5693,7 @@ }, "role": { "default": "both", - "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without moving any KV, so they pin no staging buffers and pay no per-request cost.", + "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without registering any KV cache, so it needs no GPUDirect RDMA.", "enum": [ "both", "producer", @@ -5722,8 +5733,8 @@ "title": "Segment Size" }, "stage_through_host": { - "default": true, - "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. The only transfer path supported today, so 'false' is ignored with a warning. Costs a copy each way and needs no GPUDirect RDMA.", + "default": false, + "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'.", "telemetry": { "exclude": true }, @@ -8615,6 +8626,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 0a71b3fb1b4b..7793842f6c30 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6688,8 +6688,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) @@ -6845,8 +6847,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/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 diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d95c426ee9e3..0ce83671f8c1 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -252,11 +252,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 diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 40d0f7f057e2..735b8072857c 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 --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.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 --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.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/msa/test_msa_proxy_split.py @@ -45,6 +45,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/moe/kernels/parallel - unittest/_torch/kernels/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel diff --git a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py index e989db333636..90939aa14927 100644 --- a/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py +++ b/tests/unittest/_torch/attention/sparse/glm_kpool/test_glm_kpool.py @@ -153,6 +153,8 @@ def test_indexer_runtime_cache_cost_matches_registered_buffers(fp8, snapshot_int manager.index_state_dim = 384 manager.num_local_layers = 2 manager.local_num_mamba_layers = 1 + manager.mamba_pp_layers = [0] + manager._kda_replay = None manager.num_kv_heads_per_layer = [0, 1] manager.head_dim_per_layer = [512, 512] manager.kv_cache_type = CacheTypeCpp.SELFKONLY 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 b766051d0ef0..750a406c84da 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, ) @@ -37,6 +36,7 @@ select_qsa_paged_tokens, select_qsa_tokens, ) +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import Qwen4ExpHybridCacheManager from tensorrt_llm.runtime.kv_cache_manager_v2 import PageIndexMode @@ -1451,7 +1451,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__(Qwen4ExpHybridCacheManager) + manager._qsa_enabled = True manager.qsa_index_dim = 128 manager.qsa_index_kv_heads = 1 manager.qsa_sparse_layer_ids = [7, 11] @@ -1476,13 +1477,13 @@ def get_page_stride(self, layer_idx, role): del layer_idx, role return 0 - manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) 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_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 255bc46ff80c..717424cad2e3 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 @@ -1073,6 +1073,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 @@ -1091,6 +1092,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 6f0e3f6b788f..3df75f0e07a2 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,9 @@ import torch from tensorrt_llm._device_copy import DeviceCopyBuffer +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, @@ -20,6 +23,23 @@ MambaLayerGroup, ) from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 +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, + KDAIntermediateState, + KDAReplayLayerCache, + KDAReplayState, + KimiK3HybridCacheManager, + _kda_replay_layout, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + MIN_REPLAY_HISTORY_SIZE, + Mamba2State, + NemotronHybridCacheManager, + ReplayHistory, + ReplayStateUpdateMetadata, + _advance_replay_state, +) from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata from tensorrt_llm._torch.pyexecutor._util import ( KvCacheCreator, @@ -36,23 +56,19 @@ KVCacheManagerV2, ) from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import ( - _KDA_BETA_CACHE_ALIGNMENT_BYTES, - 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, - _kda_replay_layout, _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.llm_request import ( ATTENTION_DP_DUMMY_REQUEST_ID, LlmRequest, @@ -158,10 +174,10 @@ def _hybrid_model_config(): ) return SimpleNamespace( pretrained_config=config, - mapping=Mapping(), sparse_attention_config=None, spec_config=None, get_num_mamba_layers=lambda: 1, + mapping=Mapping(), ) @@ -324,7 +340,7 @@ def _capture_kimi_v2_manager_ctor( V2 manager and capture the constructor arguments.""" captured: dict[str, object] = {} - class RecordingV2Manager(MambaHybridCacheManagerV2): + class RecordingV2Manager(KimiK3HybridCacheManager): def __init__(self, *args: object, **kwargs: object) -> None: captured["args"] = args captured["kwargs"] = kwargs @@ -334,7 +350,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 KimiK3HybridCacheManager _create_kv_cache_manager( model_engine=model_engine, @@ -380,7 +396,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) @@ -391,35 +407,19 @@ 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__(KimiK3HybridCacheManager) + mgr.spec_config = spec_config + mgr._mamba_ssm_stochastic_rounding = False + mgr._requested_num_spec = None + mgr._requested_draft_records = False + mgr._state_layout = SimpleNamespace(spec_config=spec_config, mamba_pp_layers=()) + 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 assert kwargs["conv_state_layout"] == "q_k_v" - assert "kda_draft_records" not in kwargs - - -@pytest.mark.parametrize("backbone_asks", [True, False]) -def test_kimi_explicit_v2_manager_kda_draft_records_follow_the_model( - monkeypatch: pytest.MonkeyPatch, - backbone_asks: bool, -) -> None: - """The factory asks the V2 manager for per-token KDA verify states only when - the model's backbone sets ``kda_draft_records``.""" - monkeypatch.setattr( - "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", - lambda: True, - ) - backbone = SimpleNamespace(kda_draft_records=True) if backbone_asks else SimpleNamespace() - model_engine = SimpleNamespace( - model=SimpleNamespace(model=backbone), - _max_cuda_graph_batch_size=4, - is_multimodal=True, - ) - - _, kwargs = _capture_kimi_v2_manager_ctor( - monkeypatch, MTPDecodingConfig(max_draft_len=3), model_engine - ) - - assert kwargs["kda_replay_num_spec"] == 3 - assert kwargs.get("kda_draft_records", False) is backbone_asks + assert "kda_replay_num_spec" not in kwargs @pytest.mark.parametrize("max_draft_len, replay_enabled", [(7, True), (8, False)]) @@ -436,10 +436,17 @@ def test_kimi_kda_replay_respects_fused_verify_draft_limit( _, kwargs = _capture_kimi_v2_manager_ctor(monkeypatch, spec_config) + assert "kda_replay_num_spec" not in kwargs + mgr = object.__new__(KimiK3HybridCacheManager) + mgr.spec_config = spec_config + mgr._mamba_ssm_stochastic_rounding = False + mgr._requested_num_spec = None + mgr._requested_draft_records = False + mgr._state_layout = SimpleNamespace(spec_config=spec_config, mamba_pp_layers=()) + state = mgr._initialize_spec_state() + assert isinstance(state, KDAReplayState) is replay_enabled if replay_enabled: - assert kwargs["kda_replay_num_spec"] == max_draft_len - else: - assert "kda_replay_num_spec" not in kwargs + assert state.num_speculative_tokens == max_draft_len def test_kda_mtp_max_num_spec_fits_the_kernel_chunk() -> None: @@ -466,12 +473,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): @@ -502,7 +507,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: ) kwargs = captured["kwargs"] assert "vocab_size" not in 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 @@ -538,8 +543,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, Qwen35HybridCacheManager), + (True, True, Qwen35HybridCacheManager), (False, True, CppMambaHybridCacheManager), ("auto", True, CppMambaHybridCacheManager), ], @@ -603,7 +608,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, @@ -646,19 +651,43 @@ 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 "vocab_size" not in captured_cpp - assert captured_mixed["use_replay_state_update"] is False + assert "vocab_size" not in captured_mixed + assert captured_mixed["use_replay_state_update"] is None assert captured_mixed["model_type"] == "qwen3_next" assert captured_mixed["max_num_tokens"] == 256 - assert "vocab_size" not in captured_mixed - 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__(Qwen35HybridCacheManager) + 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._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 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) @@ -701,7 +730,6 @@ def __init__(self, *args, **kwargs): dtype=torch.bfloat16, mamba_ssm_cache_dtype=torch.bfloat16, ) - monkeypatch.setattr("tensorrt_llm._torch.pyexecutor._util.get_sm_version", lambda: 90) monkeypatch.setattr( "tensorrt_llm._torch.pyexecutor._util.extract_mamba_kv_cache_params", lambda *args, **kwargs: mamba_params, @@ -822,7 +850,7 @@ def test_hybrid_cache_manager_factory_routes_explicit_v2_disagg(monkeypatch, bac backend=backend, transceiver_runtime="PYTHON" ), ) - is MambaHybridCacheManagerV2 + is Qwen35HybridCacheManager ) @@ -1020,7 +1048,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 = Qwen35HybridCacheManager if expected_v2 else CppMambaHybridCacheManager assert ( get_kv_cache_manager_cls( _hybrid_model_config(), @@ -1121,7 +1149,7 @@ def test_kimi_disagg_defaults_use_v2_manager(monkeypatch: pytest.MonkeyPatch) -> is_disagg=True, cache_transceiver_config=llm_args.cache_transceiver_config, ) - is MambaHybridCacheManagerV2 + is KimiK3HybridCacheManager ) @@ -1517,16 +1545,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_buffer = DeviceCopyBuffer(torch.zeros(64, dtype=torch.bool)) mgr.update_mamba_states( SimpleNamespace(num_seqs=2, num_contexts=0), torch.tensor([3, 2], dtype=torch.int32), @@ -1902,7 +1927,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, @@ -1958,6 +1983,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 @@ -1974,6 +2006,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 @@ -2038,7 +2071,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), ) @@ -2058,12 +2091,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 @@ -2089,12 +2124,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 @@ -2328,6 +2365,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 @@ -2456,17 +2494,21 @@ def _build_v2_hybrid_with_mamba_layer( dtype=DataType.HALF, kv_cache_dtype="auto", conv_state_layout="x_b_c", + speculative_state=None, mamba_d_conv=4, mamba_num_heads=4, mamba_n_groups=1, mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, enable_kv_pool_rebalance=False, + max_tokens=512, + max_gpu_total_bytes=0, kda_draft_records=False, cold_tier_offload_min_priority=None, host_cache_size=None, is_disagg=False, enable_stats=False, + model_manager_cls=None, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2480,8 +2522,8 @@ def _build_v2_hybrid_with_mamba_layer( enable_attention_dp=enable_attention_dp, ) kv_cache_config = KvCacheConfig( - max_tokens=512, - host_cache_size=host_cache_size, + 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( @@ -2497,8 +2539,39 @@ def _build_v2_hybrid_with_mamba_layer( dtype=kv_cache_dtype, enable_kv_pool_rebalance=enable_kv_pool_rebalance, cold_tier_offload_min_priority=cold_tier_offload_min_priority, + host_cache_size=host_cache_size, ) - return MambaHybridCacheManagerV2( + if speculative_state is None and kda_replay_num_spec is not None: + speculative_state = KDAReplayState(kda_replay_num_spec, draft_records=kda_draft_records) + model_cls = ( + KimiK3HybridCacheManager + if isinstance(speculative_state, KDAReplayState) or kda_draft_records + else Qwen35HybridCacheManager + if conv_state_layout == "q_k_v" + else NemotronHybridCacheManager + ) + if model_manager_cls is not None: + model_cls = model_manager_cls + + class ModelStateManager(model_cls): + def _initialize_spec_state(self): + if getattr(self, "kda_context_only", False): + return None + if kda_draft_records and speculative_state is None: + return KDAIntermediateState() + if speculative_state is None: + return super()._initialize_spec_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 + + validate_gdn_layout(self, speculative_state) + return speculative_state + + manager_cls = ModelStateManager + return manager_cls( mamba_d_state=8, mamba_d_conv=mamba_d_conv, mamba_num_heads=mamba_num_heads, @@ -2523,10 +2596,15 @@ 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, - kda_draft_records=kda_draft_records, is_disagg=is_disagg, enable_stats=enable_stats, + **( + {"kda_draft_records": kda_draft_records} + if issubclass(model_cls, KimiK3HybridCacheManager) + else {"index_state_dim": 16} + if issubclass(model_cls, Glm5NextCacheManager) + else {} + ), ) @@ -2565,6 +2643,8 @@ def _make_v2_conversation_request( def _run_v2_hybrid_context( manager: MambaHybridCacheManagerV2, request: LlmRequest, + *, + state_value: int | None = None, ) -> None: manager.prepare_expect_snapshot_points([request]) assert manager.prepare_context(request) @@ -2574,6 +2654,10 @@ def _run_v2_hybrid_context( batch = ScheduledRequests() batch.append_context_request(request) manager.prepare_resources(batch) + if state_value is not None: + state_index = manager.get_state_indices([request.py_request_id], [False])[0] + for buffer in (*manager.all_ssm_states, *manager.all_conv_states): + buffer[state_index].fill_(state_value) request.context_current_position = request.prompt_len assert request.context_remaining_length == 0 manager.update_context_resources(batch) @@ -2855,7 +2939,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) @@ -2923,7 +3008,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) @@ -3025,8 +3111,16 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): request_b_after_eviction, ] + def assert_reused_state(request: LlmRequest, value: int) -> None: + state_index = mgr.get_state_indices([request.py_request_id], [False])[0] + # Reuse preserves payloads even when cache offload/onboard moves their slots. + for buffer in (*mgr.all_ssm_states, *mgr.all_conv_states): + torch.testing.assert_close( + buffer[state_index], torch.full_like(buffer[state_index], value), rtol=0, atol=0 + ) + try: - _run_v2_hybrid_context(mgr, request_a) + _run_v2_hybrid_context(mgr, request_a, state_value=1) cache_a = mgr.kv_cache_map[request_a.py_request_id] attention_group = mgr.impl.get_layer_group_id(1) priorities = cache_a.get_page_priorities(attention_group) @@ -3034,10 +3128,8 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): assert mgr.impl.init_config.cold_tier_offload_min_priority == 50 assert not mgr.can_evict assert mgr.prefetch_for_context_tokens([request_a]) is True - request_a_state_index = mgr.get_state_indices([request_a.py_request_id], [False])[0] mgr.free_resources(request_a) - _run_v2_hybrid_context(mgr, request_b) - request_b_state_index = mgr.get_state_indices([request_b.py_request_id], [False])[0] + _run_v2_hybrid_context(mgr, request_b, state_value=2) mgr.free_resources(request_b) mgr.prepare_expect_snapshot_points([request_a_probe]) @@ -3046,9 +3138,7 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): probe_batch.append_context_request(request_a_probe) mgr.prepare_resources(probe_batch) assert request_a_probe.prepopulated_prompt_len == request_a_probe.prompt_len - 1 - assert mgr.get_state_indices([request_a_probe.py_request_id], [False]) == [ - request_a_state_index - ] + assert_reused_state(request_a_probe, 1) mgr.free_resources(request_a_probe) mgr.prepare_expect_snapshot_points([request_b_probe]) @@ -3057,12 +3147,10 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): probe_batch.append_context_request(request_b_probe) mgr.prepare_resources(probe_batch) assert request_b_probe.prepopulated_prompt_len == request_b_probe.prompt_len - 1 - assert mgr.get_state_indices([request_b_probe.py_request_id], [False]) == [ - request_b_state_index - ] + assert_reused_state(request_b_probe, 2) mgr.free_resources(request_b_probe) - _run_v2_hybrid_context(mgr, request_c) + _run_v2_hybrid_context(mgr, request_c, state_value=3) mgr.free_resources(request_c) mgr.prepare_expect_snapshot_points([request_a_after_eviction]) @@ -3078,9 +3166,7 @@ def test_v2_hybrid_retains_configured_number_of_conversation_turns(): request_b_after_eviction.prepopulated_prompt_len == request_b_after_eviction.prompt_len - 1 ) - assert mgr.get_state_indices([request_b_after_eviction.py_request_id], [False]) == [ - request_b_state_index - ] + assert_reused_state(request_b_after_eviction, 2) finally: for request in requests: if request.py_request_id in mgr.kv_cache_map: @@ -3719,8 +3805,12 @@ def test_v2_hybrid_intermediate_states_size_by_tokens_per_gen_step(): @skip_no_cuda -def test_v2_kda_replay_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, + ) mgr = _build_v2_hybrid_with_mamba_layer( max_batch_size=4, num_mamba_layers=2, @@ -3730,15 +3820,24 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): mamba_num_heads=6, mamba_n_groups=6, mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=2, + speculative_state=policy, + max_tokens=max_tokens, + max_gpu_total_bytes=byte_limit, ) 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 not hasattr(mgr, "get_replay_state_update_metadata") + + native_bytes = mgr.impl.get_quota(GPU_LEVEL) + assert mgr._get_external_cache_bytes(mgr.impl) == 0 + if byte_limit: + assert native_bytes <= byte_limit + 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") 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) @@ -3766,118 +3865,78 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): assert layer_cache.kda_g_cache.dtype is torch.float32 assert layer_cache.kda_v_cache.dtype is torch.bfloat16 assert layer_cache.kda_beta_cache.dtype is torch.bfloat16 - assert layer_cache.prev_num_accepted_tokens.dtype == torch.int32 - assert layer_cache.prev_num_accepted_tokens.shape == (cache_size,) - assert layer_cache.prev_num_accepted_tokens is ( - mgr.mamba_layer_cache(1).prev_num_accepted_tokens + 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 layer_cache.kda_state_tok is None - assert not mgr.keeps_kda_draft_records - assert layer_cache.has_kda_replay_caches - finally: - mgr.shutdown() - - -@skip_no_cuda -@pytest.mark.parametrize("kda_replay_num_spec", [2, None]) -def test_v2_kda_draft_records_replace_the_replay_caches(kda_replay_num_spec): - """On the KDA replay path, ``kda_draft_records`` allocates the fp32 records of every draft of each slot (vn, - beta * k, decay), per layer, in place of the per-draft replay caches; the conv caches and the replay path stay. - Without the replay path it allocates nothing.""" - mgr = _build_v2_hybrid_with_mamba_layer( - max_batch_size=4, - num_mamba_layers=2, - spec_config=MTPDecodingConfig(max_draft_len=2), - conv_state_layout="q_k_v", - mamba_d_conv=5, - mamba_num_heads=6, - mamba_n_groups=6, - mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=kda_replay_num_spec, - kda_draft_records=True, - ) - try: - if kda_replay_num_spec is None: - assert getattr(mgr, "kda_state_tok", None) is None - assert not mgr.keeps_kda_draft_records - return - assert mgr.keeps_kda_draft_records - assert set(mgr._kda_replay_layout()) == { - "kda_conv_q", - "kda_conv_k", - "kda_conv_v", - "kda_state_tok", - } - # A KVCM state buffer of the SSM layers, kept out of disaggregated transfers like the other replay buffers. - assert DataRole("kda_state_tok") in mgr.get_disagg_ignored_roles() - assert len(mgr.kda_state_tok) == 2 - for layer_idx in range(2): - layer_cache = mgr.mamba_layer_cache(layer_idx) - cache_size = layer_cache.temporal.shape[0] - assert layer_cache.has_kda_replay_caches - assert layer_cache.kda_conv_q.shape == (cache_size, 48, 6) - assert layer_cache.kda_k_cache is None - assert layer_cache.kda_g_cache is None - assert layer_cache.kda_v_cache is None - assert layer_cache.kda_beta_cache is None - states = layer_cache.kda_state_tok - num_heads, _, head_dim = layer_cache.temporal.shape[1:] - assert states.shape == (cache_size, 3, 2, num_heads, head_dim) - assert states.dtype is torch.float32 - assert states.data_ptr() == mgr.kda_state_tok[layer_idx].data_ptr() - # The records' replay addresses a slot by its stride and reads its records contiguously. - assert states[0].is_contiguous() + assert not hasattr(layer_cache, "intermediate_ssm") + assert not hasattr(layer_cache, "intermediate_conv_window") finally: mgr.shutdown() - # Released with the other replay buffers. - assert mgr.kda_state_tok is None @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, ) +@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() + + @skip_no_cuda @pytest.mark.parametrize("kda_replay", [True, False], ids=["kda_replay", "no_replay"]) def test_v2_kda_replay_refuses_kv_pool_rebalance(kda_replay): @@ -3901,7 +3960,7 @@ def test_v2_kda_replay_refuses_kv_pool_rebalance(kda_replay): def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): """A V2 hybrid manager stub for _build_cache_config: one KDA layer with the replay caches and one attention layer.""" - mgr = object.__new__(MambaHybridCacheManagerV2) + mgr = object.__new__(KimiK3HybridCacheManager) mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False mgr.kv_cache_type = CacheTypeCpp.SELF @@ -3928,14 +3987,18 @@ def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): mgr.get_layer_bytes_per_token = lambda **kwargs: 8 # Layer 1's 256-byte pages of 32 tokens (_base_attention_layer_configs; layer 0 becomes the SSM layer). mgr._attention_cache_bytes_per_token = lambda: 8 - mgr._use_kda_replay_update = True - mgr._kda_draft_records = kda_draft_records + mgr._kda_replay = KDAReplayState(2, draft_records=kda_draft_records) # The replay buffers count in each slot's state bytes (_mamba_state_bytes_per_slot); the drafts' records are, # per draft, vn, beta * k and the decay, fp32 [H, K] each (24 KiB here). - mgr._kda_replay_num_spec = 2 mgr.ssm_state_shape = [8, 256, 256] mgr.conv_section_dims = [8 * 256] * 3 mgr.conv_state_shape = [3 * 8 * 256, 3] + mgr._state_layout = SimpleNamespace( + mamba_pp_layers=tuple(mgr.mamba_pp_layers), + ssm_state_shape=tuple(mgr.ssm_state_shape), + conv_state_shape=tuple(mgr.conv_state_shape), + conv_section_dims=tuple(mgr.conv_section_dims), + ) mgr.kv_cache_config = KvCacheConfig( avg_seq_len=64, enable_block_reuse=enable_block_reuse, @@ -3944,69 +4007,6 @@ def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): return mgr -def _kda_draft_record_build_config(quota, initial_pool_ratio=None): - return KVCacheManagerConfig( - tokens_per_block=32, - cache_tiers=[GpuCacheTierConfig(quota=quota)], - layers=_base_attention_layer_configs(2), - initial_pool_ratio=initial_pool_ratio, - ) - - -@pytest.mark.parametrize("pool_sizing", ["typical_step", "block_reuse", "pool_ratio"]) -def test_v2_kda_draft_records_size_with_the_ssm_pool(pool_sizing): - """The drafts' records are a buffer of the SSM layer in each recurrent-state slot, in place of the k / g / v / - beta caches. However the SSM pool is sized (by the typical step, with block reuse, or by a pool_ratio), every slot - holds them: the build leaves the whole GPU quota to the pools, and the quotas count the records per state slot.""" - enable_block_reuse = pool_sizing == "block_reuse" - mgr = _v2_kda_draft_record_manager( - kda_draft_records=True, enable_block_reuse=enable_block_reuse - ) - plain = _v2_kda_draft_record_manager( - kda_draft_records=False, enable_block_reuse=enable_block_reuse - ) - records = 3 * 2 * 8 * 256 * 4 - replay_caches = sum( - dtype.itemsize * math.prod(shape) - for name, (dtype, shape) in _kda_replay_layout(2, 8, 256, 3).items() - if name in _KDA_REPLAY_BUFFERS[3:] - ) - slot_delta = records - replay_caches - assert mgr._mamba_state_bytes_per_slot() - plain._mamba_state_bytes_per_slot() == slot_delta - live_slots = mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots - assert ( - mgr._minimum_live_gpu_quota() - plain._minimum_live_gpu_quota() == live_slots * slot_delta - ) - - quota = 128 << 20 - pool_ratio = [0.5, 0.5] if pool_sizing == "pool_ratio" else None - built = mgr._build_cache_config(_kda_draft_record_build_config(quota, pool_ratio)) - assert built.cache_tiers[0].quota == quota - ssm_buffers = {str(buffer.role): buffer.size for buffer in built.layers[0].buffers} - assert ssm_buffers["kda_state_tok"] == records - assert not set(ssm_buffers) & set(_KDA_REPLAY_BUFFERS[3:]) - with pytest.raises(ValueError, match="too small for live recurrent states"): - mgr._build_cache_config(_kda_draft_record_build_config(mgr._minimum_live_gpu_quota() - 1)) - - -@skip_no_cuda -@pytest.mark.parametrize("enable_block_reuse", [False, True]) -def test_v2_kda_draft_records_leave_the_quota_to_the_pools(enable_block_reuse): - """With the drafts' records inside the SSM pool, the runtime pools take the whole GPU quota (rounded up to whole - 2 MiB grains) with or without block reuse: nothing is set aside beside them.""" - mgr = _v2_kda_draft_record_manager( - kda_draft_records=True, enable_block_reuse=enable_block_reuse - ) - runtime_manager = RuntimeKVCacheManager( - mgr._build_cache_config(_kda_draft_record_build_config(128 << 20)) - ) - try: - pool_bytes = runtime_manager.get_quota(GPU_LEVEL) - finally: - runtime_manager.shutdown() - assert 128 << 20 <= pool_bytes < (128 << 20) + (2 << 20) - - def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) @@ -4014,45 +4014,67 @@ def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: assert mgr.use_kda_replay_update -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__(KimiK3HybridCacheManager) 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_buffer = DeviceCopyBuffer(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"): + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(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__(KimiK3HybridCacheManager) + policy = KDAReplayState( + num_speculative_tokens=2, + ) + policy.prev_num_accepted_tokens = torch.tensor([5, 6, 7], dtype=torch.int32) + mgr._kda_replay = policy + mgr._speculative_state = policy mgr._state_indices_buffer = DeviceCopyBuffer(torch.tensor([2, 1], dtype=torch.int32)) mgr._host_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 = 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__(KimiK3HybridCacheManager) 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._kda_replay_acceptance_pending = {101, 202, 404} - 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 + mgr._speculative_state = policy + policy.acceptance_pending = {101, 202, 404} base_update = MagicMock() monkeypatch.setattr(KVCacheManagerV2, "update_resources", base_update) active = SimpleNamespace( @@ -4075,9 +4097,9 @@ 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] # The batch's acceptances are recorded; another batch's are still pending. - assert mgr._kda_replay_acceptance_pending == {404} + assert policy.acceptance_pending == {404} _KDA_REPLAY_BUFFERS = ( @@ -4105,57 +4127,76 @@ 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, draft_records=False): """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__(MambaHybridCacheManagerV2) + mgr = object.__new__(manager_cls) mgr.local_num_mamba_layers = 1 - mgr._use_kda_replay_update = True - mgr._use_replay_state_update = False - mgr._record_kda_replay_in_update_resources = False + mgr._kda_replay = KDAReplayState(2, draft_records=draft_records) + mgr._speculative_state = mgr._kda_replay + mgr.spec_config = SimpleNamespace(decoding_type="MTP") + mgr._stream = None mgr._request_id_to_state_index = {} mgr._request_id_to_is_dummy = {} - mgr._kda_replay_suspended = set() - mgr._kda_replay_acceptance_pending = set() - mgr._kda_replay_unseeded = set() + mgr._kda_replay.suspended_requests = set() + mgr._kda_replay.acceptance_pending = set() + mgr._kda_replay.unseeded_requests = set() mgr._host_state_indices = torch.zeros(num_slots, dtype=torch.int32) mgr._state_indices_buffer = DeviceCopyBuffer(torch.zeros(num_slots, dtype=torch.int32)) - mgr._dummy_request_mask_buffer = None - mgr.mamba_ssm_rand_seed = None + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(num_slots, dtype=torch.bool)) + mgr._dummy_request_mask_host = torch.zeros(num_slots, dtype=torch.bool) mgr.kv_cache_map = {} mgr.ssm_layer_group_id = 0 - mgr.prev_num_accepted_tokens = torch.zeros(num_slots, dtype=torch.int32) - mgr._kda_replay_layers = [{name: torch.zeros(num_slots, 2, 4) for name in _KDA_REPLAY_BUFFERS}] + mgr._kda_replay.prev_num_accepted_tokens = torch.zeros(num_slots, dtype=torch.int32) + for name in _KDA_REPLAY_BUFFERS: + setattr(mgr._kda_replay, name, torch.zeros(1, num_slots, 2, 4)) + if draft_records: + for name in _KDA_REPLAY_BUFFERS[3:]: + setattr(mgr._kda_replay, name, None) + mgr._kda_replay.kda_state_tok = torch.zeros(1, num_slots, 3, 2, 1, 2) mgr.conv_state_shape = [6, 3] mgr.conv_section_dims = [2, 2, 2] mgr.all_conv_states = [torch.zeros(num_slots, 6, 3)] + mgr._kda_replay._conv_states = tuple(mgr.all_conv_states) + mgr._kda_replay._committed_window = 3 + mgr._kda_replay._conv_section_dims = (2, 2, 2) mgr.committed_slots = [] - mgr._commit_kda_pending_drafts = mgr.committed_slots.append - mgr._on_kv_cache_stream = contextlib.nullcontext + mgr._kda_replay.commit_pending_drafts = mgr.committed_slots.append + mgr._kda_replay.on_cache_stream = lambda stream: contextlib.nullcontext() return mgr def _write_kda_replay_slot(mgr, slot, accepted, value): """What a verify step leaves in its request's slot for the next one.""" - mgr.prev_num_accepted_tokens[slot] = accepted - for name in _KDA_REPLAY_BUFFERS: - mgr._kda_replay_layers[0][name][slot] = value + mgr._kda_replay.prev_num_accepted_tokens[slot] = accepted + for buffer in mgr._kda_replay._layer_replay_buffers(): + buffer[:, slot] = value def _kda_replay_slot(mgr, slot): - return int(mgr.prev_num_accepted_tokens[slot]), [ - float(mgr._kda_replay_layers[0][name][slot].flatten()[0]) for name in _KDA_REPLAY_BUFFERS + return int(mgr._kda_replay.prev_num_accepted_tokens[slot]), [ + float(buffer[0, slot].flatten()[0]) for buffer in mgr._kda_replay._layer_replay_buffers() ] @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, draft_records", + [ + (KimiK3HybridCacheManager, False), + (KimiK3HybridCacheManager, True), + (Glm5NextCacheManager, False), + ], +) +def test_v2_kda_replay_suspended_request_resumes_from_its_v2_states( + resume_slot, manager_cls, draft_records +): """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, draft_records=draft_records) 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) @@ -4184,29 +4225,35 @@ def test_v2_kda_replay_suspended_request_resumes_from_its_v2_states(resume_slot) mgr.kv_cache_map[1].is_active = True mgr._setup_state_indices(resumed, num_contexts=0) - assert int(mgr.prev_num_accepted_tokens[resume_slot]) == 0 - replay = mgr._kda_replay_layers[0] + assert int(mgr._kda_replay.prev_num_accepted_tokens[resume_slot]) == 0 for section, name in enumerate(("kda_conv_q", "kda_conv_k", "kda_conv_v")): - cache = replay[name][resume_slot] + cache = getattr(mgr._kda_replay, name)[0, resume_slot] assert torch.equal(cache[:, :3], window[2 * section : 2 * section + 2]) assert not cache[:, 3:].any() for name in ("kda_k_cache", "kda_g_cache", "kda_v_cache", "kda_beta_cache"): - assert not replay[name][resume_slot].any() + buffer = getattr(mgr._kda_replay, name) + if buffer is not None: + assert not buffer[0, resume_slot].any() + if draft_records: + # A resumed slot has no pending drafts, so its old records remain unread. + expected_records = 7.0 if resume_slot == 0 else 0.0 + assert torch.all(mgr._kda_replay.kda_state_tok[0, resume_slot] == expected_records) if resume_slot != 0: - assert _kda_replay_slot(mgr, 0) == (3, [7.0] * len(_KDA_REPLAY_BUFFERS)) - assert not mgr._kda_replay_suspended + assert _kda_replay_slot(mgr, 0) == (3, [7.0] * len(mgr._kda_replay._layer_replay_buffers())) + assert not mgr._kda_replay.suspended_requests @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._record_kda_replay_in_update_resources = True + 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()) request = SimpleNamespace( @@ -4223,7 +4270,7 @@ def test_v2_kda_replay_refuses_a_suspend_before_the_host_drafter_acceptance( if recorded_by == "update_resources": mgr.update_resources(batch) - assert int(mgr.prev_num_accepted_tokens[0]) == 2 + assert int(mgr._kda_replay.prev_num_accepted_tokens[0]) == 2 else: mgr.revert_allocate_generation(request) mgr.suspend_request(request) @@ -4232,10 +4279,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) @@ -4258,9 +4306,9 @@ def test_v2_kda_replay_disagg_generation_request_suspends_its_transferred_state( mgr._setup_state_indices([request], num_contexts=0) for section, name in enumerate(("kda_conv_q", "kda_conv_k", "kda_conv_v")): assert torch.equal( - mgr._kda_replay_layers[0][name][2][:, :3], window[2 * section : 2 * section + 2] + getattr(mgr._kda_replay, name)[0, 2][:, :3], window[2 * section : 2 * section + 2] ) - assert not mgr._kda_replay_unseeded + assert not mgr._kda_replay.unseeded_requests mgr.suspend_request(request) assert mgr.committed_slots == ([0, 2] if seeded else [2]) @@ -4271,7 +4319,7 @@ def _v2_kda_replay_manager_with_large_slots(): SSM slot (~2.25 MiB) is larger than the 2 MiB allocation grain, so the SSM pool holds a handful of slots and runs out before the attention pages do. max_util_for_resume 1.0 lets resume admit caches until the pool is full.""" num_mamba_layers = 6 - return MambaHybridCacheManagerV2( + return KimiK3HybridCacheManager( mamba_d_state=128, mamba_d_conv=4, mamba_num_heads=6, @@ -4305,7 +4353,7 @@ def _take_ssm_slot(mgr, slot, others): """Create caches until the V2 runtime gives one of them ``slot``, moving the suspended owner's SSM page to the host tier; keep that one and close the rest.""" holder = None - for _ in range(mgr.prev_num_accepted_tokens.shape[0] + 1): + for _ in range(mgr._kda_replay.prev_num_accepted_tokens.shape[0] + 1): cache = mgr.impl.create_kv_cache() assert cache.resume(mgr._stream.cuda_stream) others.append(cache) @@ -4328,22 +4376,22 @@ def _fill_kda_slot(mgr, slot, generator): ssm_states[slot] = torch.randn(ssm_states.shape[1:], generator=generator, device="cuda") conv_states[slot] = torch.randn(conv_states.shape[1:], generator=generator, device="cuda") for name in _KDA_REPLAY_BUFFERS: - buffer = mgr._kda_replay_layers[layer_offset][name] + buffer = mgr._kda_replay.pool_layers[layer_offset][name] buffer[slot] = torch.randn(buffer.shape[1:], generator=generator, device="cuda") - mgr.prev_num_accepted_tokens[slot] = 1 + mgr._kda_replay.prev_num_accepted_tokens[slot] = 1 def _kda_verify_round(mgr, slot, layer_weights, round_idx, num_spec): """One verify of 1 + num_spec tokens on every KDA layer for the request in ``slot``, as the KDA mixer runs it; the inputs depend only on the round. Returns the outputs.""" tokens = 1 + num_spec - pending = mgr.prev_num_accepted_tokens[slot : slot + 1].clone() + pending = mgr._kda_replay.prev_num_accepted_tokens[slot : slot + 1].clone() slots = torch.tensor([slot], dtype=torch.int32, device="cuda") cu_seqlens = torch.tensor([0, tokens], dtype=torch.int32, device="cuda") cu_seqlens[:1].sub_(pending) outputs = [] for layer_offset, weights in enumerate(layer_weights): - replay = mgr._kda_replay_layers[layer_offset] + replay = mgr._kda_replay.pool_layers[layer_offset] generator = torch.Generator(device="cuda").manual_seed(1000 * round_idx + layer_offset) num_heads, head_dim = weights["A_log"].shape[0], 128 @@ -4424,7 +4472,7 @@ def _run_kda_verify_rounds(accepted_per_round, suspend_after=None, slot_taken=Fa for round_idx, accepted in enumerate(accepted_per_round): slot = mgr._request_id_to_state_index[1] outputs.append(_kda_verify_round(mgr, slot, layer_weights, round_idx, num_spec)) - mgr.prev_num_accepted_tokens[slot] = accepted + mgr._kda_replay.prev_num_accepted_tokens[slot] = accepted if round_idx == suspend_after: mgr.suspend_request(request) if slot_taken: @@ -4478,1438 +4526,1726 @@ def test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run(slot_taken): torch.testing.assert_close(states, expected_states, atol=1e-3, rtol=1e-3) -@skip_no_cuda -def test_v2_kda_replay_commit_without_the_replay_caches(): - """With the verify kernels' draft records (``kda_state_tok``) in place of the per-draft replay caches, the - suspend-time commit replays a slot's records onto its V2 SSM state and stores its conv window, bit for bit as the - records' replay and the conv store do on their own, and clears only that slot's draft count.""" - from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import ( - replay_kda_token_records, - store_kda_replay_conv_window, +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) - num_layers, num_slots, num_spec, num_heads, head_dim, width = 2, 4, 2, 6, 128, 4 - section_dim = num_heads * head_dim - generator = torch.Generator(device="cuda").manual_seed(13) + def dim_contiguous_conv_cache() -> torch.Tensor: + return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) - def randn(*shape): - return 0.1 * torch.randn(*shape, generator=generator, device="cuda") + 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_k_cache = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) + policy.kda_g_cache = torch.full((2, 4, 2, 2), 7.0) + policy.kda_v_cache = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) + policy.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0, dtype=torch.bfloat16) + 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 = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr._stream = torch.cuda.Stream() - mgr.all_ssm_states = [ - randn(num_slots, num_heads, head_dim, head_dim) for _ in range(num_layers) - ] - mgr.all_conv_states = [randn(num_slots, 3 * section_dim, width - 1) for _ in range(num_layers)] - mgr._kda_replay_layers = [ - { - **{ - name: randn(num_slots, width - 1 + num_spec, section_dim).transpose(-1, -2) - for name in ("kda_conv_q", "kda_conv_k", "kda_conv_v") - }, - "kda_state_tok": randn(num_slots, 3, num_spec, num_heads, head_dim), - } - for _ in range(num_layers) + 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(manager.all_conv_states): + for slot in (1, 3): + for replay_buffer, start in ( + (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(policy.kda_k_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(policy.kda_g_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 + + +def test_v2_kda_draft_records_seed_disaggregated_generation_slots(): + """Without the per-draft replay caches, seeding a transferred request's slot sets its conv caches and clears its + draft count; its drafts' records stay as they are, since with no pending drafts none of them is read.""" + mgr = object.__new__(KimiK3HybridCacheManager) + mgr._kda_replay = KDAReplayState(2, draft_records=True) + mgr._speculative_state = mgr._kda_replay + 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._kda_replay.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, ] - mgr.prev_num_accepted_tokens = torch.tensor([0, 2, 1, 0], dtype=torch.int32, device="cuda") - assert mgr.keeps_kda_draft_records - slot = 1 - slots = torch.tensor([slot], dtype=torch.int32, device="cuda") - num_accepted = mgr.prev_num_accepted_tokens[slots.long()].clone() - expected_ssm, expected_conv = [], [] - for layer_offset, replay in enumerate(mgr._kda_replay_layers): - ssm = mgr.all_ssm_states[layer_offset].clone() - conv = mgr.all_conv_states[layer_offset].clone() - replay_kda_token_records(ssm, replay["kda_state_tok"], slots, num_accepted) - store_kda_replay_conv_window( - conv, - replay["kda_conv_q"], - replay["kda_conv_k"], - replay["kda_conv_v"], - slots, - num_accepted, - ) - expected_ssm.append(ssm) - expected_conv.append(conv) + def dim_contiguous_conv_cache() -> torch.Tensor: + return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) - mgr._commit_kda_pending_drafts(slot) - torch.cuda.synchronize() + mgr._kda_replay.kda_conv_q = dim_contiguous_conv_cache() + mgr._kda_replay.kda_conv_k = dim_contiguous_conv_cache() + mgr._kda_replay.kda_conv_v = dim_contiguous_conv_cache() + mgr._kda_replay.kda_k_cache = None + mgr._kda_replay.kda_g_cache = None + mgr._kda_replay.kda_v_cache = None + mgr._kda_replay.kda_beta_cache = None + mgr._kda_replay.kda_state_tok = torch.full((2, 4, 3, 2, 1, 2), 7.0) - assert mgr.prev_num_accepted_tokens.tolist() == [0, 0, 1, 0] - for actual, expected in zip(mgr.all_ssm_states, expected_ssm): - assert torch.equal(actual, expected) - for actual, expected in zip(mgr.all_conv_states, expected_conv): - assert torch.equal(actual, expected) + mgr._kda_replay._conv_states = tuple(mgr.all_conv_states) + mgr._kda_replay._conv_section_dims = tuple(mgr.conv_section_dims) + mgr._kda_replay._committed_window = mgr.conv_state_shape[1] + mgr.on_state_transfer_complete([101, 202]) + + for layer_offset, conv_state in enumerate(mgr.all_conv_states): + for slot in (1, 3): + for replay_buffer, start in ( + (mgr._kda_replay.kda_conv_q, 0), + (mgr._kda_replay.kda_conv_k, 2), + (mgr._kda_replay.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 mgr._kda_replay.prev_num_accepted_tokens[slot] == 0 + assert mgr._kda_replay.prev_num_accepted_tokens.tolist() == [5, 0, 5, 0] + assert (mgr._kda_replay.kda_state_tok == 7.0).all() -def _build_kda_replay_manager(**kwargs): - """A real V2 hybrid manager with two KDA layers (6 heads, K = V = 8) and replay caches for two drafts.""" - return _build_v2_hybrid_with_mamba_layer( - num_mamba_layers=2, - spec_config=MTPDecodingConfig(max_draft_len=2), - conv_state_layout="q_k_v", - mamba_num_heads=6, - mamba_n_groups=6, - mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=2, - **kwargs, +def test_kda_replay_path_stays_on_with_the_drafts_records(): + """``has_kda_replay_caches`` follows the replay conv caches: it holds with the per-draft replay caches and with + the drafts' records in their place, and not without the replay path.""" + t = torch.zeros(1) + conv_caches = dict(kda_conv_q=t, kda_conv_k=t, kda_conv_v=t) + with_caches = KDAReplayLayerCache( + conv=t, + temporal=t, + kda_k_cache=t, + kda_g_cache=t, + kda_v_cache=t, + kda_beta_cache=t, + **conv_caches, ) + with_records = KDAReplayLayerCache(conv=t, temporal=t, kda_state_tok=t, **conv_caches) + assert with_caches.has_kda_replay_caches + assert with_records.has_kda_replay_caches + assert not KDAReplayLayerCache(conv=t, temporal=t).has_kda_replay_caches + assert not PythonMambaCacheManager.State(conv=t, temporal=t).has_kda_replay_caches -_KDA_REPLAY_FIELDS = (*_KDA_REPLAY_BUFFERS, "prev_num_accepted_tokens") +def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): + """The conv half of the KDA state crosses ctx->gen in bf16. + + ``extract_mamba_kv_cache_params`` forces the KDA delta-rule (ssm) pool to + fp32 but leaves the short-convolution pool at the model dtype, so what the + context rank serializes and the generation rank restores for the conv slot + is bf16, the same dtype as the replay conv caches the generation rank + rebuilds from it (k / v / beta replay caches are bf16 too, the gate cache + fp32). Seeding must copy the restored rows bit for bit, keep the logical + shape, leave the restored bf16 pool itself untouched so the decode path + reuses exactly what was transferred, and clear only the draft scratch of + the seeded slots. + """ + torch.manual_seed(0) + mgr = object.__new__(KimiK3HybridCacheManager) + 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} + 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] + + # Same dtypes as _allocate_pool_replay_buffers: bf16 conv / k / v / beta + # replay caches, fp32 gate cache. + def dim_contiguous_conv_cache() -> torch.Tensor: + return torch.full((2, 4, 6, 2), 7.0, dtype=torch.bfloat16).transpose(-1, -2) + + 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_k_cache = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) + policy.kda_g_cache = torch.full((2, 4, 2, 2), 7.0) + policy.kda_v_cache = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) + policy.kda_beta_cache = torch.full((2, 4, 2, 1), 7.0, dtype=torch.bfloat16) + + 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 ( + (policy.kda_conv_q, 0), + (policy.kda_conv_k, 2), + (policy.kda_conv_v, 4), + ): + # bf16 replay cache, logical shape unchanged by the seeding. + assert replay_buffer.dtype is torch.bfloat16 + seeded = replay_buffer[layer_offset, slot, :, :4] + assert seeded.shape == (2, 4) + # Same dtype on both sides: the restored values are copied bit + # for bit. + torch.testing.assert_close( + seeded, + conv_state[slot, start : start + 2, :], + rtol=0, + atol=0, + ) + assert torch.count_nonzero(replay_buffer[layer_offset, slot, :, 4:]) == 0 + assert torch.count_nonzero(policy.kda_k_cache[layer_offset, slot]) == 0 + assert torch.count_nonzero(policy.kda_g_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 (policy.kda_conv_q[layer_offset, slot] == 7.0).all() + assert (policy.kda_k_cache[layer_offset, slot] == 7.0).all() + assert (policy.kda_g_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 policy.prev_num_accepted_tokens[0] == 5 + assert policy.prev_num_accepted_tokens[2] == 5 @skip_no_cuda -def test_v2_kda_replay_buffers_are_kvcm_state_buffers(): - """Replay caches and the accepted-draft count are registered as SSM-layer buffers, counted in the per-slot state - bytes, and kept out of disaggregated transfers.""" - mgr = _build_kda_replay_manager() +def test_v2_hybrid_static_dynamic_tree_capacity(): + spec_config = MTPDecodingConfig( + max_draft_len=6, + max_total_draft_tokens=31, + use_dynamic_tree=True, + dynamic_tree_max_topK=10, + ) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + spec_config=spec_config, + ) try: - layout = _kda_replay_layout(2, 6, 8, 3) - replay_roles = {str(DataRole(name)) for name in layout} - for layer in mgr.mamba_pp_layers: - local_layer_id = mgr.layer_offsets[layer] - roles = { - str(buffer_id.role) - for pool_group in mgr.impl.pool_group_descs - for variant in pool_group.slot_desc.variants - for coalesced in variant.coalesced_buffers - for buffer_id in coalesced.buffer_ids - if int(buffer_id.layer_id) == local_layer_id - } - assert replay_roles <= roles - assert ("prev_num_accepted_tokens" in roles) == (layer == mgr.mamba_pp_layers[0]) - layer_bytes = sum(dtype.itemsize * math.prod(shape) for dtype, shape in layout.values()) - assert mgr._mamba_state_bytes_per_slot() == ( - 2 * (mgr.ssm_bytes + mgr.conv_bytes) + 2 * layer_bytes + 16 - ) - ignored = {str(role) for role in mgr.get_disagg_ignored_roles()} - assert replay_roles | {"prev_num_accepted_tokens"} <= ignored + assert mgr.intermediate_ssm_states.shape[2] == 32 + assert mgr.intermediate_conv_states.shape[2] == 32 + assert mgr._kv_reserve_draft_tokens == 31 + assert mgr._num_reserved_dummy_slots == 1 + assert mgr.get_replay_state_update_metadata() is None finally: mgr.shutdown() @skip_no_cuda -def test_v2_kda_draft_records_count_in_the_per_slot_budget(): - """The capacity math sees the drafts' records in place of the per-draft replay caches: per draft and layer, a slot - costs vn, beta * k and the decay, fp32 [H, K] each, instead of the k / g / v / beta entries.""" - kwargs = dict( - num_mamba_layers=2, - spec_config=MTPDecodingConfig(max_draft_len=2), - conv_state_layout="q_k_v", - mamba_d_conv=5, - mamba_num_heads=6, - mamba_n_groups=6, - mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=2, +@pytest.mark.parametrize( + "builder", + [_build_hybrid_with_mamba_layer, _build_v2_hybrid_with_mamba_layer], + ids=["cpp", "v2"], +) +def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): + spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) + mgr = builder( + max_batch_size=4, + spec_config=spec_config, + use_replay_state_update=True, ) - plain = _build_v2_hybrid_with_mamba_layer(**kwargs) - try: - plain_bytes = plain._mamba_state_bytes_per_slot() - finally: - plain.shutdown() - with_states = _build_v2_hybrid_with_mamba_layer(kda_draft_records=True, **kwargs) try: - num_heads, _, head_dim = with_states.ssm_state_shape - per_token = 2 * 2 * 3 * num_heads * head_dim * 4 - replay_caches = 2 * sum( - dtype.itemsize * math.prod(shape) - for name, (dtype, shape) in _kda_replay_layout(2, num_heads, head_dim, 4).items() - if name in _KDA_REPLAY_BUFFERS[3:] + replay_metadata = mgr.get_replay_state_update_metadata() + assert replay_metadata is not None + assert replay_metadata.replay_step_width == spec_config.tokens_per_gen_step + assert replay_metadata.replay_history_size == max( + MIN_REPLAY_HISTORY_SIZE, spec_config.tokens_per_gen_step + ) + layer_cache = mgr.mamba_layer_cache(0) + _assert_replay_layer_cache_uses_history_size( + layer_cache, replay_metadata.replay_history_size ) - assert with_states._mamba_state_bytes_per_slot() == plain_bytes - replay_caches + per_token finally: - with_states.shutdown() + mgr.shutdown() @skip_no_cuda -def test_v2_kda_replay_acceptance_and_cuda_graph(): - mgr = _build_kda_replay_manager() +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), + conv_state_layout="q_k_v", + speculative_state=policy, + ) try: - requests = [_make_v2_conversation_request(i, list(range(8)), str(i)) for i in (101, 102)] - for req in requests: - _run_v2_hybrid_context(mgr, req) - slots = torch.tensor(mgr.get_state_indices([101, 102]), device="cuda") - # Mark the second row as padding through the public metadata interface. - mgr.get_state_indices([101, 102], [False, True]) - accepted = torch.tensor([3, 2], dtype=torch.int32, device="cuda") - metadata = SimpleNamespace(num_seqs=2, num_contexts=0) - for layer_id in range(2): - counts = mgr.mamba_layer_cache(layer_id).prev_num_accepted_tokens - assert counts[slots].tolist() == [0, 0] - counts[slots[1]] = 1 - # Warm up on a side stream before capture. - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(stream): - mgr.update_mamba_states(metadata, accepted, state_indices=slots) - torch.cuda.current_stream().wait_stream(stream) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - mgr.update_mamba_states(metadata, accepted, state_indices=slots) - accepted[0] = 2 - graph.replay() - for layer_id in range(2): - assert mgr.mamba_layer_cache(layer_id).prev_num_accepted_tokens[slots].tolist() == [ - 1, - 1, - ] + assert mgr.use_gdn_cached_replay_all_layer_commit + 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] + expected_layer_stride = ( + states[1].data_ptr() - first_state.data_ptr() + ) // first_state.element_size() + assert state_strides == ( + mgr.local_num_mamba_layers, + expected_layer_stride, + first_state.stride(0), + ) finally: mgr.shutdown() @skip_no_cuda -def test_v2_state_index_setup_skips_unchanged_uploads(monkeypatch): - """A step whose state indices and dummy flags did not change enqueues no upload of either: the device buffers - keep values the host never sent. A changed step uploads both again.""" - # The test writes the buffers behind _setup_state_indices to see the skip, which the check of skipped copies - # reports. - monkeypatch.setattr("tensorrt_llm._device_copy._verify_skipped_device_copies", False) - - class StateCache: - def __init__(self, slot: int) -> None: - self.slot = slot - - def get_ssm_block_base_index(self, _layer_group_id) -> int: - return self.slot - - def close(self) -> None: - pass +def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): + from tensorrt_llm._torch.modules.fla.cached_replay import ( + commit_gdn_cached_replay_history_layers, + ) + batch_size = 16 + policy = GDNReplayState(tokens_per_gen_step=4) mgr = _build_v2_hybrid_with_mamba_layer( - spec_config=MTPDecodingConfig(max_draft_len=2), + max_batch_size=batch_size, + num_mamba_layers=3, + spec_config=MTPDecodingConfig(max_draft_len=3), conv_state_layout="q_k_v", - mamba_n_groups=4, - mamba_ssm_cache_dtype=torch.float32, - kda_replay_num_spec=2, + speculative_state=policy, ) - live = SimpleNamespace(py_request_id=101, is_dummy=False) - padding = SimpleNamespace(py_request_id=202, is_dummy=True) - mgr.kv_cache_map.update({101: StateCache(0), 202: StateCache(1)}) try: - mgr._setup_state_indices([live, padding]) - torch.cuda.synchronize() - assert mgr.cuda_state_indices[:2].tolist() == [0, 1] - assert mgr._dummy_request_mask[:2].tolist() == [False, True] + torch.manual_seed(1234) + for state in mgr.all_ssm_states: + state.normal_() + policy.old_x.normal_() + policy.old_B.normal_() + policy.old_dt.uniform_(-0.2, -0.01) - mgr.cuda_state_indices.fill_(-1) - mgr._dummy_request_mask.fill_(True) - mgr._setup_state_indices([live, padding]) - torch.cuda.synchronize() - assert mgr.cuda_state_indices[:2].tolist() == [-1, -1] - assert mgr._dummy_request_mask[:2].tolist() == [True, True] + positions = torch.arange(batch_size, device="cuda", dtype=torch.int32) + replay_work_items = torch.stack( + ( + positions, + positions, + torch.full_like(positions, policy.replay_history_size), + torch.zeros_like(positions), + ), + dim=1, + ) + n_writes = torch.tensor([batch_size], device="cuda", dtype=torch.int32) + expected = torch.stack(mgr.all_ssm_states) + commit_gdn_cached_replay_history_layers( + ssm_states=expected, + 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=policy.replay_history_size, + ) - mgr.kv_cache_map[202].slot = 2 - padding.is_dummy = False - mgr._setup_state_indices([live, padding]) - torch.cuda.synchronize() - assert mgr.cuda_state_indices[:2].tolist() == [0, 2] - assert mgr._dummy_request_mask[:2].tolist() == [False, False] + policy.commit_all_layers( + SimpleNamespace( + mamba_metadata=SimpleNamespace( + replay_num_decodes=batch_size, + replay_work_items=replay_work_items, + replay_n_writes=n_writes, + ) + ), + batch_size, + ) + + torch.testing.assert_close(torch.stack(mgr.all_ssm_states), expected, rtol=0, atol=0) finally: mgr.shutdown() -@skip_no_cuda -@pytest.mark.parametrize( - "role,is_disagg", [("context", True), ("generation", True), (None, True), ("context", False)] -) -def test_v2_kda_worker_role_cache(monkeypatch, role, is_disagg): - """Only a disaggregated context server drops the KDA replay and draft-verify state.""" - if role is None: - monkeypatch.delenv("TRTLLM_DISAGG_ROLE", raising=False) - else: - monkeypatch.setenv("TRTLLM_DISAGG_ROLE", role) - mgr = _build_kda_replay_manager(is_disagg=is_disagg) - try: - req = _make_v2_conversation_request(123, list(range(8)), "role") - _run_v2_hybrid_context(mgr, req) - cache = mgr.mamba_layer_cache(0) - ctx_only = is_disagg and role == "context" - assert mgr.kda_context_only == ctx_only - assert mgr.use_kda_replay_update != ctx_only - assert isinstance(cache, PythonMambaCacheManager.SpeculativeState) != ctx_only - assert cache.conv.shape[1:] == (144, 3) - assert cache.temporal.shape[1:] == (6, 8, 8) - if ctx_only: - assert mgr.get_intermediate_ssm_states(0) is None - assert mgr.get_intermediate_conv_states(0) is None - assert not mgr.get_disagg_ignored_roles() & { - DataRole(name) for name in _KDA_REPLAY_FIELDS - } - mgr.update_mamba_states( - SimpleNamespace(num_seqs=1, num_contexts=1), - torch.ones(1, device="cuda", dtype=torch.int32), - ) - else: - slot = mgr.get_state_indices([123])[0] - assert cache.prev_num_accepted_tokens[slot].item() == 0 - finally: - mgr.shutdown() +def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): + mgr = object.__new__(Qwen35HybridCacheManager) + batch_size = 16 + mgr.local_num_mamba_layers = 1 + 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_buffer = DeviceCopyBuffer(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)] + policy._ssm_states = tuple(mgr.all_ssm_states) + policy._conv_states = tuple(mgr.all_conv_states) + events = [] + 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.modules.fla.cache_manager._advance_replay_state", + lambda *_args, **_kwargs: events.append("advance"), + ) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", + lambda *_args, **_kwargs: None, + ) -@skip_no_cuda -@pytest.mark.parametrize("kda_draft_records", [False, True], ids=["replay_caches", "draft_records"]) -@pytest.mark.parametrize("snapshot", [False, True]) -def test_v2_kda_replay_survives_cache_lifecycle(snapshot, kda_draft_records): - """KVCM moves the replay caches, or the drafts' records in their place, with the recurrent state through - suspend, offload to host and snapshot reuse.""" - mgr = _build_kda_replay_manager( - host_cache_size=16 << 20, enable_stats=True, kda_draft_records=kda_draft_records + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(64, dtype=torch.bool)) + mgr.update_mamba_states( + 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), ) - replay_fields = (*mgr._kda_replay_layout(), "prev_num_accepted_tokens") - assert ("kda_state_tok" in replay_fields) == kda_draft_records - kv = None - reused = None - pressure = [] - try: - stream = torch.cuda.current_stream().cuda_stream - kv = mgr.impl.create_kv_cache() - assert kv.resume(stream) - assert kv.resize(32) - group = mgr.impl.get_layer_group_id(LayerId(0)) - slot = kv.get_ssm_block_base_index(group) - expected = [] - for layer_id in range(2): - cache = mgr.mamba_layer_cache(layer_id) - fields = {} - for offset, name in enumerate((*replay_fields, "conv", "temporal")): - tensor = getattr(cache, name) - # Accepted counts describe the request and are shared across layers. - value = offset + 1 if name == "prev_num_accepted_tokens" else offset + layer_id + 1 - tensor[slot].fill_(value) - fields[name] = tensor[slot].clone() - expected.append(fields) - if snapshot: - kv.commit(list(range(32))) - kv.close() - kv = None - else: - kv.suspend() - # Fill the recurrent pool with newer suspended requests, forcing - # the original state to host while respecting KVCM's quota floor. - pool = next( - pool - for pool in mgr.impl.pool_group_descs - if any(variant.layer_group_id == group for variant in pool.slot_desc.variants) - ) - for _ in range(pool.num_slots + 1): - other = mgr.impl.create_kv_cache() - pressure.append(other) - assert other.resume(stream) - other.suspend() - for other in pressure: - other.commit_pending_stats() - other.close() - pressure.clear() - stats = mgr.impl.get_and_reset_iteration_stats() - assert stats[group].iter_offload_blocks > 0 - if snapshot: - reused = mgr.impl.create_kv_cache(input_tokens=list(range(32))) - assert reused.num_committed_tokens == 32 - active = reused - else: - active = kv - assert active.resume(stream) - assert active.resize(33) - new_slot = active.get_ssm_block_base_index(group) - for layer_id, fields in enumerate(expected): - cache = mgr.mamba_layer_cache(layer_id) - for name, value in fields.items(): - torch.testing.assert_close(getattr(cache, name)[new_slot], value, rtol=0, atol=0) - torch.cuda.synchronize() - finally: - for other in pressure: - other.close() - if reused is not None: - reused.close() - if kv is not None: - kv.close() - mgr.shutdown() + assert events == ["commit", "advance"] -@skip_no_cuda -def test_v2_kda_asymmetric_context_generation_transfer(monkeypatch): - """A context server without replay caches and a generation server with them describe the same transferable - recurrent state, and the generation server seeds its replay caches from the transferred conv state.""" - from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( - MambaPolicy, - mamba_receiver_payload_bytes, - ) - from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo - monkeypatch.setenv("TRTLLM_DISAGG_ROLE", "context") - ctx = _build_kda_replay_manager(is_disagg=True) - monkeypatch.setenv("TRTLLM_DISAGG_ROLE", "generation") - gen = _build_kda_replay_manager(is_disagg=True) - try: - assert ctx.kda_context_only and not gen.kda_context_only - for mgr in (ctx, gen): - _run_v2_hybrid_context(mgr, _make_v2_conversation_request(101, [1, 2], "transfer")) - src_slot = ctx.get_state_indices([101])[0] - dst_slot = gen.get_state_indices([101])[0] - src_table = build_page_table_from_manager(ctx) - dst_table = build_page_table_from_manager(gen) - src_rank = RankInfo.from_kv_cache_manager("ctx", ctx, device_id=0) - dst_rank = RankInfo.from_kv_cache_manager("gen", gen, device_id=0) - MambaPolicy.validate_peer_compatible(src_rank, dst_rank, src_table, dst_table) - src_group = next(g for g in src_table.layer_groups if isinstance(g, MambaLayerGroup)) - dst_group = next(g for g in dst_table.layer_groups if isinstance(g, MambaLayerGroup)) - for group in (src_group, dst_group): - assert {pv.pool_role for pv in group.pool_views} == {MAMBA_CONV_ROLE, MAMBA_SSM_ROLE} - assert mamba_receiver_payload_bytes(src_table, dst_table, dst_slot) == sum( - gen.mamba_layer_cache(i).conv[dst_slot].numel() * 2 - + gen.mamba_layer_cache(i).temporal[dst_slot].numel() * 4 - for i in range(2) - ) +def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): + mgr = object.__new__(NemotronHybridCacheManager) + mgr.local_num_mamba_layers = 1 + mgr._request_id_to_state_index = { + 100: 0, + 101: 1, + 102: 2, + 103: 3, + } + mgr._request_id_to_is_dummy = { + 100: False, + 101: False, + 102: True, + 103: False, + } + policy = ReplayHistory(tokens_per_gen_step=5) + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.bool)) + mgr._dummy_request_mask_host = torch.zeros(4, dtype=torch.bool) + 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)] + 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, + ) - # What the transfer writes into the generation slot. - for layer_id in range(2): - source = ctx.mamba_layer_cache(layer_id) - target = gen.mamba_layer_cache(layer_id) - source.conv[src_slot].fill_(layer_id + 1) - source.temporal[src_slot].fill_(layer_id + 2) - target.conv[dst_slot].copy_(source.conv[src_slot]) - target.temporal[dst_slot].copy_(source.temporal[src_slot]) - gen.seed_kda_replay_caches_for_disagg_gen([101]) - for layer_id in range(2): - source = ctx.mamba_layer_cache(layer_id) - target = gen.mamba_layer_cache(layer_id) - assert target.prev_num_accepted_tokens[dst_slot] == 0 - for section, name in enumerate(_KDA_REPLAY_BUFFERS[:3]): - torch.testing.assert_close( - getattr(target, name)[dst_slot, :, :3], - source.conv[src_slot, section * 48 : (section + 1) * 48].to(torch.bfloat16), - ) - finally: - ctx.shutdown() - gen.shutdown() + request_ids = [100, 101, 102, 103] + state_indices = torch.tensor( + mgr.get_state_indices(request_ids, [False, False, False, True]), + dtype=torch.int32, + ) + assert mgr._dummy_request_mask.tolist() == [False, False, True, True] + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(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, + ) -def _kda_cache_budget(monkeypatch, available, model, mapping, **kwargs): - from tensorrt_llm._torch.modules.kimi_kda import _kda_kernels + assert policy.prev_num_accepted_tokens.tolist() == [13, 3, 13, 13] + assert policy.cache_buf_idx.tolist() == [1, 0, 1, 1] - monkeypatch.setattr(_kda_kernels, "is_kda_mtp_verify_available", lambda: available) - monkeypatch.setattr(_kda_kernels, "kda_mtp_max_num_spec", lambda: 7) - return MambaHybridCacheManagerV2.get_cache_size_per_token(model, mapping, **kwargs) +def test_v2_hybrid_dynamic_tree_promotes_accepted_leaf_state(monkeypatch): + mgr = object.__new__(NemotronHybridCacheManager) + mgr.local_num_mamba_layers = 1 + 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)) + mgr._speculative_state = policy + mgr.all_ssm_states = [torch.empty(0)] + mgr.all_conv_states = [torch.empty(0)] + policy._ssm_states = tuple(mgr.all_ssm_states) + policy._conv_states = tuple(mgr.all_conv_states) -@pytest.mark.parametrize( - "role,is_disagg,available", - [ - ("context", True, True), - ("generation", True, True), - (None, True, True), - ("context", False, True), - ("generation", True, False), - ], -) -@pytest.mark.parametrize("interval,replay_slope", [(0, 0), (32, 95), (33, 92), (48, 63), (256, 0)]) -def test_kda_cache_budget_accounts_for_replay_role( - monkeypatch, role, is_disagg, available, interval, replay_slope -): - if role is None: - monkeypatch.delenv("TRTLLM_DISAGG_ROLE", raising=False) - else: - monkeypatch.setenv("TRTLLM_DISAGG_ROLE", role) - model = _kimi_model_config() - spec_config = MTPDecodingConfig(max_draft_len=2) - mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) - kwargs = dict( - max_batch_size=2, - max_seq_len=128, - tokens_per_block=32, - kv_cache_config=KvCacheConfig( - enable_block_reuse=True, - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=interval), - ), - spec_config=spec_config, - is_disagg=is_disagg, - ) - baseline = _kda_cache_budget(monkeypatch, False, model, mapping, **kwargs) - result = _kda_cache_budget(monkeypatch, available, model, mapping, **kwargs) - if available and not (is_disagg and role == "context"): - # Two layers: three BF16 [5,32] conv windows, two BF16 and one - # FP32 [2,32] histories and BF16 [2,8] padded beta per layer; - # one shared count per slot, padded to 16 bytes. - per_layer = 3 * 5 * 32 * 2 + 2 * 32 * 8 + 2 * 8 * 2 - # At interval 33, round the combined recurrent + replay snapshot - # cost once: ceil(6224 / 33) - ceil(3200 / 33) = 92 bytes/token. - assert result[0] - baseline[0] == replay_slope - live_and_dummy_slots = 2 + _get_num_cuda_graph_padding_dummy_slots(spec_config, 2) - assert result[1] - baseline[1] == live_and_dummy_slots * (2 * per_layer + 16) - else: - assert result == baseline + promoted_positions = [] + def capture_promoted_position(_dst, _src, _src_indices, positions, _dst_indices): + promoted_positions.append(positions.clone()) -@pytest.mark.parametrize("enable_block_reuse", [False, True]) -@pytest.mark.parametrize("enable_attention_dp", [False, True]) -@pytest.mark.parametrize("rank", [0, 1]) -@pytest.mark.parametrize("pp_partition,stage_kda_layers", [(None, [2, 0]), ([2, 2], [1, 1])]) -def test_kda_cache_budget_counts_live_dummy_and_fixed_snapshot_slots( - monkeypatch, enable_block_reuse, enable_attention_dp, rank, pp_partition, stage_kda_layers -): - model = _kimi_model_config() - mapping = Mapping( - world_size=2, - rank=rank, - tp_size=1, - pp_size=2, - enable_attention_dp=enable_attention_dp, - pp_partition=pp_partition, + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", + capture_promoted_position, ) - spec_config = MTPDecodingConfig(max_draft_len=2, draft_len_schedule={1: 2, 2: 1}) - kwargs = dict( - max_batch_size=4, - max_seq_len=128, - tokens_per_block=32, - kv_cache_config=KvCacheConfig( - enable_block_reuse=enable_block_reuse, - mamba_state_config=MambaStateConfig( - periodic_snapshot_interval=0, - additional_snapshot_offsets_from_start=[32, 64, 64, 256], - additional_snapshot_offsets_from_end=[0, 32, 128], - ), - ), - spec_config=spec_config, - is_disagg=False, + + mgr._generation_state_indices = torch.arange(64, dtype=torch.int32) + if not hasattr(mgr, "_dummy_request_mask"): + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(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), + state_indices=torch.tensor([9, 10, 11], dtype=torch.int32), + accepted_leaf_positions=torch.tensor([4, 7], dtype=torch.int64), ) - baseline = _kda_cache_budget(monkeypatch, False, model, mapping, **kwargs) - result = _kda_cache_budget(monkeypatch, True, model, mapping, **kwargs) - # Without an explicit partition, splitting the four target layers plus - # one MTP layer puts both KDA layers on rank 0. Partition [2,2] places one - # KDA layer on each rank, with the MTP layer appended to rank 1. - # Every live slot (max_batch_size * pp_size), every CUDA-graph padding - # dummy plus one more for attention DP, and with reuse four reachable, - # distinct snapshot rules per resident sequence carry replay caches. - local_kda_layers = stage_kda_layers[rank] - replay_bytes_per_slot = 1504 * local_kda_layers + (16 if local_kda_layers else 0) - slots = ( - 8 - + _get_num_cuda_graph_padding_dummy_slots(spec_config, 4) - + int(enable_attention_dp) - + (8 * 4 if enable_block_reuse else 0) + + assert len(promoted_positions) == 2 + assert all( + torch.equal(positions, torch.tensor([4, 7], dtype=torch.int32)) + for positions in promoted_positions ) - assert result == (baseline[0], baseline[1] + slots * replay_bytes_per_slot) -def test_v2_kda_replay_seeds_disaggregated_generation_slots(): - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr._kda_replay_unseeded = set() - 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, - ] +@skip_no_cuda +@pytest.mark.parametrize( + "mamba_ssm_cache_dtype", + [torch.float16, torch.float32, torch.bfloat16], +) +def test_cpp_hybrid_passes_per_window_pool_dtypes_for_nvfp4_kv_cache( + mamba_ssm_cache_dtype, +): + mgr = _build_hybrid_with_mamba_layer( + dtype=DataType.NVFP4, + mamba_ssm_cache_dtype=mamba_ssm_cache_dtype, + ) + recurrent_pool_dtype = torch_dtype_to_binding(mamba_ssm_cache_dtype) - def dim_contiguous_conv_cache() -> torch.Tensor: - return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) + expected_dtypes = [ + (LinearCacheType.RECURRENT_STATES.value, recurrent_pool_dtype), + (128, DataType.NVFP4), + ] + assert [ + (config.window_size, config.dtype) for config in mgr.pool_configurations + ] == expected_dtypes + assert [ + (config.window_size, config.dtype) for config in mgr.impl.pool_configurations + ] == expected_dtypes + assert mgr._layer_to_pool_idx == {0: 0, 1: 1} + assert mgr.recurrent_states_pool_index == 0 + assert mgr.impl.get_recurrent_states_pool().dtype == mamba_ssm_cache_dtype - replay = {} - replay["kda_conv_q"] = dim_contiguous_conv_cache() - replay["kda_conv_k"] = dim_contiguous_conv_cache() - replay["kda_conv_v"] = dim_contiguous_conv_cache() - replay["kda_k_cache"] = torch.full((2, 4, 2, 2), 7.0) - replay["kda_g_cache"] = torch.full((2, 4, 2, 2), 7.0) - replay["kda_v_cache"] = torch.full((2, 4, 2, 2), 7.0) - replay["kda_beta_cache"] = torch.full((2, 4, 2, 1), 7.0) + compact_scale_pointers = mgr.impl.get_block_scale_pool_pointers() + assert mgr.impl.get_block_pool_pointers().shape == (2, 2) + assert compact_scale_pointers.shape == (1, 2) + assert mgr.kv_cache_pool_pointers.shape == (2, 2, 2) + assert torch.count_nonzero(mgr.kv_cache_pool_pointers[0, :, 1]) == 0 + assert torch.equal(mgr.kv_cache_pool_pointers[1, :, 1], compact_scale_pointers[0]) - mgr._kda_replay_layers = [ - {name: buffer[layer] for name, buffer in replay.items()} for layer in range(2) - ] - mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) +@skip_no_cuda +def test_cpp_hybrid_merges_compact_scale_rows_with_unmanaged_layers(): + mgr = _build_hybrid_with_mamba_layer( + dtype=DataType.NVFP4, + mamba_layer_mask=[True, False, True, False], + attention_layer_mask=[False, False, False, True], + ) - for layer_offset, conv_state in enumerate(mgr.all_conv_states): - for slot in (1, 3): - for replay_buffer, start in ( - (replay["kda_conv_q"], 0), - (replay["kda_conv_k"], 2), - (replay["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(replay["kda_k_cache"][layer_offset, slot]) == 0 - assert torch.count_nonzero(replay["kda_g_cache"][layer_offset, slot]) == 0 - assert torch.count_nonzero(replay["kda_v_cache"][layer_offset, slot]) == 0 - assert torch.count_nonzero(replay["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 mgr.pp_layers == [0, 2, 3] + recurrent_window = LinearCacheType.RECURRENT_STATES.value + assert mgr.max_attention_window_vec == [ + recurrent_window, + recurrent_window, + mgr.max_seq_len, + ] + assert mgr.kv_cache_pool_mapping[:, 0].tolist() == [0, 0, 1] + compact_scale_pointers = mgr.impl.get_block_scale_pool_pointers() + assert compact_scale_pointers.shape == (1, 2) + assert mgr.kv_cache_pool_pointers.shape == (2, 2, 2) + assert torch.count_nonzero(mgr.kv_cache_pool_pointers[0, :, 1]) == 0 + assert torch.equal(mgr.kv_cache_pool_pointers[1, :, 1], compact_scale_pointers[0]) -def test_v2_kda_draft_records_seed_disaggregated_generation_slots(): - """Without the per-draft replay caches, seeding a transferred request's slot sets its conv caches and clears its - draft count; its drafts' records stay as they are, since with no pending drafts none of them is read.""" - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr._kda_replay_unseeded = set() - 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, - ] +@skip_no_cuda +def test_cpp_hybrid_recurrent_pool_reserves_cuda_graph_padding_slot(): + """Without spec decoding, the recurrent-state snapshot pool must + have at least max_batch_size + 1 slots — one extra for the + CUDA-graph padding sentinel (CUDA_GRAPH_DUMMY_REQUEST_ID). Without + it, the padding sentinel evicts live recurrent state under load.""" + max_batch_size = 4 + mgr = _build_hybrid_with_mamba_layer(spec_config=None, max_batch_size=max_batch_size) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + assert recurrent_primary >= max_batch_size + 1, ( + f"recurrent-state pool has {recurrent_primary} slots, " + f"need >= max_batch_size + 1 = {max_batch_size + 1} to host the " + f"CUDA-graph padding sentinel without evicting live state" + ) - def dim_contiguous_conv_cache() -> torch.Tensor: - return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) - replay = {} - replay["kda_conv_q"] = dim_contiguous_conv_cache() - replay["kda_conv_k"] = dim_contiguous_conv_cache() - replay["kda_conv_v"] = dim_contiguous_conv_cache() - replay["kda_state_tok"] = torch.full((2, 4, 3, 2, 1, 2), 7.0) +@skip_no_cuda +def test_cpp_hybrid_recurrent_pool_reserves_draft_len_sentinel_slots(): + """With spec decoding, CUDAGraphRunner._get_padded_batch issues a + distinct dummy request id for each runtime_draft_len in + [0, max_draft_len], so the recurrent-state snapshot pool must reserve + one slot per draft length on top of the CUDA-graph padding slot.""" + max_batch_size, max_draft_len = 4, 2 + spec_config = MTPDecodingConfig(max_draft_len=max_draft_len) + mgr = _build_hybrid_with_mamba_layer(spec_config=spec_config, max_batch_size=max_batch_size) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + expected_min = max_batch_size + 1 + max_draft_len + assert recurrent_primary >= expected_min, ( + f"recurrent-state pool has {recurrent_primary} slots, " + f"need >= max_batch_size + 1 + max_draft_len = {expected_min} so " + f"per-draft-len sentinels don't collide with live state" + ) - mgr._kda_replay_layers = [ - {name: buffer[layer] for name, buffer in replay.items()} for layer in range(2) - ] - assert mgr.keeps_kda_draft_records - mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) +def _build_hybrid_with_mamba_layer_pp( + spec_config=None, max_batch_size=4, enable_block_reuse=False, pp_size=2 +): + """Same as ``_build_hybrid_with_mamba_layer`` but with ``pp_size`` >= 1. - for layer_offset, conv_state in enumerate(mgr.all_conv_states): - for slot in (1, 3): - for replay_buffer, start in ( - (replay["kda_conv_q"], 0), - (replay["kda_conv_k"], 2), - (replay["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 mgr.prev_num_accepted_tokens[slot] == 0 - assert mgr.prev_num_accepted_tokens.tolist() == [5, 0, 5, 0] - assert (replay["kda_state_tok"] == 7.0).all() + Uses ``world_size = pp_size`` and ``rank = 0`` so the real C++ KVCacheManager + still goes through its single-process path while the Python pool-sizing + code sees ``mapping.pp_size > 1``. Constructs ``pp_size * 2`` total layers + (alternating mamba/attn) so that each PP slice has both a mamba and an + attention layer — otherwise some ranks would hit a slope=0 edge case in + the affine memory model when block reuse is disabled. + """ + pairs = pp_size # one (mamba, attn) pair per PP stage so every rank has both kinds + mamba_mask = [True, False] * pairs + attn_mask = [False, True] * pairs + mamba_num_layers = sum(mamba_mask) + num_layers = sum(attn_mask) + mapping = Mapping(world_size=pp_size, rank=0, tp_size=1, pp_size=pp_size) + kv_cache_config = KvCacheConfig(max_tokens=512, enable_block_reuse=enable_block_reuse) + return CppMambaHybridCacheManager( + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=mamba_num_layers, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=num_layers, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=max_batch_size, + mapping=mapping, + spec_config=spec_config, + layer_mask=attn_mask, + ) -def test_kda_replay_path_stays_on_with_the_drafts_records(): - """``has_kda_replay_caches`` follows the replay conv caches: it holds with the per-draft replay caches and with - the drafts' records in their place, and not without the replay path.""" - t = torch.zeros(1) - conv_caches = dict(kda_conv_q=t, kda_conv_k=t, kda_conv_v=t) - with_caches = PythonMambaCacheManager.SpeculativeState( - conv=t, - temporal=t, - kda_k_cache=t, - kda_g_cache=t, - kda_v_cache=t, - kda_beta_cache=t, - **conv_caches, +# Skip when running under pytest --run-ray (sets TLLM_DISABLE_MPI=1). In that +# mode ``Mapping`` resolves to ``DeviceMeshTopology``, whose ``pp_rank`` +# requires torch.distributed initialisation that isn't available in this +# single-process unit test. The pp-sharding behaviour exercised here is +# orthogonal to the Ray orchestrator. +_skip_under_ray = pytest.mark.skipif( + os.environ.get("TLLM_DISABLE_MPI") == "1", + reason="pp_size>1 helper builds Mapping with world_size>1 which needs " + "torch.distributed under TLLM_DISABLE_MPI=1 (Ray) sessions", +) + + +@_skip_under_ray +@skip_no_cuda +@pytest.mark.parametrize("pp_size", [2, 4]) +def test_cpp_hybrid_recurrent_pool_scales_with_pp_size(pp_size): + """With pipeline parallelism, multiple microbatches are in-flight on the + same rank concurrently, each holding up to ``max_batch_size`` sequences' + Mamba state. The recurrent-state pool must therefore size for + ``max_batch_size * pp_size`` live slots (plus the CUDA-graph padding + sentinel). Without this scaling, the first inference batch under PP>1 + trips ``No free block found`` once requests beyond the first microbatch + enter the pool (cf. TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2]). + """ + max_batch_size = 4 + mgr = _build_hybrid_with_mamba_layer_pp( + spec_config=None, max_batch_size=max_batch_size, pp_size=pp_size ) - with_records = PythonMambaCacheManager.SpeculativeState( - conv=t, temporal=t, kda_state_tok=t, **conv_caches + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + expected_min = max_batch_size * pp_size + 1 + assert recurrent_primary >= expected_min, ( + f"recurrent-state pool has {recurrent_primary} slots with pp_size={pp_size}, " + f"need >= max_batch_size * pp_size + 1 = {expected_min} so concurrent " + f"in-flight microbatches don't exhaust live-state slots" ) - assert with_caches.has_kda_replay_caches - assert with_records.has_kda_replay_caches - assert not PythonMambaCacheManager.SpeculativeState(conv=t, temporal=t).has_kda_replay_caches - assert not PythonMambaCacheManager.State(conv=t, temporal=t).has_kda_replay_caches -def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): - """The conv half of the KDA state crosses ctx->gen in bf16. +@skip_no_cuda +def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): + """With block reuse enabled, the block-reuse branch must not drop the + live-state + CUDA-graph-padding floor. - ``extract_mamba_kv_cache_params`` forces the KDA delta-rule (ssm) pool to - fp32 but leaves the short-convolution pool at the model dtype, so what the - context rank serializes and the generation rank restores for the conv slot - is bf16, the same dtype as the replay conv caches the generation rank - rebuilds from it (k / v / beta replay caches are bf16 too, the gate cache - fp32). Seeding must copy the restored rows bit for bit, keep the logical - shape, leave the restored bf16 pool itself untouched so the decode path - reuses exactly what was transferred, and clear only the draft scratch of - the seeded slots. + With max_batch_size=4, periodic_snapshot_interval=256, max_tokens=512: + naive: max_snapshots = 512 // 256 = 2 (drops live-state floor!) + fixed: max_snapshots = max(2, 4 + 1) = 5 """ - torch.manual_seed(0) - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr._kda_replay_unseeded = set() - 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) - # 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)] - restored = [conv.clone() for conv in mgr.all_conv_states] + max_batch_size = 4 + mgr = _build_hybrid_with_mamba_layer( + spec_config=None, + max_batch_size=max_batch_size, + enable_block_reuse=True, + periodic_snapshot_interval=256, + ) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + assert recurrent_primary >= max_batch_size + 1, ( + f"recurrent-state pool has {recurrent_primary} slots with block reuse enabled, " + f"need >= max_batch_size + 1 = {max_batch_size + 1} to prevent the padding " + f"sentinel from evicting live recurrent state" + ) - # Same dtypes as _allocate_pool_replay_buffers: bf16 conv / k / v / beta - # replay caches, fp32 gate cache. - def dim_contiguous_conv_cache() -> torch.Tensor: - return torch.full((2, 4, 6, 2), 7.0, dtype=torch.bfloat16).transpose(-1, -2) - replay = {} - replay["kda_conv_q"] = dim_contiguous_conv_cache() - replay["kda_conv_k"] = dim_contiguous_conv_cache() - replay["kda_conv_v"] = dim_contiguous_conv_cache() - replay["kda_k_cache"] = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) - replay["kda_g_cache"] = torch.full((2, 4, 2, 2), 7.0, dtype=torch.float32) - replay["kda_v_cache"] = torch.full((2, 4, 2, 2), 7.0, dtype=torch.bfloat16) - replay["kda_beta_cache"] = torch.full((2, 4, 2, 1), 7.0, dtype=torch.bfloat16) +@skip_no_cuda +def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): + """Dry-run path (is_estimating_kv_cache=True) under block reuse must + keep the live-state floor *plus* room for snapshots, not collapse to + max(snapshots, live). With max_batch_size=4, interval=256, max_tokens=512: + old: max_snapshots = max(512//256, 4) = 4 (no headroom for snapshots) + new: max_snapshots = 4 + 512//256 = 6 (live + snapshots) + """ + max_batch_size = 4 + mgr = _build_hybrid_with_mamba_layer( + spec_config=None, + max_batch_size=max_batch_size, + enable_block_reuse=True, + periodic_snapshot_interval=256, + is_estimating_kv_cache=True, + ) + recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] + # 4 live state slots + 2 reuse snapshots = 6. + expected_min = max_batch_size + (512 // 256) + assert recurrent_primary >= expected_min, ( + f"dry-run recurrent-state pool has {recurrent_primary} slots, " + f"need >= live_state + reuse_snapshots = {expected_min}; the old " + f"max(reuse, live) formula dropped reuse headroom" + ) - mgr._kda_replay_layers = [ - {name: buffer[layer] for name, buffer in replay.items()} for layer in range(2) - ] - mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) +# --------------------------------------------------------------------------- +# CppMambaHybridCacheManager: rank with zero local mamba layers +# +# Regression test for the early-exit path added when a rank ends up with no +# mamba layers (e.g. under PP sharding when all mamba layers fall on other +# ranks). On that path, the constructor must: +# - call the real parent KVCacheManager with the union layer_mask and +# num_layers=num_layers (not mamba_num_layers + num_layers), +# - skip allocating any mamba-only state, and +# - leave self.requests = [] so Mamba-only hooks and metadata preparation +# can no-op without touching uninitialized state. +# +# We exercise the same Python branch with world_size=1 (so the real C++ +# KVCacheManager init doesn't need MPI) and a layer mask that contains zero +# mamba layers. +# --------------------------------------------------------------------------- - 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 ( - (replay["kda_conv_q"], 0), - (replay["kda_conv_k"], 2), - (replay["kda_conv_v"], 4), - ): - # bf16 replay cache, logical shape unchanged by the seeding. - assert replay_buffer.dtype is torch.bfloat16 - seeded = replay_buffer[layer_offset, slot, :, :4] - assert seeded.shape == (2, 4) - # Same dtype on both sides: the restored values are copied bit - # for bit. - torch.testing.assert_close( - seeded, - conv_state[slot, start : start + 2, :], - rtol=0, - atol=0, - ) - assert torch.count_nonzero(replay_buffer[layer_offset, slot, :, 4:]) == 0 - assert torch.count_nonzero(replay["kda_k_cache"][layer_offset, slot]) == 0 - assert torch.count_nonzero(replay["kda_g_cache"][layer_offset, slot]) == 0 - assert torch.count_nonzero(replay["kda_v_cache"][layer_offset, slot]) == 0 - assert torch.count_nonzero(replay["kda_beta_cache"][layer_offset, slot]) == 0 - assert mgr.prev_num_accepted_tokens[slot] == 0 - # Slots no request restored keep their pre-seed contents. - for slot in (0, 2): - assert (replay["kda_conv_q"][layer_offset, slot] == 7.0).all() - assert (replay["kda_k_cache"][layer_offset, slot] == 7.0).all() - assert (replay["kda_g_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 +def _build_zero_mamba_hybrid( + enable_block_reuse=False, + periodic_snapshot_interval=256, +): + """Construct a real CppMambaHybridCacheManager whose this-rank slice has + no mamba layers. world_size=1 / pp_size=1 keeps the real parent + KVCacheManager off the MPI path.""" + # [other, other, full_attn, full_attn] + mamba_mask = [False, False, False, False] + attn_mask = [False, False, True, True] + mamba_num_layers = sum(mamba_mask) # 0 + num_layers = sum(attn_mask) # 4 -@skip_no_cuda -def test_v2_hybrid_static_dynamic_tree_capacity(): - spec_config = MTPDecodingConfig( - max_draft_len=6, - max_total_draft_tokens=31, - use_dynamic_tree=True, - dynamic_tree_max_topK=10, + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + # Cap KV pool size so the real C++ allocator only takes a tiny slice of + # GPU memory; we don't actually use the cache. + kv_cache_config = KvCacheConfig( + max_tokens=128, + enable_block_reuse=enable_block_reuse, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=periodic_snapshot_interval), ) - mgr = _build_v2_hybrid_with_mamba_layer( - max_batch_size=4, - spec_config=spec_config, + + mgr = CppMambaHybridCacheManager( + # mamba cache parameters — values are unused on the early-exit path + # but must be type-valid. + mamba_d_state=8, + mamba_d_conv=4, + mamba_num_heads=4, + mamba_n_groups=1, + mamba_head_dim=8, + mamba_num_layers=mamba_num_layers, + mamba_layer_mask=mamba_mask, + mamba_cache_dtype=torch.float16, + mamba_ssm_cache_dtype=torch.float16, + # kv cache parameters + kv_cache_config=kv_cache_config, + kv_cache_type=CacheTypeCpp.SELF, + num_layers=num_layers, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_num_tokens=96, + max_batch_size=2, + mapping=mapping, + spec_config=None, + layer_mask=attn_mask, ) - try: - assert mgr.intermediate_ssm_states.shape[2] == 32 - assert mgr.intermediate_conv_states.shape[2] == 32 - assert mgr._kv_reserve_draft_tokens == 31 - assert mgr._num_reserved_dummy_slots == 1 - assert not mgr.use_replay_state_update - finally: - mgr.shutdown() + return mgr @skip_no_cuda -@pytest.mark.parametrize( - "builder", - [_build_hybrid_with_mamba_layer, _build_v2_hybrid_with_mamba_layer], - ids=["cpp", "v2"], -) -def test_hybrid_replay_buffers_size_by_tokens_per_gen_step(builder): - spec_config = _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5) - mgr = builder( - max_batch_size=4, - spec_config=spec_config, - use_replay_state_update=True, +def test_cpp_hybrid_zero_local_mamba_layers(): + """End-to-end: real parent KVCacheManager + real early-exit. Verifies + early-exit invariants on the manager state and that Mamba hooks and + metadata preparation do not touch uninitialized Mamba-only state.""" + mgr = _build_zero_mamba_hybrid( + enable_block_reuse=True, + periodic_snapshot_interval=64, ) - try: - replay_metadata = mgr.get_replay_state_update_metadata() - assert mgr.use_replay_state_update is True - assert replay_metadata is not None - assert replay_metadata.replay_step_width == spec_config.tokens_per_gen_step - assert replay_metadata.replay_history_size == max( - MIN_REPLAY_HISTORY_SIZE, spec_config.tokens_per_gen_step - ) - layer_cache = mgr.mamba_layer_cache(0) - _assert_replay_layer_cache_uses_history_size( - layer_cache, replay_metadata.replay_history_size - ) - finally: - mgr.shutdown() + # Early-exit indicators. + assert mgr.local_num_mamba_layers == 0 + assert mgr.mamba_pp_layers == [] + assert mgr.requests == [] + assert mgr.pp_layers == [2, 3] -@skip_no_cuda -def test_v2_gdn_replay_builds_affine_state_layout(): - 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", + # Parent KVCacheManager was really initialized. self.impl is the C++ + # KVCacheManagerCpp object; blocks_per_window is set up by it. + assert hasattr(mgr, "impl") + assert hasattr(mgr, "blocks_per_window") + # Parent saw num_layers = num_layers (4), not mamba_num_layers + num_layers. + # On the early-exit branch, num_layers is forwarded as-is. + assert mgr.num_layers == 4 + assert mgr.num_local_layers == 2 + assert mgr.max_num_tokens == 96 + assert all( + config.window_size != LinearCacheType.RECURRENT_STATES.value + for config in mgr.pool_configurations + ) + assert all( + config.window_size != LinearCacheType.RECURRENT_STATES.value + for config in mgr.impl.pool_configurations ) - 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 state_strides is not None - states = mgr.all_ssm_states - first_state = states[0] - expected_layer_stride = ( - states[1].data_ptr() - first_state.data_ptr() - ) // first_state.element_size() - assert state_strides == ( - mgr.local_num_mamba_layers, - expected_layer_stride, - first_state.stride(0), - ) - finally: - mgr.shutdown() + # No mamba-only state was allocated. + for attr in ( + "ssm_state_shape", + "conv_state_shape", + "mamba_layer_offsets", + "cuda_state_indices", + "host_block_offsets", + "recurrent_states_pool_index", + ): + assert not hasattr(mgr, attr), f"{attr} must not be set on the zero-mamba early-exit path" + # Parent must not have been told to treat this as linear attention. + assert mgr.is_linear_attention is False -@skip_no_cuda -def test_v2_gdn_replay_all_layer_commit_matches_contiguous_layout(): - from tensorrt_llm._torch.modules.fla.cached_replay import ( - commit_gdn_cached_replay_history_layers, - ) + # The scheduler hook is still advertised on ranks without local Mamba + # layers, so its inputs must be initialized with the same interval as + # Mamba-owning PP ranks to keep their scheduling decisions aligned. + assert mgr.kv_cache_config.enable_block_reuse is True + assert mgr.linear_attention_metadata.states_snapshot_interval == 64 + request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[]) + mgr.prepare_expect_snapshot_points([request]) + assert request.expect_snapshot_points == [64, 128] + # The shared interval must not make this attention-only rank consult its + # nonexistent recurrent-state pool and report zero KV capacity. + attention_capacity = KVCacheManager.get_num_available_tokens(mgr, 128) + assert attention_capacity > 0 + assert mgr.get_num_available_tokens(128) == attention_capacity - batch_size = 16 - 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", - ) - 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) + # Guards on the three mamba-only methods must turn them into no-ops + # instead of crashing on the missing state above. + empty_batch = ScheduledRequests() + mgr.prepare_resources(empty_batch) # super() runs, then guard returns + mgr.update_mamba_states(attn_metadata=None, num_accepted_tokens=None, state_indices=None) + mgr._setup_state_indices() - 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.zeros_like(positions), + metadata = Mamba2Metadata(max_batch_size=1, chunk_size=8) + seq_lens = torch.tensor([8], dtype=torch.int32) + metadata.prepare( + SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=1, + num_ctx_tokens=8, + kv_cache_manager=mgr, + request_ids=[123], + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int32) ), - dim=1, - ) - n_writes = torch.tensor([batch_size], device="cuda", dtype=torch.int32) - 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, - replay_work_items=replay_work_items, - n_writes=n_writes, - history_size=mgr.replay_history_size, ) + ) + assert metadata.state_indices[0].item() == 0 - mgr._commit_gdn_cached_replay_history_layers( - SimpleNamespace( - mamba_metadata=SimpleNamespace( - replay_num_decodes=batch_size, - replay_work_items=replay_work_items, - replay_n_writes=n_writes, - ) - ), - batch_size, - ) - torch.testing.assert_close(torch.stack(mgr.all_ssm_states), expected, rtol=0, atol=0) - finally: - mgr.shutdown() +@skip_no_cuda +def test_v2_stack_state_views_detects_affine_layout(): + pool = torch.zeros(3, 4, 8, device="cuda") + layers = [pool[i] for i in range(3)] + stacked = _stack_state_views(layers) -def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): - mgr = object.__new__(MambaHybridCacheManagerV2) - 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) - mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(batch_size, dtype=torch.bool)) - mgr.all_conv_states = [torch.empty(0)] - mgr._stacked_conv_states = None - mgr.intermediate_conv_states = torch.empty(0) + assert stacked.shape == (3, 4, 8) + stacked[2, 1, 3] = 5.0 + assert layers[2][1, 3] == 5.0 - events = [] - mgr._commit_gdn_cached_replay_history_layers = lambda *_args, **_kwargs: events.append("commit") - monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._advance_replay_state", - lambda *_args, **_kwargs: events.append("advance"), - ) - monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", - lambda *_args, **_kwargs: None, - ) - mgr.update_mamba_states( - SimpleNamespace(num_seqs=batch_size, num_contexts=0), - torch.full((batch_size,), 3, dtype=torch.int32), - state_indices=torch.arange(batch_size, dtype=torch.int32), - ) +def test_v2_stack_state_views_rejects_non_affine_layout(): + pool = torch.zeros(4, 4, 8) + stack = _stack_state_views - assert events == ["commit", "advance"] + assert stack([pool[0], pool[1], pool[3]]) is None + assert stack([pool[0], torch.zeros(5, 8)]) is None + assert stack([]) is None -def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr.local_num_mamba_layers = 1 - mgr._request_id_to_state_index = { - 100: 0, - 101: 1, - 102: 2, - 103: 3, - } - mgr._request_id_to_is_dummy = { - 100: False, - 101: False, - 102: True, - 103: False, - } - mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.bool)) - mgr._dummy_request_mask_host = torch.zeros(4, dtype=torch.bool) - mgr._use_replay_state_update = True - mgr.replay_step_width = 5 - mgr.replay_history_size = 16 - mgr.prev_num_accepted_tokens = torch.full((4,), 13, dtype=torch.int32) - mgr.cache_buf_idx = torch.ones(4, dtype=torch.int32) - mgr.intermediate_state_indices = torch.arange(4, dtype=torch.int32) - mgr.all_ssm_states = [] - mgr.all_conv_states = [torch.empty(0)] - mgr._stacked_conv_states = None - mgr.intermediate_conv_states = torch.empty(0) - monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", - lambda *args, **kwargs: None, +@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=skip_no_cuda)]) +def test_v2_kda_admission_slot_updates_without_sync(device): + """Disagg GEN admission resets, relocates and seeds KDA replay slots; on CUDA none of it may + synchronize the stream (the previous step is still queued).""" + num_slots, num_layers, num_spec, dim, window = 8, 2, 2, 2, 3 + mgr = object.__new__(KimiK3HybridCacheManager) + policy = KDAReplayState(num_spec) + mgr._kda_replay = policy + mgr._speculative_state = policy + mgr.local_num_mamba_layers = num_layers + policy.prev_num_accepted_tokens = torch.arange( + 1, num_slots + 1, dtype=torch.int32, device=device ) - - request_ids = [100, 101, 102, 103] - state_indices = torch.tensor( - mgr.get_state_indices(request_ids, [False, False, False, True]), - dtype=torch.int32, + # Stale device indices: the reset must take its slots from _host_state_indices. + mgr._state_indices_buffer = DeviceCopyBuffer( + torch.full((num_slots,), 7, dtype=torch.int32, device=device) ) - assert mgr._dummy_request_mask.tolist() == [False, False, True, True] - - mgr.update_mamba_states( - SimpleNamespace(num_seqs=4, num_contexts=1), - torch.tensor([1, 3, 3, 3], dtype=torch.int32), - state_indices=state_indices, + mgr._host_state_indices = torch.tensor([5, 1, 0, 0, 0, 0, 0, 0], dtype=torch.int32) + mgr._request_id_to_state_index = {10: 4, 11: 2} + mgr.conv_state_shape = [3 * dim, window] + mgr.conv_section_dims = [dim] * 3 + mgr.all_conv_states = [ + torch.randn(num_slots, 3 * dim, window, device=device) for _ in range(num_layers) + ] + policy.kda_conv_q, policy.kda_conv_k, policy.kda_conv_v = ( + torch.randn(num_layers, num_slots, dim, window + num_spec, device=device) for _ in range(3) + ) + policy.kda_k_cache, policy.kda_g_cache, policy.kda_v_cache, policy.kda_beta_cache = ( + torch.randn(num_layers, num_slots, num_spec, dim, device=device) for _ in range(4) ) - assert mgr.prev_num_accepted_tokens.tolist() == [13, 3, 13, 13] - assert mgr.cache_buf_idx.tolist() == [1, 0, 1, 1] - + policy._conv_states = tuple(mgr.all_conv_states) + policy._conv_section_dims = tuple(mgr.conv_section_dims) + policy._committed_window = window -def test_v2_hybrid_dynamic_tree_promotes_accepted_leaf_state(monkeypatch): - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr.local_num_mamba_layers = 1 - mgr._use_replay_state_update = False - mgr.intermediate_state_indices = torch.arange(2, dtype=torch.int32) - mgr.all_ssm_states = [torch.empty(0)] - mgr.all_conv_states = [torch.empty(0)] - mgr._stacked_ssm_states = None - mgr._stacked_conv_states = None - mgr.intermediate_ssm_states = torch.empty((1, 2, 8)) - mgr.intermediate_conv_states = torch.empty((1, 2, 8)) + if device == "cuda": + torch.cuda.synchronize() + previous_mode = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + mgr._reset_context_mamba_slots(2) + policy.reset_slots(torch.tensor([3]), [3]) + policy.relocate_slots([1], [6]) + mgr.seed_kda_replay_caches_for_disagg_gen([10, 11]) + finally: + if device == "cuda": + torch.cuda.set_sync_debug_mode(previous_mode) - promoted_positions = [] + # Reset 5, 1 and 3; slot 6 takes slot 1's (reset) count; seeding resets 4 and 2; 7 untouched. + assert policy.prev_num_accepted_tokens.tolist() == [1, 0, 0, 0, 0, 0, 0, 8] - def capture_promoted_position(_dst, _src, _src_indices, positions, _dst_indices): - promoted_positions.append(positions.clone()) +@pytest.mark.parametrize("backbone_asks", [True, False]) +def test_kimi_explicit_v2_manager_kda_draft_records_follow_the_model( + monkeypatch: pytest.MonkeyPatch, + backbone_asks: bool, +) -> None: + """The factory asks the V2 manager for per-token KDA verify states only when + the model's backbone sets ``kda_draft_records``.""" monkeypatch.setattr( - "tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager._promote_mamba_state_triton", - capture_promoted_position, + "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", + lambda: True, ) - - mgr.update_mamba_states( - SimpleNamespace(num_seqs=3, num_contexts=1), - torch.tensor([1, 2, 3], dtype=torch.int32), - state_indices=torch.tensor([9, 10, 11], dtype=torch.int32), - accepted_leaf_positions=torch.tensor([4, 7], dtype=torch.int64), + backbone = SimpleNamespace(kda_draft_records=True) if backbone_asks else SimpleNamespace() + model_engine = SimpleNamespace( + model=SimpleNamespace(model=backbone), + _max_cuda_graph_batch_size=4, + is_multimodal=True, ) - assert len(promoted_positions) == 2 - assert all( - torch.equal(positions, torch.tensor([4, 7], dtype=torch.int32)) - for positions in promoted_positions + _, kwargs = _capture_kimi_v2_manager_ctor( + monkeypatch, MTPDecodingConfig(max_draft_len=3), model_engine ) + assert "kda_replay_num_spec" not in kwargs + assert kwargs.get("kda_draft_records", False) is backbone_asks + @skip_no_cuda -@pytest.mark.parametrize( - "mamba_ssm_cache_dtype", - [torch.float16, torch.float32, torch.bfloat16], -) -def test_cpp_hybrid_passes_per_window_pool_dtypes_for_nvfp4_kv_cache( - mamba_ssm_cache_dtype, -): - mgr = _build_hybrid_with_mamba_layer( - dtype=DataType.NVFP4, - mamba_ssm_cache_dtype=mamba_ssm_cache_dtype, +@pytest.mark.parametrize("kda_replay_num_spec", [2, None]) +def test_v2_kda_draft_records_replace_the_replay_caches(kda_replay_num_spec): + """On the KDA replay path, ``kda_draft_records`` allocates the fp32 records of every draft of each slot (vn, + beta * k, decay), per layer, in place of the per-draft replay caches; the conv caches and the replay path stay. + Without the replay path it allocates nothing.""" + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + num_mamba_layers=2, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_d_conv=5, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=kda_replay_num_spec, + kda_draft_records=True, ) - recurrent_pool_dtype = torch_dtype_to_binding(mamba_ssm_cache_dtype) - - expected_dtypes = [ - (LinearCacheType.RECURRENT_STATES.value, recurrent_pool_dtype), - (128, DataType.NVFP4), - ] - assert [ - (config.window_size, config.dtype) for config in mgr.pool_configurations - ] == expected_dtypes - assert [ - (config.window_size, config.dtype) for config in mgr.impl.pool_configurations - ] == expected_dtypes - assert mgr._layer_to_pool_idx == {0: 0, 1: 1} - assert mgr.recurrent_states_pool_index == 0 - assert mgr.impl.get_recurrent_states_pool().dtype == mamba_ssm_cache_dtype - - compact_scale_pointers = mgr.impl.get_block_scale_pool_pointers() - assert mgr.impl.get_block_pool_pointers().shape == (2, 2) - assert compact_scale_pointers.shape == (1, 2) - assert mgr.kv_cache_pool_pointers.shape == (2, 2, 2) - assert torch.count_nonzero(mgr.kv_cache_pool_pointers[0, :, 1]) == 0 - assert torch.equal(mgr.kv_cache_pool_pointers[1, :, 1], compact_scale_pointers[0]) + try: + if kda_replay_num_spec is None: + assert getattr(mgr, "kda_state_tok", None) is None + assert not mgr.keeps_kda_draft_records + return + assert mgr.keeps_kda_draft_records + assert set(mgr._kda_replay_layout()) == { + "kda_conv_q", + "kda_conv_k", + "kda_conv_v", + "kda_state_tok", + } + # A KVCM state buffer of the SSM layers, kept out of disaggregated transfers like the other replay buffers. + assert DataRole("kda_state_tok") in mgr.get_disagg_ignored_roles() + assert len(mgr.kda_state_tok) == 2 + for layer_idx in range(2): + layer_cache = mgr.mamba_layer_cache(layer_idx) + cache_size = layer_cache.temporal.shape[0] + assert layer_cache.has_kda_replay_caches + assert layer_cache.kda_conv_q.shape == (cache_size, 48, 6) + assert layer_cache.kda_k_cache is None + assert layer_cache.kda_g_cache is None + assert layer_cache.kda_v_cache is None + assert layer_cache.kda_beta_cache is None + states = layer_cache.kda_state_tok + num_heads, _, head_dim = layer_cache.temporal.shape[1:] + assert states.shape == (cache_size, 3, 2, num_heads, head_dim) + assert states.dtype is torch.float32 + assert states.data_ptr() == mgr.kda_state_tok[layer_idx].data_ptr() + # The records' replay addresses a slot by its stride and reads its records contiguously. + assert states[0].is_contiguous() + finally: + mgr.shutdown() + # Released with the other replay buffers. + assert mgr.kda_state_tok is None @skip_no_cuda -def test_cpp_hybrid_merges_compact_scale_rows_with_unmanaged_layers(): - mgr = _build_hybrid_with_mamba_layer( - dtype=DataType.NVFP4, - mamba_layer_mask=[True, False, True, False], - attention_layer_mask=[False, False, False, True], +def test_v2_kda_draft_records_count_in_the_per_slot_budget(): + """The capacity math sees the drafts' records in place of the per-draft replay caches: per draft and layer, a slot + costs vn, beta * k and the decay, fp32 [H, K] each, instead of the k / g / v / beta entries.""" + kwargs = dict( + num_mamba_layers=2, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_d_conv=5, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, ) - - assert mgr.pp_layers == [0, 2, 3] - recurrent_window = LinearCacheType.RECURRENT_STATES.value - assert mgr.max_attention_window_vec == [ - recurrent_window, - recurrent_window, - mgr.max_seq_len, - ] - assert mgr.kv_cache_pool_mapping[:, 0].tolist() == [0, 0, 1] - compact_scale_pointers = mgr.impl.get_block_scale_pool_pointers() - assert compact_scale_pointers.shape == (1, 2) - assert mgr.kv_cache_pool_pointers.shape == (2, 2, 2) - assert torch.count_nonzero(mgr.kv_cache_pool_pointers[0, :, 1]) == 0 - assert torch.equal(mgr.kv_cache_pool_pointers[1, :, 1], compact_scale_pointers[0]) + plain = _build_v2_hybrid_with_mamba_layer(**kwargs) + try: + plain_bytes = plain._mamba_state_bytes_per_slot() + finally: + plain.shutdown() + with_states = _build_v2_hybrid_with_mamba_layer(kda_draft_records=True, **kwargs) + try: + num_heads, _, head_dim = with_states.ssm_state_shape + per_token = 2 * 2 * 3 * num_heads * head_dim * 4 + replay_caches = 2 * sum( + dtype.itemsize * math.prod(shape) + for name, (dtype, shape) in _kda_replay_layout(2, num_heads, head_dim, 4).items() + if name in _KDA_REPLAY_BUFFERS[3:] + ) + assert with_states._mamba_state_bytes_per_slot() == plain_bytes - replay_caches + per_token + finally: + with_states.shutdown() @skip_no_cuda -def test_cpp_hybrid_recurrent_pool_reserves_cuda_graph_padding_slot(): - """Without spec decoding, the recurrent-state snapshot pool must - have at least max_batch_size + 1 slots — one extra for the - CUDA-graph padding sentinel (CUDA_GRAPH_DUMMY_REQUEST_ID). Without - it, the padding sentinel evicts live recurrent state under load.""" - max_batch_size = 4 - mgr = _build_hybrid_with_mamba_layer(spec_config=None, max_batch_size=max_batch_size) - recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] - assert recurrent_primary >= max_batch_size + 1, ( - f"recurrent-state pool has {recurrent_primary} slots, " - f"need >= max_batch_size + 1 = {max_batch_size + 1} to host the " - f"CUDA-graph padding sentinel without evicting live state" - ) - +def test_v2_state_index_setup_skips_unchanged_uploads(monkeypatch): + """A step whose state indices and dummy flags did not change enqueues no upload of either: the device buffers + keep values the host never sent. A changed step uploads both again.""" + # The test writes the buffers behind _setup_state_indices to see the skip, which the check of skipped copies + # reports. + monkeypatch.setattr("tensorrt_llm._device_copy._verify_skipped_device_copies", False) -@skip_no_cuda -def test_cpp_hybrid_recurrent_pool_reserves_draft_len_sentinel_slots(): - """With spec decoding, CUDAGraphRunner._get_padded_batch issues a - distinct dummy request id for each runtime_draft_len in - [0, max_draft_len], so the recurrent-state snapshot pool must reserve - one slot per draft length on top of the CUDA-graph padding slot.""" - max_batch_size, max_draft_len = 4, 2 - spec_config = MTPDecodingConfig(max_draft_len=max_draft_len) - mgr = _build_hybrid_with_mamba_layer(spec_config=spec_config, max_batch_size=max_batch_size) - recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] - expected_min = max_batch_size + 1 + max_draft_len - assert recurrent_primary >= expected_min, ( - f"recurrent-state pool has {recurrent_primary} slots, " - f"need >= max_batch_size + 1 + max_draft_len = {expected_min} so " - f"per-draft-len sentinels don't collide with live state" - ) + class StateCache: + def __init__(self, slot: int) -> None: + self.slot = slot + def get_ssm_block_base_index(self, _layer_group_id) -> int: + return self.slot -def _build_hybrid_with_mamba_layer_pp( - spec_config=None, max_batch_size=4, enable_block_reuse=False, pp_size=2 -): - """Same as ``_build_hybrid_with_mamba_layer`` but with ``pp_size`` >= 1. + def close(self) -> None: + pass - Uses ``world_size = pp_size`` and ``rank = 0`` so the real C++ KVCacheManager - still goes through its single-process path while the Python pool-sizing - code sees ``mapping.pp_size > 1``. Constructs ``pp_size * 2`` total layers - (alternating mamba/attn) so that each PP slice has both a mamba and an - attention layer — otherwise some ranks would hit a slope=0 edge case in - the affine memory model when block reuse is disabled. - """ - pairs = pp_size # one (mamba, attn) pair per PP stage so every rank has both kinds - mamba_mask = [True, False] * pairs - attn_mask = [False, True] * pairs - mamba_num_layers = sum(mamba_mask) - num_layers = sum(attn_mask) - mapping = Mapping(world_size=pp_size, rank=0, tp_size=1, pp_size=pp_size) - kv_cache_config = KvCacheConfig(max_tokens=512, enable_block_reuse=enable_block_reuse) - return CppMambaHybridCacheManager( - mamba_d_state=8, - mamba_d_conv=4, - mamba_num_heads=4, - mamba_n_groups=1, - mamba_head_dim=8, - mamba_num_layers=mamba_num_layers, - mamba_layer_mask=mamba_mask, - mamba_cache_dtype=torch.float16, - mamba_ssm_cache_dtype=torch.float16, - kv_cache_config=kv_cache_config, - kv_cache_type=CacheTypeCpp.SELF, - num_layers=num_layers, - num_kv_heads=4, - head_dim=64, - tokens_per_block=32, - max_seq_len=128, - max_batch_size=max_batch_size, - mapping=mapping, - spec_config=spec_config, - layer_mask=attn_mask, + 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, ) + live = SimpleNamespace(py_request_id=101, is_dummy=False) + padding = SimpleNamespace(py_request_id=202, is_dummy=True) + mgr.kv_cache_map.update({101: StateCache(0), 202: StateCache(1)}) + try: + mgr._setup_state_indices([live, padding]) + torch.cuda.synchronize() + assert mgr.cuda_state_indices[:2].tolist() == [0, 1] + assert mgr._dummy_request_mask[:2].tolist() == [False, True] + mgr.cuda_state_indices.fill_(-1) + mgr._dummy_request_mask.fill_(True) + mgr._setup_state_indices([live, padding]) + torch.cuda.synchronize() + assert mgr.cuda_state_indices[:2].tolist() == [-1, -1] + assert mgr._dummy_request_mask[:2].tolist() == [True, True] -# Skip when running under pytest --run-ray (sets TLLM_DISABLE_MPI=1). In that -# mode ``Mapping`` resolves to ``DeviceMeshTopology``, whose ``pp_rank`` -# requires torch.distributed initialisation that isn't available in this -# single-process unit test. The pp-sharding behaviour exercised here is -# orthogonal to the Ray orchestrator. -_skip_under_ray = pytest.mark.skipif( - os.environ.get("TLLM_DISABLE_MPI") == "1", - reason="pp_size>1 helper builds Mapping with world_size>1 which needs " - "torch.distributed under TLLM_DISABLE_MPI=1 (Ray) sessions", -) + mgr.kv_cache_map[202].slot = 2 + padding.is_dummy = False + mgr._setup_state_indices([live, padding]) + torch.cuda.synchronize() + assert mgr.cuda_state_indices[:2].tolist() == [0, 2] + assert mgr._dummy_request_mask[:2].tolist() == [False, False] + finally: + mgr.shutdown() -@_skip_under_ray @skip_no_cuda -@pytest.mark.parametrize("pp_size", [2, 4]) -def test_cpp_hybrid_recurrent_pool_scales_with_pp_size(pp_size): - """With pipeline parallelism, multiple microbatches are in-flight on the - same rank concurrently, each holding up to ``max_batch_size`` sequences' - Mamba state. The recurrent-state pool must therefore size for - ``max_batch_size * pp_size`` live slots (plus the CUDA-graph padding - sentinel). Without this scaling, the first inference batch under PP>1 - trips ``No free block found`` once requests beyond the first microbatch - enter the pool (cf. TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2]). - """ - max_batch_size = 4 - mgr = _build_hybrid_with_mamba_layer_pp( - spec_config=None, max_batch_size=max_batch_size, pp_size=pp_size - ) - recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] - expected_min = max_batch_size * pp_size + 1 - assert recurrent_primary >= expected_min, ( - f"recurrent-state pool has {recurrent_primary} slots with pp_size={pp_size}, " - f"need >= max_batch_size * pp_size + 1 = {expected_min} so concurrent " - f"in-flight microbatches don't exhaust live-state slots" +def test_v2_kda_replay_allocates_logical_slot_caches(): + spec_config = MTPDecodingConfig(max_draft_len=2) + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + num_mamba_layers=2, + spec_config=spec_config, + conv_state_layout="q_k_v", + mamba_d_conv=5, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, ) + try: + assert mgr.intermediate_ssm_states is None + assert mgr.intermediate_conv_states is None + assert mgr.intermediate_state_indices.shape == (4,) + assert not hasattr(mgr, "get_replay_state_update_metadata") + layer_cache = mgr.mamba_layer_cache(0) + 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) + assert layer_cache.kda_conv_v.shape == (cache_size, 48, 6) + assert layer_cache.kda_conv_q.dtype is torch.bfloat16 + assert layer_cache.kda_conv_q.stride(-2) == 1 + assert layer_cache.kda_k_cache.shape == (cache_size, 2, 48) + assert layer_cache.kda_g_cache.shape == (cache_size, 2, 48) + assert layer_cache.kda_v_cache.shape == (cache_size, 2, 48) + # Six bf16 heads are 12 bytes, so the row is padded to 16 (stride 8) + # to keep every nested per-draft view 16-byte aligned for CuTe. + assert layer_cache.kda_beta_cache.shape == (cache_size, 2, 6) + assert layer_cache.kda_beta_cache.stride(-2) == 8 + for layer_idx in range(2): + for slot_idx in range(cache_size): + for draft_idx in range(2): + assert ( + mgr.mamba_layer_cache(layer_idx) + .kda_beta_cache[slot_idx, draft_idx] + .data_ptr() + % _KDA_BETA_CACHE_ALIGNMENT_BYTES + == 0 + ) + assert layer_cache.kda_k_cache.dtype is torch.bfloat16 + assert layer_cache.kda_g_cache.dtype is torch.float32 + assert layer_cache.kda_v_cache.dtype is torch.bfloat16 + assert layer_cache.kda_beta_cache.dtype is torch.bfloat16 + assert layer_cache.prev_num_accepted_tokens.dtype == torch.int32 + assert layer_cache.prev_num_accepted_tokens.shape == (cache_size,) + assert layer_cache.prev_num_accepted_tokens is ( + mgr.mamba_layer_cache(1).prev_num_accepted_tokens + ) + assert not hasattr(layer_cache, "intermediate_ssm") + assert not hasattr(layer_cache, "intermediate_conv_window") + assert layer_cache.kda_state_tok is None + assert not mgr.keeps_kda_draft_records + assert layer_cache.has_kda_replay_caches + finally: + mgr.shutdown() -@skip_no_cuda -def test_cpp_hybrid_recurrent_pool_floor_with_block_reuse(): - """With block reuse enabled, the block-reuse branch must not drop the - live-state + CUDA-graph-padding floor. - With max_batch_size=4, periodic_snapshot_interval=256, max_tokens=512: - naive: max_snapshots = 512 // 256 = 2 (drops live-state floor!) - fixed: max_snapshots = max(2, 4 + 1) = 5 - """ - max_batch_size = 4 - mgr = _build_hybrid_with_mamba_layer( - spec_config=None, - max_batch_size=max_batch_size, - enable_block_reuse=True, - periodic_snapshot_interval=256, - ) - recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] - assert recurrent_primary >= max_batch_size + 1, ( - f"recurrent-state pool has {recurrent_primary} slots with block reuse enabled, " - f"need >= max_batch_size + 1 = {max_batch_size + 1} to prevent the padding " - f"sentinel from evicting live recurrent state" +def _kda_draft_record_build_config(quota, initial_pool_ratio=None): + return KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + initial_pool_ratio=initial_pool_ratio, ) -@skip_no_cuda -def test_cpp_hybrid_dry_run_recurrent_pool_additive_with_block_reuse(): - """Dry-run path (is_estimating_kv_cache=True) under block reuse must - keep the live-state floor *plus* room for snapshots, not collapse to - max(snapshots, live). With max_batch_size=4, interval=256, max_tokens=512: - old: max_snapshots = max(512//256, 4) = 4 (no headroom for snapshots) - new: max_snapshots = 4 + 512//256 = 6 (live + snapshots) - """ - max_batch_size = 4 - mgr = _build_hybrid_with_mamba_layer( - spec_config=None, - max_batch_size=max_batch_size, - enable_block_reuse=True, - periodic_snapshot_interval=256, - is_estimating_kv_cache=True, +@pytest.mark.parametrize("pool_sizing", ["typical_step", "block_reuse", "pool_ratio"]) +def test_v2_kda_draft_records_size_with_the_ssm_pool(pool_sizing): + """The drafts' records are a buffer of the SSM layer in each recurrent-state slot, in place of the k / g / v / + beta caches. However the SSM pool is sized (by the typical step, with block reuse, or by a pool_ratio), every slot + holds them: the build leaves the whole GPU quota to the pools, and the quotas count the records per state slot.""" + enable_block_reuse = pool_sizing == "block_reuse" + mgr = _v2_kda_draft_record_manager( + kda_draft_records=True, enable_block_reuse=enable_block_reuse ) - recurrent_primary, _ = mgr.blocks_per_window[LinearCacheType.RECURRENT_STATES.value] - # 4 live state slots + 2 reuse snapshots = 6. - expected_min = max_batch_size + (512 // 256) - assert recurrent_primary >= expected_min, ( - f"dry-run recurrent-state pool has {recurrent_primary} slots, " - f"need >= live_state + reuse_snapshots = {expected_min}; the old " - f"max(reuse, live) formula dropped reuse headroom" + plain = _v2_kda_draft_record_manager( + kda_draft_records=False, enable_block_reuse=enable_block_reuse ) - - -# --------------------------------------------------------------------------- -# CppMambaHybridCacheManager: rank with zero local mamba layers -# -# Regression test for the early-exit path added when a rank ends up with no -# mamba layers (e.g. under PP sharding when all mamba layers fall on other -# ranks). On that path, the constructor must: -# - call the real parent KVCacheManager with the union layer_mask and -# num_layers=num_layers (not mamba_num_layers + num_layers), -# - skip allocating any mamba-only state, and -# - leave self.requests = [] so Mamba-only hooks and metadata preparation -# can no-op without touching uninitialized state. -# -# We exercise the same Python branch with world_size=1 (so the real C++ -# KVCacheManager init doesn't need MPI) and a layer mask that contains zero -# mamba layers. -# --------------------------------------------------------------------------- - - -def _build_zero_mamba_hybrid( - enable_block_reuse=False, - periodic_snapshot_interval=256, -): - """Construct a real CppMambaHybridCacheManager whose this-rank slice has - no mamba layers. world_size=1 / pp_size=1 keeps the real parent - KVCacheManager off the MPI path.""" - # [other, other, full_attn, full_attn] - mamba_mask = [False, False, False, False] - attn_mask = [False, False, True, True] - mamba_num_layers = sum(mamba_mask) # 0 - num_layers = sum(attn_mask) # 4 - - mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) - # Cap KV pool size so the real C++ allocator only takes a tiny slice of - # GPU memory; we don't actually use the cache. - kv_cache_config = KvCacheConfig( - max_tokens=128, - enable_block_reuse=enable_block_reuse, - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=periodic_snapshot_interval), + records = 3 * 2 * 8 * 256 * 4 + replay_caches = sum( + dtype.itemsize * math.prod(shape) + for name, (dtype, shape) in _kda_replay_layout(2, 8, 256, 3).items() + if name in _KDA_REPLAY_BUFFERS[3:] ) - - mgr = CppMambaHybridCacheManager( - # mamba cache parameters — values are unused on the early-exit path - # but must be type-valid. - mamba_d_state=8, - mamba_d_conv=4, - mamba_num_heads=4, - mamba_n_groups=1, - mamba_head_dim=8, - mamba_num_layers=mamba_num_layers, - mamba_layer_mask=mamba_mask, - mamba_cache_dtype=torch.float16, - mamba_ssm_cache_dtype=torch.float16, - # kv cache parameters - kv_cache_config=kv_cache_config, - kv_cache_type=CacheTypeCpp.SELF, - num_layers=num_layers, - num_kv_heads=4, - head_dim=64, - tokens_per_block=32, - max_seq_len=128, - max_num_tokens=96, - max_batch_size=2, - mapping=mapping, - spec_config=None, - layer_mask=attn_mask, + slot_delta = records - replay_caches + assert mgr._mamba_state_bytes_per_slot() - plain._mamba_state_bytes_per_slot() == slot_delta + live_slots = mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots + assert ( + mgr._minimum_live_gpu_quota() - plain._minimum_live_gpu_quota() == live_slots * slot_delta ) - return mgr + quota = 128 << 20 + pool_ratio = [0.5, 0.5] if pool_sizing == "pool_ratio" else None + built = mgr._build_cache_config(_kda_draft_record_build_config(quota, pool_ratio)) + assert built.cache_tiers[0].quota == quota + ssm_buffers = {str(buffer.role): buffer.size for buffer in built.layers[0].buffers} + assert ssm_buffers["kda_state_tok"] == records + assert not set(ssm_buffers) & set(_KDA_REPLAY_BUFFERS[3:]) + with pytest.raises(ValueError, match="too small for live recurrent states"): + mgr._build_cache_config(_kda_draft_record_build_config(mgr._minimum_live_gpu_quota() - 1)) -@skip_no_cuda -def test_cpp_hybrid_zero_local_mamba_layers(): - """End-to-end: real parent KVCacheManager + real early-exit. Verifies - early-exit invariants on the manager state and that Mamba hooks and - metadata preparation do not touch uninitialized Mamba-only state.""" - mgr = _build_zero_mamba_hybrid( - enable_block_reuse=True, - periodic_snapshot_interval=64, - ) - - # Early-exit indicators. - assert mgr.local_num_mamba_layers == 0 - assert mgr.mamba_pp_layers == [] - assert mgr.requests == [] - assert mgr.pp_layers == [2, 3] - # Parent KVCacheManager was really initialized. self.impl is the C++ - # KVCacheManagerCpp object; blocks_per_window is set up by it. - assert hasattr(mgr, "impl") - assert hasattr(mgr, "blocks_per_window") - # Parent saw num_layers = num_layers (4), not mamba_num_layers + num_layers. - # On the early-exit branch, num_layers is forwarded as-is. - assert mgr.num_layers == 4 - assert mgr.num_local_layers == 2 - assert mgr.max_num_tokens == 96 - assert all( - config.window_size != LinearCacheType.RECURRENT_STATES.value - for config in mgr.pool_configurations - ) - assert all( - config.window_size != LinearCacheType.RECURRENT_STATES.value - for config in mgr.impl.pool_configurations +@skip_no_cuda +@pytest.mark.parametrize("enable_block_reuse", [False, True]) +def test_v2_kda_draft_records_leave_the_quota_to_the_pools(enable_block_reuse): + """With the drafts' records inside the SSM pool, the runtime pools take the whole GPU quota (rounded up to whole + 2 MiB grains) with or without block reuse: nothing is set aside beside them.""" + mgr = _v2_kda_draft_record_manager( + kda_draft_records=True, enable_block_reuse=enable_block_reuse ) - - # No mamba-only state was allocated. - for attr in ( - "ssm_state_shape", - "conv_state_shape", - "mamba_layer_offsets", - "cuda_state_indices", - "host_block_offsets", - "recurrent_states_pool_index", - ): - assert not hasattr(mgr, attr), f"{attr} must not be set on the zero-mamba early-exit path" - # Parent must not have been told to treat this as linear attention. - assert mgr.is_linear_attention is False - - # The scheduler hook is still advertised on ranks without local Mamba - # layers, so its inputs must be initialized with the same interval as - # Mamba-owning PP ranks to keep their scheduling decisions aligned. - assert mgr.kv_cache_config.enable_block_reuse is True - assert mgr.linear_attention_metadata.states_snapshot_interval == 64 - request = SimpleNamespace(prompt_len=150, expect_snapshot_points=[]) - mgr.prepare_expect_snapshot_points([request]) - assert request.expect_snapshot_points == [64, 128] - # The shared interval must not make this attention-only rank consult its - # nonexistent recurrent-state pool and report zero KV capacity. - attention_capacity = KVCacheManager.get_num_available_tokens(mgr, 128) - assert attention_capacity > 0 - assert mgr.get_num_available_tokens(128) == attention_capacity - - # Guards on the three mamba-only methods must turn them into no-ops - # instead of crashing on the missing state above. - empty_batch = ScheduledRequests() - mgr.prepare_resources(empty_batch) # super() runs, then guard returns - mgr.update_mamba_states(attn_metadata=None, num_accepted_tokens=None, state_indices=None) - mgr._setup_state_indices() - - metadata = Mamba2Metadata(max_batch_size=1, chunk_size=8) - seq_lens = torch.tensor([8], dtype=torch.int32) - metadata.prepare( - SimpleNamespace( - seq_lens=seq_lens, - seq_lens_cuda=seq_lens.cuda(), - num_contexts=1, - num_ctx_tokens=8, - kv_cache_manager=mgr, - request_ids=[123], - kv_cache_params=SimpleNamespace( - num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int32) - ), - ) + runtime_manager = RuntimeKVCacheManager( + mgr._build_cache_config(_kda_draft_record_build_config(128 << 20)) ) - assert metadata.state_indices[0].item() == 0 + try: + pool_bytes = runtime_manager.get_quota(GPU_LEVEL) + finally: + runtime_manager.shutdown() + assert 128 << 20 <= pool_bytes < (128 << 20) + (2 << 20) @skip_no_cuda -def test_v2_stack_state_views_detects_affine_layout(): - pool = torch.zeros(3, 4, 8, device="cuda") - layers = [pool[i] for i in range(3)] - - stacked = MambaHybridCacheManagerV2._stack_state_views(layers) - - assert stacked.shape == (3, 4, 8) - stacked[2, 1, 3] = 5.0 - assert layers[2][1, 3] == 5.0 - - -def test_v2_stack_state_views_rejects_non_affine_layout(): - pool = torch.zeros(4, 4, 8) - stack = MambaHybridCacheManagerV2._stack_state_views +def test_v2_kda_replay_commit_without_the_replay_caches(): + """With the verify kernels' draft records (``kda_state_tok``) in place of the per-draft replay caches, the + suspend-time commit replays a slot's records onto its V2 SSM state and stores its conv window, bit for bit as the + records' replay and the conv store do on their own, and clears only that slot's draft count.""" + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import ( + replay_kda_token_records, + store_kda_replay_conv_window, + ) - assert stack([pool[0], pool[1], pool[3]]) is None - assert stack([pool[0], torch.zeros(5, 8)]) is None - assert stack([]) is None + num_layers, num_slots, num_spec, num_heads, head_dim, width = 2, 4, 2, 6, 128, 4 + section_dim = num_heads * head_dim + generator = torch.Generator(device="cuda").manual_seed(13) + def randn(*shape): + return 0.1 * torch.randn(*shape, generator=generator, device="cuda") -@pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=skip_no_cuda)]) -def test_v2_kda_admission_slot_updates_without_sync(device): - """Disagg GEN admission resets and seeds KDA replay slots; on CUDA none of it may - synchronize the stream (the previous step is still queued).""" - num_slots, num_layers, num_spec, dim, window = 8, 2, 2, 2, 3 - mgr = object.__new__(MambaHybridCacheManagerV2) - mgr._use_kda_replay_update = True - mgr._use_replay_state_update = False - mgr.mamba_ssm_rand_seed = None - mgr.local_num_mamba_layers = num_layers - mgr.prev_num_accepted_tokens = torch.arange(1, num_slots + 1, dtype=torch.int32, device=device) - # Stale device indices: the reset must take its slots from _host_state_indices. - mgr._state_indices_buffer = DeviceCopyBuffer( - torch.full((num_slots,), 7, dtype=torch.int32, device=device) - ) - mgr._host_state_indices = torch.tensor([5, 1, 0, 0, 0, 0, 0, 0], dtype=torch.int32) - mgr._request_id_to_state_index = {10: 4, 11: 2} - mgr._kda_replay_unseeded = set() - mgr.conv_state_shape = [3 * dim, window] - mgr.conv_section_dims = [dim] * 3 - mgr.all_conv_states = [ - torch.randn(num_slots, 3 * dim, window, device=device) for _ in range(num_layers) + mgr = object.__new__(KimiK3HybridCacheManager) + mgr._kda_replay = KDAReplayState(num_spec, draft_records=True) + mgr._stream = torch.cuda.Stream() + mgr.all_ssm_states = [ + randn(num_slots, num_heads, head_dim, head_dim) for _ in range(num_layers) ] - mgr._kda_replay_layers = [ + mgr.all_conv_states = [randn(num_slots, 3 * section_dim, width - 1) for _ in range(num_layers)] + mgr._kda_replay.pool_layers = [ { **{ - name: torch.randn(num_slots, dim, window + num_spec, device=device) + name: randn(num_slots, width - 1 + num_spec, section_dim).transpose(-1, -2) for name in ("kda_conv_q", "kda_conv_k", "kda_conv_v") }, - **{ - name: torch.randn(num_slots, num_spec, dim, device=device) - for name in ("kda_k_cache", "kda_g_cache", "kda_v_cache", "kda_beta_cache") - }, + "kda_state_tok": randn(num_slots, 3, num_spec, num_heads, head_dim), } for _ in range(num_layers) ] + mgr._kda_replay.prev_num_accepted_tokens = torch.tensor( + [0, 2, 1, 0], dtype=torch.int32, device="cuda" + ) + assert mgr.keeps_kda_draft_records - if device == "cuda": + mgr._kda_replay._ssm_states = tuple(mgr.all_ssm_states) + mgr._kda_replay._conv_states = tuple(mgr.all_conv_states) + slot = 1 + slots = torch.tensor([slot], dtype=torch.int32, device="cuda") + num_accepted = mgr._kda_replay.prev_num_accepted_tokens[slots.long()].clone() + expected_ssm, expected_conv = [], [] + for layer_offset, replay in enumerate(mgr._kda_replay.pool_layers): + ssm = mgr.all_ssm_states[layer_offset].clone() + conv = mgr.all_conv_states[layer_offset].clone() + replay_kda_token_records(ssm, replay["kda_state_tok"], slots, num_accepted) + store_kda_replay_conv_window( + conv, + replay["kda_conv_q"], + replay["kda_conv_k"], + replay["kda_conv_v"], + slots, + num_accepted, + ) + expected_ssm.append(ssm) + expected_conv.append(conv) + + mgr._kda_replay.commit_pending_drafts(slot) + torch.cuda.synchronize() + + assert mgr._kda_replay.prev_num_accepted_tokens.tolist() == [0, 0, 1, 0] + for actual, expected in zip(mgr.all_ssm_states, expected_ssm): + assert torch.equal(actual, expected) + for actual, expected in zip(mgr.all_conv_states, expected_conv): + assert torch.equal(actual, expected) + + +def _build_kda_replay_manager(**kwargs): + """A real V2 hybrid manager with two KDA layers (6 heads, K = V = 8) and replay caches for two drafts.""" + return _build_v2_hybrid_with_mamba_layer( + num_mamba_layers=2, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + **kwargs, + ) + + +@skip_no_cuda +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) +def test_v2_kda_replay_buffers_are_kvcm_state_buffers(manager_cls): + """Replay caches and the accepted-draft count are registered as SSM-layer buffers, counted in the per-slot state + bytes, and kept out of disaggregated transfers.""" + mgr = _build_kda_replay_manager(model_manager_cls=manager_cls) + try: + layout = _kda_replay_layout(2, 6, 8, 3) + replay_roles = {str(DataRole(name)) for name in layout} + for layer in mgr.mamba_pp_layers: + local_layer_id = mgr.layer_offsets[layer] + roles = { + str(buffer_id.role) + for pool_group in mgr.impl.pool_group_descs + for variant in pool_group.slot_desc.variants + for coalesced in variant.coalesced_buffers + for buffer_id in coalesced.buffer_ids + if int(buffer_id.layer_id) == local_layer_id + } + assert replay_roles <= roles + assert ("prev_num_accepted_tokens" in roles) == (layer == mgr.mamba_pp_layers[0]) + layer_bytes = sum(dtype.itemsize * math.prod(shape) for dtype, shape in layout.values()) + assert mgr._mamba_state_bytes_per_slot() == ( + 2 * (mgr.ssm_bytes + mgr.conv_bytes) + 2 * layer_bytes + 16 + ) + ignored = {str(role) for role in mgr.get_disagg_ignored_roles()} + assert replay_roles | {"prev_num_accepted_tokens"} <= ignored + finally: + mgr.shutdown() + + +@skip_no_cuda +def test_v2_kda_replay_acceptance_and_cuda_graph(): + mgr = _build_kda_replay_manager() + try: + requests = [_make_v2_conversation_request(i, list(range(8)), str(i)) for i in (101, 102)] + for req in requests: + _run_v2_hybrid_context(mgr, req) + slots = torch.tensor(mgr.get_state_indices([101, 102]), device="cuda") + # Mark the second row as padding through the public metadata interface. + mgr.get_state_indices([101, 102], [False, True]) + accepted = torch.tensor([3, 2], dtype=torch.int32, device="cuda") + metadata = SimpleNamespace(num_seqs=2, num_contexts=0) + for layer_id in range(2): + counts = mgr.mamba_layer_cache(layer_id).prev_num_accepted_tokens + assert counts[slots].tolist() == [0, 0] + counts[slots[1]] = 1 + # Warm up on a side stream before capture. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + mgr.update_mamba_states(metadata, accepted, state_indices=slots) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + mgr.update_mamba_states(metadata, accepted, state_indices=slots) + accepted[0] = 2 + graph.replay() + for layer_id in range(2): + assert mgr.mamba_layer_cache(layer_id).prev_num_accepted_tokens[slots].tolist() == [ + 1, + 1, + ] + finally: + mgr.shutdown() + + +@skip_no_cuda +@pytest.mark.parametrize("manager_cls", [KimiK3HybridCacheManager, Glm5NextCacheManager]) +@pytest.mark.parametrize( + "role,is_disagg", [("context", True), ("generation", True), (None, True), ("context", False)] +) +def test_v2_kda_worker_role_cache(monkeypatch, role, is_disagg, manager_cls): + """Only a disaggregated context server drops the KDA replay and draft-verify state.""" + if role is None: + monkeypatch.delenv("TRTLLM_DISAGG_ROLE", raising=False) + else: + monkeypatch.setenv("TRTLLM_DISAGG_ROLE", role) + mgr = _build_kda_replay_manager(is_disagg=is_disagg, model_manager_cls=manager_cls) + try: + req = _make_v2_conversation_request(123, list(range(8)), "role") + _run_v2_hybrid_context(mgr, req) + cache = mgr.mamba_layer_cache(0) + ctx_only = is_disagg and role == "context" + assert mgr.kda_context_only == ctx_only + assert mgr.use_kda_replay_update != ctx_only + assert isinstance(cache, KDAReplayLayerCache) != ctx_only + assert cache.conv.shape[1:] == (144, 3) + assert cache.temporal.shape[1:] == (6, 8, 8) + if ctx_only: + assert mgr.intermediate_ssm_states is None + assert mgr.intermediate_conv_states is None + assert not mgr.get_disagg_ignored_roles() & { + DataRole(name) for name in _KDA_REPLAY_FIELDS + } + mgr.update_mamba_states( + SimpleNamespace(num_seqs=1, num_contexts=1), + torch.ones(1, device="cuda", dtype=torch.int32), + ) + else: + slot = mgr.get_state_indices([123])[0] + assert cache.prev_num_accepted_tokens[slot].item() == 0 + finally: + mgr.shutdown() + + +@skip_no_cuda +@pytest.mark.parametrize("kda_draft_records", [False, True], ids=["replay_caches", "draft_records"]) +@pytest.mark.parametrize("snapshot", [False, True]) +def test_v2_kda_replay_survives_cache_lifecycle(snapshot, kda_draft_records): + """KVCM moves the replay caches, or the drafts' records in their place, with the recurrent state through + suspend, offload to host and snapshot reuse.""" + mgr = _build_kda_replay_manager( + host_cache_size=16 << 20, + enable_stats=True, + kda_draft_records=kda_draft_records, + enable_block_reuse=True, + ) + replay_fields = (*mgr._kda_replay_layout(), "prev_num_accepted_tokens") + assert ("kda_state_tok" in replay_fields) == kda_draft_records + assert mgr.reuse_match_backoff == 0 + kv = None + reused = None + pressure = [] + try: + stream = torch.cuda.current_stream().cuda_stream + kv = mgr.impl.create_kv_cache() + assert kv.resume(stream) + assert kv.resize(32) + group = mgr.impl.get_layer_group_id(LayerId(0)) + slot = kv.get_ssm_block_base_index(group) + expected = [] + for layer_id in range(2): + cache = mgr.mamba_layer_cache(layer_id) + fields = {} + for offset, name in enumerate((*replay_fields, "conv", "temporal")): + tensor = getattr(cache, name) + # Accepted counts describe the request and are shared across layers. + value = offset + 1 if name == "prev_num_accepted_tokens" else offset + layer_id + 1 + tensor[slot].fill_(value) + fields[name] = tensor[slot].clone() + expected.append(fields) + if snapshot: + kv.commit(list(range(32))) + kv.close() + kv = None + else: + kv.suspend() + # Fill the recurrent pool with newer suspended requests, forcing + # the original state to host while respecting KVCM's quota floor. + pool = next( + pool + for pool in mgr.impl.pool_group_descs + if any(variant.layer_group_id == group for variant in pool.slot_desc.variants) + ) + for _ in range(pool.num_slots + 1): + other = mgr.impl.create_kv_cache() + pressure.append(other) + assert other.resume(stream) + other.suspend() + for other in pressure: + other.commit_pending_stats() + other.close() + pressure.clear() + stats = mgr.impl.get_and_reset_iteration_stats() + assert stats[group].iter_offload_blocks > 0 + if snapshot: + reused = mgr.impl.create_kv_cache(input_tokens=list(range(32))) + assert reused.num_committed_tokens == 32 + active = reused + else: + active = kv + assert active.resume(stream) + assert active.resize(33) + new_slot = active.get_ssm_block_base_index(group) + for layer_id, fields in enumerate(expected): + cache = mgr.mamba_layer_cache(layer_id) + for name, value in fields.items(): + torch.testing.assert_close(getattr(cache, name)[new_slot], value, rtol=0, atol=0) torch.cuda.synchronize() - previous_mode = torch.cuda.get_sync_debug_mode() - torch.cuda.set_sync_debug_mode("error") + finally: + for other in pressure: + other.close() + if reused is not None: + reused.close() + if kv is not None: + kv.close() + mgr.shutdown() + + +@skip_no_cuda +def test_v2_kda_asymmetric_context_generation_transfer(monkeypatch): + """A context server without replay caches and a generation server with them describe the same transferable + recurrent state, and the generation server seeds its replay caches from the transferred conv state.""" + from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( + MambaPolicy, + mamba_receiver_payload_bytes, + ) + from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo + + monkeypatch.setenv("TRTLLM_DISAGG_ROLE", "context") + ctx = _build_kda_replay_manager(is_disagg=True) + monkeypatch.setenv("TRTLLM_DISAGG_ROLE", "generation") + gen = _build_kda_replay_manager(is_disagg=True) try: - mgr._reset_context_mamba_slots(2) - mgr._reset_kda_replay_slots([3]) - mgr.seed_kda_replay_caches_for_disagg_gen([10, 11]) + assert ctx.kda_context_only and not gen.kda_context_only + for mgr in (ctx, gen): + _run_v2_hybrid_context(mgr, _make_v2_conversation_request(101, [1, 2], "transfer")) + src_slot = ctx.get_state_indices([101])[0] + dst_slot = gen.get_state_indices([101])[0] + src_table = build_page_table_from_manager(ctx) + dst_table = build_page_table_from_manager(gen) + src_rank = RankInfo.from_kv_cache_manager("ctx", ctx, device_id=0) + dst_rank = RankInfo.from_kv_cache_manager("gen", gen, device_id=0) + MambaPolicy.validate_peer_compatible(src_rank, dst_rank, src_table, dst_table) + src_group = next(g for g in src_table.layer_groups if isinstance(g, MambaLayerGroup)) + dst_group = next(g for g in dst_table.layer_groups if isinstance(g, MambaLayerGroup)) + for group in (src_group, dst_group): + assert {pv.pool_role for pv in group.pool_views} == {MAMBA_CONV_ROLE, MAMBA_SSM_ROLE} + assert mamba_receiver_payload_bytes(src_table, dst_table, dst_slot) == sum( + gen.mamba_layer_cache(i).conv[dst_slot].numel() * 2 + + gen.mamba_layer_cache(i).temporal[dst_slot].numel() * 4 + for i in range(2) + ) + + # What the transfer writes into the generation slot. + for layer_id in range(2): + source = ctx.mamba_layer_cache(layer_id) + target = gen.mamba_layer_cache(layer_id) + source.conv[src_slot].fill_(layer_id + 1) + source.temporal[src_slot].fill_(layer_id + 2) + target.conv[dst_slot].copy_(source.conv[src_slot]) + target.temporal[dst_slot].copy_(source.temporal[src_slot]) + gen.seed_kda_replay_caches_for_disagg_gen([101]) + for layer_id in range(2): + source = ctx.mamba_layer_cache(layer_id) + target = gen.mamba_layer_cache(layer_id) + assert target.prev_num_accepted_tokens[dst_slot] == 0 + for section, name in enumerate(_KDA_REPLAY_BUFFERS[:3]): + torch.testing.assert_close( + getattr(target, name)[dst_slot, :, :3], + source.conv[src_slot, section * 48 : (section + 1) * 48].to(torch.bfloat16), + ) finally: - if device == "cuda": - torch.cuda.set_sync_debug_mode(previous_mode) + ctx.shutdown() + gen.shutdown() - # Reset 5, 1 and 3; seeding resets 4 and 2; 0, 6 and 7 untouched. - assert mgr.prev_num_accepted_tokens.tolist() == [1, 0, 0, 0, 0, 0, 7, 8] + +def _kda_cache_budget(monkeypatch, available, model, mapping, **kwargs): + from tensorrt_llm._torch.modules.kimi_kda import _kda_kernels + + monkeypatch.setattr(_kda_kernels, "is_kda_mtp_verify_available", lambda: available) + monkeypatch.setattr(_kda_kernels, "kda_mtp_max_num_spec", lambda: 7) + return KimiK3HybridCacheManager.get_cache_size_per_token(model, mapping, **kwargs) + + +@pytest.mark.parametrize( + "role,is_disagg,available", + [ + ("context", True, True), + ("generation", True, True), + (None, True, True), + ("context", False, True), + ("generation", True, False), + ], +) +@pytest.mark.parametrize("interval,replay_slope", [(0, 0), (32, 95), (33, 92), (48, 63), (256, 0)]) +def test_kda_cache_budget_accounts_for_replay_role( + monkeypatch, role, is_disagg, available, interval, replay_slope +): + if role is None: + monkeypatch.delenv("TRTLLM_DISAGG_ROLE", raising=False) + else: + monkeypatch.setenv("TRTLLM_DISAGG_ROLE", role) + model = _kimi_model_config() + spec_config = MTPDecodingConfig(max_draft_len=2) + mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + kwargs = dict( + max_batch_size=2, + max_seq_len=128, + tokens_per_block=32, + kv_cache_config=KvCacheConfig( + enable_block_reuse=True, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=interval), + ), + spec_config=spec_config, + is_disagg=is_disagg, + ) + baseline = _kda_cache_budget(monkeypatch, False, model, mapping, **kwargs) + result = _kda_cache_budget(monkeypatch, available, model, mapping, **kwargs) + if available and not (is_disagg and role == "context"): + # Two layers: three BF16 [5,32] conv windows, two BF16 and one + # FP32 [2,32] histories and BF16 [2,8] padded beta per layer; + # one shared count per slot, padded to 16 bytes. + per_layer = 3 * 5 * 32 * 2 + 2 * 32 * 8 + 2 * 8 * 2 + # At interval 33, round the combined recurrent + replay snapshot + # cost once: ceil(6224 / 33) - ceil(3200 / 33) = 92 bytes/token. + assert result[0] - baseline[0] == replay_slope + live_and_dummy_slots = 2 + _get_num_cuda_graph_padding_dummy_slots(spec_config, 2) + assert result[1] - baseline[1] == live_and_dummy_slots * (2 * per_layer + 16) + else: + assert result == baseline + + +@pytest.mark.parametrize("enable_block_reuse", [False, True]) +@pytest.mark.parametrize("enable_attention_dp", [False, True]) +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("pp_partition,stage_kda_layers", [(None, [2, 0]), ([2, 2], [1, 1])]) +def test_kda_cache_budget_counts_live_dummy_and_fixed_snapshot_slots( + monkeypatch, enable_block_reuse, enable_attention_dp, rank, pp_partition, stage_kda_layers +): + model = _kimi_model_config() + mapping = Mapping( + world_size=2, + rank=rank, + tp_size=1, + pp_size=2, + enable_attention_dp=enable_attention_dp, + pp_partition=pp_partition, + ) + spec_config = MTPDecodingConfig(max_draft_len=2, draft_len_schedule={1: 2, 2: 1}) + kwargs = dict( + max_batch_size=4, + max_seq_len=128, + tokens_per_block=32, + kv_cache_config=KvCacheConfig( + enable_block_reuse=enable_block_reuse, + mamba_state_config=MambaStateConfig( + periodic_snapshot_interval=0, + additional_snapshot_offsets_from_start=[32, 64, 64, 256], + additional_snapshot_offsets_from_end=[0, 32, 128], + ), + ), + spec_config=spec_config, + is_disagg=False, + ) + baseline = _kda_cache_budget(monkeypatch, False, model, mapping, **kwargs) + result = _kda_cache_budget(monkeypatch, True, model, mapping, **kwargs) + # Without an explicit partition, splitting the four target layers plus + # one MTP layer puts both KDA layers on rank 0. Partition [2,2] places one + # KDA layer on each rank, with the MTP layer appended to rank 1. + # Every live slot (max_batch_size * pp_size), every CUDA-graph padding + # dummy plus one more for attention DP, and with reuse four reachable, + # distinct snapshot rules per resident sequence carry replay caches. + local_kda_layers = stage_kda_layers[rank] + replay_bytes_per_slot = 1504 * local_kda_layers + (16 if local_kda_layers else 0) + slots = ( + 8 + + _get_num_cuda_graph_padding_dummy_slots(spec_config, 4) + + int(enable_attention_dp) + + (8 * 4 if enable_block_reuse else 0) + ) + assert result == (baseline[0], baseline[1] + slots * replay_bytes_per_slot) + + +def test_v2_kda_replay_seeds_disaggregated_generation_slots(): + mgr = object.__new__(KimiK3HybridCacheManager) + mgr._kda_replay = KDAReplayState(2) + 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._kda_replay.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 dim_contiguous_conv_cache() -> torch.Tensor: + return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) + + replay = {} + replay["kda_conv_q"] = dim_contiguous_conv_cache() + replay["kda_conv_k"] = dim_contiguous_conv_cache() + replay["kda_conv_v"] = dim_contiguous_conv_cache() + replay["kda_k_cache"] = torch.full((2, 4, 2, 2), 7.0) + replay["kda_g_cache"] = torch.full((2, 4, 2, 2), 7.0) + replay["kda_v_cache"] = torch.full((2, 4, 2, 2), 7.0) + replay["kda_beta_cache"] = torch.full((2, 4, 2, 1), 7.0) + + mgr._kda_replay.pool_layers = [ + {name: buffer[layer] for name, buffer in replay.items()} for layer in range(2) + ] + + mgr._kda_replay._conv_states = tuple(mgr.all_conv_states) + mgr._kda_replay._conv_section_dims = tuple(mgr.conv_section_dims) + mgr._kda_replay._committed_window = mgr.conv_state_shape[1] + mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) + + for layer_offset, conv_state in enumerate(mgr.all_conv_states): + for slot in (1, 3): + for replay_buffer, start in ( + (replay["kda_conv_q"], 0), + (replay["kda_conv_k"], 2), + (replay["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(replay["kda_k_cache"][layer_offset, slot]) == 0 + assert torch.count_nonzero(replay["kda_g_cache"][layer_offset, slot]) == 0 + assert torch.count_nonzero(replay["kda_v_cache"][layer_offset, slot]) == 0 + assert torch.count_nonzero(replay["kda_beta_cache"][layer_offset, slot]) == 0 + assert mgr._kda_replay.prev_num_accepted_tokens[slot] == 0 + assert mgr._kda_replay.prev_num_accepted_tokens[0] == 5 + assert mgr._kda_replay.prev_num_accepted_tokens[2] == 5 + + +_KDA_REPLAY_FIELDS = (*_KDA_REPLAY_BUFFERS, "prev_num_accepted_tokens") 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..9ff123faec8b --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_model_managers.py @@ -0,0 +1,1625 @@ +# 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._device_copy import DeviceCopyBuffer +from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.cache_manager import ( + Glm5NextCacheManager, +) +from tensorrt_llm._torch.modules.fla.cache_manager import ( + GDNIntermediateState, + GDNLayerCache, + GDNReplayState, + Qwen35HybridCacheManager, +) +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import ( + KDAIntermediateLayerCache, + KDAIntermediateState, + KDAReplayLayerCache, + KDAReplayState, + KimiK3HybridCacheManager, +) +from tensorrt_llm._torch.modules.mamba.cache_manager import ( + Mamba2LayerCache, + Mamba2State, + NemotronHybridCacheManager, + ReplayHistory, +) +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + PLE_CONV_STATE, + PLE_NGRAM_CONTEXT, + Qwen4ExpHybridCacheManager, + 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 ( + MambaAcceptanceBatch, + MambaHybridCacheManager, + MambaLayerCache, + MambaRole, + MambaStateLayout, +) +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, + BufferConfig, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheManagerConfig, + LayerId, +) + + +@pytest.mark.parametrize( + "manager_cls", + [ + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, + Glm5NextCacheManager, + ], +) +def test_model_managers_inherit_the_common_lifecycle(manager_cls): + assert manager_cls.__bases__ == (MambaHybridCacheManagerV2,) + for method in ( + "try_commit_blocks", + "_setup_state_pool", + ): + assert getattr(manager_cls, method) is getattr(MambaHybridCacheManagerV2, method) + + +@pytest.mark.parametrize( + "manager_cls", + [ + MambaHybridCacheManagerV2, + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, + Glm5NextCacheManager, + ], +) +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, + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, + Glm5NextCacheManager, + 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 _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, Qwen4ExpHybridCacheManager): + 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, Qwen4ExpHybridCacheManager): + 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 + manager._state_layout = replace(_layout(), spec_config=manager.spec_config) + manager.local_num_mamba_layers = 2 + manager._global_n_groups = 2 + if isinstance(manager, Qwen4ExpHybridCacheManager): + 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 (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_spec_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", + [ + NemotronHybridCacheManager, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, + Glm5NextCacheManager, + ], +) +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._requested_draft_records = False + 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_spec_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), + ], +) +@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 +): + manager = object.__new__(manager_cls) + manager._requested_num_spec = requested_num_spec + manager._requested_draft_records = False + 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=KDAIntermediateState) + replay_factory = Mock(wraps=KDAReplayState) + selector = Mock(return_value=selected_num_spec) + monkeypatch.setitem( + manager_cls._initialize_spec_state.__globals__, + "KDAIntermediateState", + intermediate_factory, + ) + monkeypatch.setitem( + manager_cls._initialize_spec_state.__globals__, + "KDAReplayState", + replay_factory, + ) + monkeypatch.setitem( + manager_cls._initialize_spec_state.__globals__, + "get_kda_replay_num_spec", + selector, + ) + + 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() + replay_kwargs = {"stochastic_rounding": False} + if manager_cls is KimiK3HybridCacheManager: + replay_kwargs["draft_records"] = False + replay_factory.assert_called_once_with(num_spec, **replay_kwargs) + assert manager._speculative_state is manager._kda_replay + 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 = [], [] + pool_bind = Mock() + monkeypatch.setitem( + manager_cls._setup_model_state.__globals__, "bind_kda_replay_pool", pool_bind + ) + bind = Mock() + monkeypatch.setattr(state, "bind", bind) + manager._setup_model_state() + if num_spec is not None: + pool_bind.assert_called_once_with(state, manager) + else: + pool_bind.assert_not_called() + bind.assert_called_once_with( + manager._state_layout, manager.all_ssm_states, manager.all_conv_states + ) + + shutdown = Mock(wraps=state.shutdown) + monkeypatch.setattr(state, "shutdown", shutdown) + _shutdown_manager(manager, monkeypatch) + shutdown.assert_called_once_with() + _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(monkeypatch): + manager = object.__new__(MambaHybridCacheManagerV2) + manager.spec_config = None + 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) + _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 + assert not hasattr(manager, "get_replay_state_update_metadata") + + +@pytest.mark.parametrize( + "manager_cls", + [ + MambaHybridCacheManager, + MambaHybridCacheManagerV2, + Qwen35HybridCacheManager, + Qwen4ExpHybridCacheManager, + KimiK3HybridCacheManager, + Glm5NextCacheManager, + ], +) +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(monkeypatch): + manager = object.__new__(NemotronHybridCacheManager) + manager.spec_config = None + manager._requested_replay = False + 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 + manager._speculative_state = manager._initialize_spec_state() + manager._setup_model_state() + seeds = manager.get_mamba_ssm_rand_seed() + before = seeds.clone() + manager._state_indices_buffer = DeviceCopyBuffer(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 + _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, monkeypatch): + 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_k_cache is not None + assert payload.kda_g_cache is not None + elif isinstance(state, (ReplayHistory, GDNReplayState)): + assert payload.old_x is not None + else: + assert payload.intermediate_ssm is not None + assert not hasattr(manager, "get_replay_state_update_metadata") + _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 + + +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_spec_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_buffer = DeviceCopyBuffer(torch.empty(0, dtype=torch.bool)) + with pytest.raises(NotImplementedError, match="model manager"): + 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( + "manager_cls", + [NemotronHybridCacheManager, Qwen35HybridCacheManager, Qwen4ExpHybridCacheManager], +) +@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._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 Qwen4ExpHybridCacheManager: + 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) + manager.local_num_mamba_layers = 1 + manager._state_indices_buffer = DeviceCopyBuffer(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_buffer = DeviceCopyBuffer(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) + _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": + assert state.intermediate_ssm is None + + +def test_ple_adds_buffers_without_overriding_snapshot_or_budget_solver(): + manager = object.__new__(Qwen4ExpHybridCacheManager) + 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__(Qwen4ExpHybridCacheManager) + 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._guard_page_by_layer = {} + 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 GDNIntermediateState() + manager._setup_model_state() + state = manager._speculative_state + scratch_before = {} + for name, tensor in vars(state).items(): + if isinstance(tensor, torch.Tensor) and not name.startswith("_stacked_"): + 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 (Qwen35HybridCacheManager, Qwen4ExpHybridCacheManager): + 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", [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), [], []) + state.reset_slots(torch.tensor([0]), [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)): + assert replay.dtype is torch.bfloat16 + torch.testing.assert_close( + replay[layer, 2, :, :3], + conv[2, section * 8 : (section + 1) * 8].to(torch.bfloat16), + ) + assert torch.count_nonzero(replay[layer, 2, :, 3:]) == 0 + assert state.prev_num_accepted_tokens[2] == 0 + assert torch.count_nonzero(state.kda_k_cache[:, 2]) == 0 + assert torch.count_nonzero(state.kda_g_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) + # Beta cache rows retain their logical head count but include alignment padding. + actual = sum(buffer.untyped_storage().nbytes() 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 + + +@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", "both"]) +@pytest.mark.parametrize("granularity", [1, 4096]) +def test_kda_pool_allocation_includes_replay_in_native_quota( + monkeypatch, manager_cls, replay, pool_share, quota_source, granularity +): + 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 = manager._mamba_state_bytes_per_slot() + total_quota = 1 << 20 + 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 + granularity - 1) // granularity * granularity + ) + 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 + granularity - 1) // granularity * granularity + return True + + monkeypatch.setattr(module, "KVCacheManagerPy", Pool) + config = KVCacheManagerConfig( + tokens_per_block=32, + layers=[], + cache_tiers=[GpuCacheTierConfig(quota=total_quota), HostCacheTierConfig(quota=12345)], + ) + 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 + assert pool.resize_calls == 0 + assert pool.quota == (total_quota + granularity - 1) // granularity * granularity + assert manager._get_external_cache_bytes(pool) == 0 + # 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, 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, 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 + from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import common + + layout = _layout() + params = SimpleNamespace( + n_groups=2, + state_size=4, + head_dim=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, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=interval), + ) + 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 + ) + # KVCM pads the shared accepted-draft count from 4 to 16 bytes. + state_bytes = 2 * 416 + scratch_per_slot + (12 if replay else 0) + interval = kv_config.mamba_state_config.periodic_snapshot_interval + slope = (state_bytes + interval - 1) // interval if reuse and interval > 0 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"] +) +def test_layer_payload_preserves_storage_and_strides(kind, layer): + states = { + "intermediate": GDNIntermediateState, + "kda_intermediate": KDAIntermediateState, + "mamba2": lambda: Mamba2State(stochastic_rounding=True), + "replay": lambda: ReplayHistory(3), + "gdn": lambda: GDNReplayState(3), + "kda": lambda: KDAReplayState(2, stochastic_rounding=True), + } + state = states[kind]() + 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": 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 + 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_k_cache", + "kda_g_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 + 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: + 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, 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() + 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, 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) + 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 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__ == (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 = 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(), mapping=Mapping(world_size=2, tp_size=2, rank=1)) + views = _state_views() + 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) + 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", + [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 + + 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 state.intermediate_ssm is None + state.shutdown() + + +@pytest.mark.parametrize( + "state", + [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, + (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, + (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. + assert views.all_conv_states[0].shape == (5, 24, 3) + + +@pytest.mark.parametrize( + "cls", [GDNIntermediateState, 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) + + +@pytest.mark.parametrize( + "manager_cls", + [ + 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 KimiK3HybridCacheManager: + 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__(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) + 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__(KimiK3HybridCacheManager) + 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__(Qwen4ExpHybridCacheManager) + 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_ = Qwen4ExpHybridCacheManager._get_view_by_role_and_layer.__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_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 + 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_view_by_role_and_layer(42, roles, dtype=torch.float32, page_shape=shape) + + +def test_qwen4_position_uses_model_owned_slot_view(monkeypatch): + manager = object.__new__(Qwen4ExpHybridCacheManager) + manager.qsa_position_layer_id = 42 + manager.tokens_per_block = 4 + 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_view_by_role_and_layer.assert_called_once() + manager.qsa_position_layer_id = None + manager._get_view_by_role_and_layer.reset_mock() + assert manager.get_qsa_position_buffer() is None + manager._get_view_by_role_and_layer.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__(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 + 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_view_by_role_and_layer = Mock(return_value=sentinel) + assert manager.get_buffers(42, kv_layout) is sentinel + manager._get_view_by_role_and_layer.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_qwen4_attention_pool_layout_validates_backend_contract(invalid): + manager = object.__new__(Qwen4ExpHybridCacheManager) + 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) + 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_qsa_attention_pool_layout() == (0, 6) + else: + with pytest.raises(RuntimeError): + manager.get_qsa_attention_pool_layout() + + +def test_qwen4_requires_one_mapping_and_skips_nonlocal_layers(): + 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} + 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 {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() + + +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(Qwen4ExpHybridCacheManager): + def _initialize_spec_state(self): + state = super()._initialize_spec_state() + events.append("validate") + return state + + 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) + + +@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/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index fb7dcf8470a6..6fb0aecd316b 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -2877,6 +2877,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) @@ -2896,6 +2897,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) @@ -2915,6 +2917,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") @@ -2932,6 +2935,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") diff --git a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py index 5763a76c2ed8..6239f8fb2cec 100644 --- a/tests/unittest/_torch/modeling/test_glm5_next_contracts.py +++ b/tests/unittest/_torch/modeling/test_glm5_next_contracts.py @@ -830,10 +830,28 @@ def __init__(self, *args, **kwargs): ) allocate.assert_called_once() assert allocate.call_args.kwargs["dtype"] == (DataType.FP8 if fp8_kv_cache else DataType.BF16) + assert allocate.call_args.kwargs["conv_state_layout"] == "q_k_v" + assert "kda_replay_num_spec" not in allocate.call_args.kwargs + manager = object.__new__(Glm5NextCacheManager) + manager._requested_num_spec = None + manager._requested_draft_records = False + manager._kda_replay = None + manager._mamba_ssm_stochastic_rounding = False + manager.spec_config = spec_config + manager._state_layout = SimpleNamespace( + spec_config=spec_config, + mamba_pp_layers=(), + conv_state_layout="q_k_v", + conv_section_dims=(), + ) + with patch( + "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", + return_value=verify_kernel is True, + ): + state = manager._initialize_spec_state() + assert manager.use_kda_replay_update is (verify_kernel is True) if verify_kernel is True: - assert allocate.call_args.kwargs["kda_replay_num_spec"] == 3 - else: - assert "kda_replay_num_spec" not in allocate.call_args.kwargs + assert state.num_speculative_tokens == 3 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") 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/_torch/modeling/test_qsa_runtime_wiring.py b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py index ee5b97703635..3438be616700 100644 --- a/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py +++ b/tests/unittest/_torch/modeling/test_qsa_runtime_wiring.py @@ -10,18 +10,22 @@ 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 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 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: @@ -222,7 +226,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__(Qwen4ExpHybridCacheManager) manager.qsa_position_layer_id = None metadata = object.__new__(QSAAttentionMetadata) metadata.kv_cache_manager = manager @@ -253,7 +257,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 Qwen4ExpHybridCacheManager def test_qsa_hybrid_rejects_kv_cache_manager_v1(monkeypatch) -> None: @@ -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 @@ -327,24 +336,23 @@ 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( model_engine=None, - kv_cache_manager_cls=QSAMambaHybridCacheManagerV2, + 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, 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,21 +364,28 @@ 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_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 = Qwen4ExpHybridCacheManager(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)) def test_qsa_cache_manager_delegates_regular_kv_layout( monkeypatch: pytest.MonkeyPatch, dtype: DataType ) -> None: - manager = object.__new__(QSAMambaHybridCacheManagerV2) + manager = object.__new__(Qwen4ExpHybridCacheManager) 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 ddf59ea8ca34..7dac22c46df8 100644 --- a/tests/unittest/_torch/modeling/test_qwen4_exp_support.py +++ b/tests/unittest/_torch/modeling/test_qwen4_exp_support.py @@ -483,7 +483,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, ) @@ -496,6 +495,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 @@ -506,7 +509,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 @@ -518,7 +521,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()) @@ -547,9 +552,9 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state( offsets_from_start, offsets_from_end, expected_state_slots ) -> 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 ( + Qwen4ExpHybridCacheManager, + extract_qwen4_exp_ple_cache_params, ) from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MambaStateConfig from tensorrt_llm.mapping import Mapping @@ -568,10 +573,10 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state( ), ), } - with_ple = MambaHybridCacheManagerV2.get_cache_size_per_token( + with_ple = Qwen4ExpHybridCacheManager.get_cache_size_per_token( SimpleNamespace(pretrained_config=config, quant_config=None), **common ) - without_ple = MambaHybridCacheManagerV2.get_cache_size_per_token( + without_ple = Qwen4ExpHybridCacheManager.get_cache_size_per_token( SimpleNamespace(pretrained_config=no_ple_config, quant_config=None), **common ) @@ -585,9 +590,10 @@ def test_v2_cache_estimator_counts_ple_lifecycle_state( 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, + Qwen4ExpHybridCacheManager, ) ngram_context = torch.full((12, 2), 11, dtype=torch.int64) @@ -597,20 +603,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(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] @@ -625,8 +631,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/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", [ diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 04f06439febd..6c75de02a7ea 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -15,12 +15,17 @@ """Unit tests for Mamba2 metadata preparation optimizations.""" from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from tensorrt_llm._device_copy import DeviceCopyBuffer 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 +37,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(), @@ -44,7 +45,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): @@ -60,6 +60,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 + + class _StateIndexCacheManager: """Returns ``state_indices`` as set, a list or a CPU tensor, for a batch of decode requests.""" @@ -266,8 +314,6 @@ def counted(dst, src, *args, **kwargs): 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( 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 diff --git a/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py b/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py index a2c8ae831694..16306eecfa54 100644 --- a/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py +++ b/tests/unittest/_torch/moe/kernels/parallel/test_cute_dsl_moe.py @@ -1313,6 +1313,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, @@ -1322,10 +1323,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) diff --git a/tests/unittest/disaggregated/test_extractor.py b/tests/unittest/disaggregated/test_extractor.py index 7e7e3f7f0803..ec6ef42ad0f7 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 KVCacheManagerV2, 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 KVCacheManagerV2, Role +from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaRole from tensorrt_llm._torch.pyexecutor.resource_manager import ( CacheTypeCpp, DataType, @@ -531,7 +532,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 @@ -568,7 +569,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)], ), ], ) @@ -594,7 +595,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( @@ -647,7 +648,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] @@ -666,7 +667,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 @@ -677,7 +678,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, diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index cfbecb8b5a77..3538f65a86ca 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -66,8 +66,8 @@ build_page_table_from_manager, ) from tensorrt_llm._torch.disaggregation.resource.page import MambaLayerGroup +from tensorrt_llm._torch.modules.kimi_kda.cache_manager import KimiK3HybridCacheManager from tensorrt_llm._torch.pyexecutor.config_utils import extract_mamba_kv_cache_params -from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import DataType @@ -126,7 +126,7 @@ def _kimi_linear_config(): def _create_kda_managers( tp: int, enable_attention_dp: bool = False, max_batch_size: int = MAX_BATCH_SIZE ): - """Create MambaHybridCacheManagerV2s with K3-style KDA slots. + """Create KimiK3HybridCacheManagers with K3-style KDA slots. Mirrors the ``is_kimi_linear`` route in ``_util.py``: the mamba params come from ``extract_mamba_kv_cache_params`` UNSCALED (global @@ -150,7 +150,7 @@ def _create_kda_managers( pp_size=1, enable_attention_dp=enable_attention_dp, ) - mgr = MambaHybridCacheManagerV2( + mgr = KimiK3HybridCacheManager( mamba_d_state=params.state_size, mamba_d_conv=params.conv_kernel, mamba_num_heads=params.num_heads, diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index ebdddf36c802..8cde60c9f05e 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 Qwen35HybridCacheManager +from tensorrt_llm._torch.modules.mamba.cache_manager import NemotronHybridCacheManager +from tensorrt_llm._torch.modules.qwen4_exp.cache_manager import ( + Qwen4ExpHybridCacheManager, + 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 = Qwen4ExpHybridCacheManager + elif conv_state_layout == "q_k_v": + manager_cls = Qwen35HybridCacheManager + else: + manager_cls = NemotronHybridCacheManager 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, Qwen4ExpHybridCacheManager): for state in manager._ple_conv_states.values(): state.zero_() for state in manager._ple_ngram_contexts.values(): @@ -844,7 +856,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):