diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f830d3c23191..98ef12781bac 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3048,6 +3048,14 @@ def _create_kv_cache_manager( if is_kda_mtp_verify_available(): kda_extra_kwargs["kda_replay_num_spec"] = ( spec_config.tokens_per_gen_step - 1) + # 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_token_states`. + backbone = getattr(getattr(model_engine, "model", None), + "model", None) + if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) + and getattr(backbone, "kda_token_states", False)): + kda_extra_kwargs["kda_token_states"] = True if is_glm5_next: # The manager places an indexer buffer on every attention layer. from ..attention.backends.sparse.glm_kpool import \ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index ad417618950b..8f4767ca0efb 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -20,6 +20,7 @@ from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, Optional, Protocol, Sequence, Tuple, Union, cast) +import numpy as np import torch import triton import triton.language as tl @@ -52,7 +53,9 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, BatchDesc, BufferConfig, - DataRole, KVCacheDesc) + DataRole, + GpuCacheTierConfig, + KVCacheDesc) from tensorrt_llm.runtime.kv_cache_manager_v2 import \ KVCacheManagerConfig as KVCacheManagerConfigPy from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, @@ -330,6 +333,13 @@ 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_token_states(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. @@ -407,7 +417,9 @@ class SpeculativeState(State): - 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) + 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" @@ -434,10 +446,21 @@ class SpeculativeState(State): kda_conv_v: torch.Tensor | None = None # Post-processed per-draft quantities for replay: # [layers, slots, num_spec, 3, H*K] (q/k/gate), [.., num_spec, H*V], - # [.., num_spec, H] — all fp32. + # [.., num_spec, H] — all fp32. None with kda_state_tok. kda_qkg_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). @@ -456,9 +479,11 @@ class SpeculativeState(State): @property def has_kda_replay_caches(self) -> bool: - """True when the KDA fused-verify replay caches were allocated - (the fused ``trtllm::kda_mtp_decode`` verify path).""" - return self.kda_qkg_cache is not None + """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: @@ -3024,6 +3049,7 @@ def __init__( 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_token_states: bool = False, **kwargs, ) -> None: if conv_state_layout not in ("x_b_c", "q_k_v"): @@ -3040,6 +3066,9 @@ def __init__( 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_token_states = kda_token_states and self._use_kda_replay_update if self._use_kda_replay_update: if use_replay_state_update: raise ValueError( @@ -3420,6 +3449,10 @@ def get_disagg_role_layouts(self) -> Dict[DataRole, RoleLayout]: def use_gdn_cached_replay_all_layer_commit(self) -> bool: return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) + @property + def keeps_kda_token_states(self) -> bool: + return getattr(self, "kda_state_tok", None) is not None + def _setup_mtp_intermediate_states(self, spec_config, max_batch_size: int) -> None: if not self.use_kda_replay_update: @@ -3450,6 +3483,7 @@ def _allocate_pool_replay_buffers( self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None + self.kda_state_tok = None if (not self.use_kda_replay_update or self.local_num_mamba_layers == 0): return allocated @@ -3482,30 +3516,44 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: self.kda_conv_q = allocate_dim_contiguous_conv_cache() self.kda_conv_k = allocate_dim_contiguous_conv_cache() self.kda_conv_v = allocate_dim_contiguous_conv_cache() - self.kda_qkg_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - 3, - section_dim, - dtype=torch.float32, - device=device, - ) - self.kda_v_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - section_dim, - dtype=torch.float32, - device=device, - ) - self.kda_beta_cache = _allocate_kda_beta_cache( - (self.local_num_mamba_layers, cache_size, num_spec, - self.ssm_state_shape[0]), - device, - ) + if self._kda_token_states: + # Beside the V2 pools. With the SSM pool at its live floor, + # _build_cache_config took these bytes out of the GPU quota. The + # verify replays the accepted drafts from these records, so the + # per-draft replay caches are not allocated. + self.kda_state_tok = torch.zeros( + (self.local_num_mamba_layers, cache_size, + *self._kda_draft_records_shape()), + dtype=torch.float32, + device=device, + ) + else: + self.kda_qkg_cache = torch.zeros( + self.local_num_mamba_layers, + cache_size, + num_spec, + 3, + section_dim, + dtype=torch.float32, + device=device, + ) + self.kda_v_cache = torch.zeros( + self.local_num_mamba_layers, + cache_size, + num_spec, + section_dim, + dtype=torch.float32, + device=device, + ) + self.kda_beta_cache = _allocate_kda_beta_cache( + (self.local_num_mamba_layers, cache_size, num_spec, + self.ssm_state_shape[0]), + device, + ) + per_token = (", with the drafts' records in place of the replay " + "caches" if self.kda_state_tok is not None else "") logger.info("Mamba Cache (kda-replay) is allocated for " - f"{cache_size} state slots") + f"{cache_size} state slots{per_token}") return True def _commit_gdn_cached_replay_history_layers( @@ -3798,9 +3846,26 @@ def _get_pool_roles(self, 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 _kda_token_state_bytes_per_slot(self) -> int: + """Bytes of ``kda_state_tok`` per SSM slot: the fp32 records of every + draft of every local Mamba layer, or 0 without them.""" + if not getattr(self, "_kda_token_states", False): + return 0 + return (self.local_num_mamba_layers * + math.prod(self._kda_draft_records_shape()) * + torch.float32.itemsize) + def _mamba_state_bytes_per_slot(self) -> int: - base_bytes = self.local_num_mamba_layers * (self.ssm_bytes + - self.conv_bytes) + base_bytes = (self.local_num_mamba_layers * + (self.ssm_bytes + self.conv_bytes) + + self._kda_token_state_bytes_per_slot()) local_ple_layers = sum(layer_id in self.pp_layers for layer_id in self._ple_layer_ids) if local_ple_layers == 0: @@ -3896,6 +3961,16 @@ def _get_typical_request_capacity( "pool ratio.") return fallback_capacity + def _state_quota(self, state_slots: int) -> int: + """Quota for ``state_slots`` recurrent-state slots, or for every slot + the live floor keeps (``_live_floor_token_state_slots``) when that is + more: each holds its states and its drafts' records.""" + slots = max( + state_slots, + self._live_floor_token_state_slots( + self.kv_cache_config.max_util_for_resume)) + return slots * self._mamba_state_bytes_per_slot() + 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() @@ -3903,7 +3978,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: 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() + state_quota = self._state_quota(state_slots) # 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. @@ -3939,15 +4014,77 @@ def _minimum_live_gpu_quota(self) -> int: 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() + state_quota = self._state_quota(num_state_slots) return max( self._get_quota_from_max_tokens(0), state_quota + attention_block_quota, ) + def _ssm_pool_at_live_floor(self, kv_cache_config: KvCacheConfig) -> bool: + """Whether the SSM pool group keeps only its min-slots floor (one slot + per resident sequence plus the reserved dummies) instead of its share + of the typical step's ratio. + + With per-token KDA states, ``kda_state_tok`` takes num_spec drafts' + records per SSM slot beside the V2 pools + (``_allocate_pool_replay_buffers``), addressed by the request's SSM + slot. Without block reuse no SSM slot holds anything but a live + request or a reserved dummy, so slots past the floor are never used + and would only add that memory. At the floor the pool's slot count is + known before sizing, so ``_build_cache_config`` takes those records' + bytes out of the GPU quota. + """ + return (getattr(self, "_kda_token_states", False) + and self.local_num_mamba_layers > 0 + and not kv_cache_config.enable_block_reuse + and self._attention_cache_bytes_per_token() > 0) + + def _live_floor_token_state_slots(self, max_util_for_resume: float) -> int: + """SSM slots whose ``kda_state_tok`` entries come out of the GPU quota. + + With the SSM pool at its live floor and no ``pool_ratio``, that is + every slot the pool keeps: the floor's min-slots constraint as the + runtime scales it by 1 / max_util_for_resume. (Rounding each pool up + to whole allocation grains can add slots when they are smaller than a + grain.) ``_build_cache_config`` takes their records out of the GPU + quota before the pools are sized. Otherwise 0. + """ + kv_cache_config = self.kv_cache_config + if (kv_cache_config.pool_ratio is not None + or not self._ssm_pool_at_live_floor(kv_cache_config)): + return 0 + floor_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + # The runtime holds max_util_for_resume as a C++ float; round the + # Python value the same way, or the ceiling can differ by one slot + # when floor_slots / max_util_for_resume is a whole number. + return math.ceil(floor_slots / float(np.float32(max_util_for_resume))) + + def _quota_filling_request_capacity(self, gpu_quota: int) -> int: + """A typical request capacity whose resident requests' attention pages + alone take ``gpu_quota``. The typical step's ratio then gives the SSM + pool group fewer slots than its floor, so the min-slots constraint + sets its size and attention gets the rest of the quota.""" + bytes_per_token = (self._max_resident_sequences() * + self._attention_cache_bytes_per_token()) + # KVCacheDesc.capacity is a 32-bit int. + return min(-(-gpu_quota // bytes_per_token), 1 << 30) + def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: kv_cache_config = self.kv_cache_config + if (getattr(self, "_kda_token_states", False) + and self.local_num_mamba_layers > 0 + and (config.initial_pool_ratio is not None + or not self._ssm_pool_at_live_floor(kv_cache_config))): + # kda_state_tok is allocated for every slot the SSM pool keeps; + # only at the live floor is that count known before sizing, so + # that its bytes can come out of the quota. + raise ValueError( + "Per-token KDA states (kda_token_states) need the SSM pool at " + "its live floor: kv_cache_config.enable_block_reuse off, " + "kv_cache_config.pool_ratio unset, and attention layers on " + "this rank.") cache_tiers = config.cache_tiers gpu_quota = cache_tiers[0].quota minimum_live_quota = self._minimum_live_gpu_quota() @@ -3956,6 +4093,25 @@ def _build_cache_config( "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.") + ssm_pool_at_live_floor = (config.initial_pool_ratio is None and + self._ssm_pool_at_live_floor(kv_cache_config)) + if ssm_pool_at_live_floor: + # kda_state_tok is allocated beside the pools, one entry per SSM + # slot: take those records' bytes out of the GPU quota before the + # pools are sized. + token_state_slots = self._live_floor_token_state_slots( + config.max_util_for_resume) + token_state_quota = (token_state_slots * + self._kda_token_state_bytes_per_slot()) + gpu_quota -= token_state_quota + cache_tiers = [ + GpuCacheTierConfig(quota=gpu_quota), *cache_tiers[1:] + ] + logger.info( + "KV cache manager v2: the KDA draft records of " + f"{token_state_slots} SSM slots take " + f"{token_state_quota / (1 << 30)}GiB of the device quota; the " + f"pools get {gpu_quota / (1 << 30)}GiB") # _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. @@ -3990,6 +4146,10 @@ def _build_cache_config( if config.initial_pool_ratio is None: typical_capacity = self._get_typical_request_capacity( kv_cache_config) + if ssm_pool_at_live_floor: + typical_capacity = max( + typical_capacity, + self._quota_filling_request_capacity(gpu_quota)) request_descs = self._typical_request_descs(typical_capacity, kv_cache_config) typical_step = BatchDesc(request_descs * @@ -4016,6 +4176,7 @@ def _build_cache_config( ] return replace( config, + cache_tiers=cache_tiers, layers=layers, typical_step=typical_step, constraints=constraints, @@ -4285,16 +4446,17 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], destination_slots = torch.tensor([new for _, new in moves], dtype=torch.long, device=device) - replay_buffers = ( - self.kda_conv_q, - self.kda_conv_k, - self.kda_conv_v, - self.kda_qkg_cache, - self.kda_v_cache, - self.kda_beta_cache, - ) + replay_buffers = [self.kda_conv_q, self.kda_conv_k, self.kda_conv_v] + assert all(replay_buffer is not None + for replay_buffer in replay_buffers) + # The per-draft replay caches, or the drafts' records in their place. + replay_buffers += [ + replay_buffer + for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, + self.kda_beta_cache, self.kda_state_tok) + if replay_buffer is not None + ] for replay_buffer in replay_buffers: - assert replay_buffer is not None replay_buffer.index_copy_( 1, destination_slots, @@ -4468,10 +4630,12 @@ def seed_kda_replay_caches_for_disagg_gen(self, replay_layer[:, :, :committed_window].index_copy_( 0, state_indices, section.to(replay_layer.dtype)) + # The drafts' records need no reset: with no pending drafts, none of + # them is read. for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, self.kda_beta_cache): - assert replay_buffer is not None - replay_buffer.index_fill_(1, state_indices, 0) + if replay_buffer is not None: + replay_buffer.index_fill_(1, state_indices, 0) self.prev_num_accepted_tokens[state_indices] = 0 def mamba_layer_cache( @@ -4486,18 +4650,21 @@ def mamba_layer_cache( assert self.kda_conv_q is not None assert self.kda_conv_k is not None assert self.kda_conv_v is not None - assert self.kda_qkg_cache is not None - assert self.kda_v_cache is not None - assert self.kda_beta_cache is not None spec_kwargs = { "prev_num_accepted_tokens": self.prev_num_accepted_tokens, "kda_conv_q": self.kda_conv_q[layer_offset], "kda_conv_k": self.kda_conv_k[layer_offset], "kda_conv_v": self.kda_conv_v[layer_offset], - "kda_qkg_cache": self.kda_qkg_cache[layer_offset], - "kda_v_cache": self.kda_v_cache[layer_offset], - "kda_beta_cache": self.kda_beta_cache[layer_offset], } + if self.kda_state_tok is not None: + spec_kwargs["kda_state_tok"] = self.kda_state_tok[layer_offset] + else: + assert self.kda_qkg_cache is not None + assert self.kda_v_cache is not None + assert self.kda_beta_cache is not None + spec_kwargs["kda_qkg_cache"] = self.kda_qkg_cache[layer_offset] + spec_kwargs["kda_v_cache"] = self.kda_v_cache[layer_offset] + spec_kwargs["kda_beta_cache"] = self.kda_beta_cache[layer_offset] if self.mamba_ssm_rand_seed is not None: spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed return PythonMambaCacheManager.SpeculativeState( @@ -4781,4 +4948,5 @@ def shutdown(self): self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None + self.kda_state_tok = None super().shutdown() 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 1f920c7b718c..074c5c19ad16 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 @@ -77,6 +77,7 @@ ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + GPU_LEVEL, AttentionLayerConfig, BatchDesc, BufferConfig, @@ -300,12 +301,16 @@ def _kimi_model_config() -> SimpleNamespace: quant_config=None, sparse_attention_config=None, get_num_mamba_layers=lambda: 2, + # get_kv_cache_manager_cls reads the mapping and spec config (helix check). + mapping=None, + spec_config=None, ) def _capture_kimi_v2_manager_ctor( monkeypatch: pytest.MonkeyPatch, spec_config=None, + model_engine=None, ) -> tuple[tuple, dict]: """Route a Kimi config through _create_kv_cache_manager with an explicit V2 manager and capture the constructor arguments.""" @@ -324,7 +329,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: assert get_kv_cache_manager_cls(model_config, kv_cache_config) is MambaHybridCacheManagerV2 _create_kv_cache_manager( - model_engine=None, + model_engine=model_engine, kv_cache_manager_cls=RecordingV2Manager, mapping=Mapping(world_size=1, tp_size=1, pp_size=1), kv_cache_config=kv_cache_config, @@ -380,6 +385,33 @@ def test_kimi_explicit_v2_manager_enables_kda_replay( assert kwargs["kda_replay_num_spec"] == spec_config.tokens_per_gen_step - 1 assert kwargs["conv_state_layout"] == "q_k_v" + assert "kda_token_states" not in kwargs + + +@pytest.mark.parametrize("backbone_asks", [True, False]) +def test_kimi_explicit_v2_manager_kda_token_states_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_token_states``.""" + monkeypatch.setattr( + "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", + lambda: True, + ) + backbone = SimpleNamespace(kda_token_states=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_token_states", False) is backbone_asks def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout( @@ -2353,6 +2385,7 @@ def _build_v2_hybrid_with_mamba_layer( mamba_n_groups=1, mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, + kda_token_states=False, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2407,6 +2440,7 @@ def _build_v2_hybrid_with_mamba_layer( dtype=dtype, conv_state_layout=conv_state_layout, kda_replay_num_spec=kda_replay_num_spec, + kda_token_states=kda_token_states, ) @@ -3629,6 +3663,54 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): ) 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_token_states + 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_token_states_replace_the_replay_caches(kda_replay_num_spec): + """On the KDA replay path, ``kda_token_states`` 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_token_states=True, + ) + try: + if kda_replay_num_spec is None: + assert getattr(mgr, "kda_state_tok", None) is None + assert not mgr.keeps_kda_token_states + return + assert mgr.keeps_kda_token_states + assert mgr.kda_qkg_cache is None + assert mgr.kda_v_cache is None + assert mgr.kda_beta_cache is None + 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_qkg_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() + assert not states.any() finally: mgr.shutdown() @@ -3676,6 +3758,189 @@ def test_v2_kda_replay_validates_configuration( ) +def _v2_kda_token_state_manager(kda_token_states, 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._generation_kv_capacity_headroom = 1 + mgr._has_cp_helix = False + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + # One 2 MiB grain per state slot, so the pool's slot count follows its grains (a Kimi K3 slot holds 27 MB). + mgr.ssm_bytes = 2 << 20 + mgr.conv_bytes = 64 << 10 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 1 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + # 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_token_states = kda_token_states + # The drafts' records count in each slot's state bytes (_mamba_state_bytes_per_slot): 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.kv_cache_config = KvCacheConfig( + avg_seq_len=64, + enable_block_reuse=enable_block_reuse, + enable_partial_reuse=False, + ) + return mgr + + +@pytest.mark.parametrize( + ("kda_token_states", "enable_block_reuse", "at_floor"), + [(True, False, True), (False, False, False), (False, True, False)], +) +def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( + kda_token_states, enable_block_reuse, at_floor +): + """The KDA drafts' records take memory per SSM slot beside the V2 pools. Without block reuse the SSM pool keeps + only its live floor, the records' bytes come out of the quota, and attention gets the rest; without them (replay + caches only, with or without block reuse) the typical step's ratio sizes the pool.""" + mgr = _v2_kda_token_state_manager(kda_token_states, enable_block_reuse) + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=128 << 20)], + layers=_base_attention_layer_configs(2), + ) + runtime_manager = RuntimeKVCacheManager(mgr._build_cache_config(base_config)) + try: + slots = {} + for stats in runtime_manager.get_storage_statistics(): + sizes = stats.slot_sizes if hasattr(stats, "slot_sizes") else stats.slot_size + role = "ssm" if mgr.ssm_bytes in [int(s) for s in sizes] else "attention" + slots[role] = int(stats.total) + pool_bytes = runtime_manager.get_quota(GPU_LEVEL) + finally: + runtime_manager.shutdown() + + floor = mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots + if at_floor: + # The floor's min-slots constraint is scaled by 1 / max_util_for_resume (0.97). + assert floor <= slots["ssm"] <= floor + 1 + # kda_state_tok, allocated beside the pools: (layers, SSM slots, 3, num_spec, H, K) fp32. + token_state_bytes = ( + mgr.local_num_mamba_layers + * slots["ssm"] + * 3 + * mgr._kda_replay_num_spec + * mgr.ssm_state_shape[0] + * mgr.ssm_state_shape[-1] + * 4 + ) + # The records' bytes come out of the quota and attention takes the rest of it. The runtime rounds the pools' + # quota up to whole 2 MiB grains, so the pools can take back less than one grain of the records' bytes. + assert 128 << 20 <= pool_bytes + token_state_bytes < (128 << 20) + (2 << 20) + else: + assert slots["ssm"] > 10 * floor + + +@pytest.mark.parametrize("off_the_floor", ["block_reuse", "pool_ratio", "no_attention"]) +def test_v2_kda_token_states_refuse_a_pool_off_the_live_floor(off_the_floor): + """kda_state_tok is allocated for every slot the SSM pool keeps, so with block reuse, a pool_ratio, or no + attention layers on the rank (the SSM pool sized by its ratio, its slot count unknown before sizing) the build + refuses the per-token states instead of allocating them outside the quota.""" + mgr = _v2_kda_token_state_manager( + kda_token_states=True, enable_block_reuse=off_the_floor == "block_reuse" + ) + pool_ratio = None + if off_the_floor == "pool_ratio": + pool_ratio = [0.5, 0.5] + elif off_the_floor == "no_attention": + mgr._attention_cache_bytes_per_token = lambda: 0 + with pytest.raises(ValueError, match="need the SSM pool at its live floor"): + mgr._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=128 << 20)], + layers=_base_attention_layer_configs(2), + initial_pool_ratio=pool_ratio, + ) + ) + + +def test_v2_kda_token_states_reserve_quota_before_sizing(): + """At the live floor the drafts' records of every slot the SSM pool keeps come out of the GPU quota before the + pools are sized, and the minimum quota and the quota for max_tokens count them all.""" + mgr = _v2_kda_token_state_manager(kda_token_states=True, enable_block_reuse=False) + plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) + assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 3 + # The pool keeps ceil(3 / max_util_for_resume) = 4 slots, each with the records of 2 drafts (3 x 8 x 256 fp32 + # each). The quotas count the 4 slots' records and the 4th slot's states, which the quotas without the records + # leave out; the build takes the records out of the quota. + reserved = 4 * 2 * 3 * 8 * 256 * 4 + extra_slot_states = plain._mamba_state_bytes_per_slot() + minimum = mgr._minimum_live_gpu_quota() + assert minimum - plain._minimum_live_gpu_quota() == reserved + extra_slot_states + for max_tokens in (0, 4096): + assert ( + mgr._get_quota_from_max_tokens(max_tokens) + - plain._get_quota_from_max_tokens(max_tokens) + == reserved + extra_slot_states + ) + + def build(quota): + return mgr._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + ) + ) + + with pytest.raises(ValueError, match="too small for live recurrent states"): + build(minimum - 1) + assert build(minimum).cache_tiers[0].quota == minimum - reserved + assert build(128 << 20).cache_tiers[0].quota == (128 << 20) - reserved + + +def test_v2_kda_token_states_reserve_the_slots_the_runtime_keeps(): + """The quota checks and the build count the same slots when floor / max_util_for_resume is a whole number in + exact arithmetic: the runtime holds the ratio as a C++ float, so 19 / 0.95 keeps 21 slots, not 20.""" + mgr = _v2_kda_token_state_manager(kda_token_states=True, enable_block_reuse=False) + plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) + for m in (mgr, plain): + m.max_batch_size = 18 + m.kv_cache_config = KvCacheConfig( + avg_seq_len=64, + enable_block_reuse=False, + enable_partial_reuse=False, + max_util_for_resume=0.95, + ) + assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 19 + reserved = 21 * 2 * 3 * 8 * 256 * 4 + minimum = mgr._minimum_live_gpu_quota() + # The 2 slots past the 19 resident ones count with their states as well. + assert ( + minimum - plain._minimum_live_gpu_quota() + == reserved + 2 * plain._mamba_state_bytes_per_slot() + ) + quota = minimum + (64 << 20) + built = mgr._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + max_util_for_resume=0.95, + ) + ) + assert quota - built.cache_tiers[0].quota == reserved + + def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) @@ -3785,6 +4050,73 @@ def test_v2_kda_replay_relocates_live_slot_history(): mgr.shutdown() +@skip_no_cuda +def test_v2_kda_token_states_relocate_with_their_slot(): + """Without the per-draft replay caches, a slot's conv caches, drafts' records and draft count move with it.""" + 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, + kda_token_states=True, + ) + try: + assert mgr.kda_qkg_cache is None + assert mgr.kda_v_cache is None + assert mgr.kda_beta_cache is None + mgr.prev_num_accepted_tokens[:3].copy_( + torch.tensor([10, 20, 30], dtype=torch.int32, device="cuda") + ) + slot_buffers = (mgr.kda_conv_q, mgr.kda_conv_k, mgr.kda_conv_v, mgr.kda_state_tok) + source_rows = [] + for offset, slot_buffer in enumerate(slot_buffers): + assert slot_buffer is not None + slot_buffer.zero_() + slot_buffer[:, 0].fill_(offset * 10 + 1) + slot_buffer[:, 1].fill_(offset * 10 + 2) + source_rows.append((slot_buffer[:, 0].clone(), slot_buffer[:, 1].clone())) + + mgr._relocate_kda_replay_slots([0, 1], [1, 2]) + + assert mgr.prev_num_accepted_tokens[:3].tolist() == [10, 10, 20] + for slot_buffer, (source_zero, source_one) in zip(slot_buffers, source_rows): + torch.testing.assert_close(slot_buffer[:, 1], source_zero, rtol=0, atol=0) + torch.testing.assert_close(slot_buffer[:, 2], source_one, rtol=0, atol=0) + finally: + mgr.shutdown() + # Released with the other replay buffers. + assert mgr.kda_state_tok is None + + +@skip_no_cuda +def test_v2_kda_token_states_count_in_the_per_slot_budget(): + """The capacity math sees the drafts' records: per draft and layer, a slot costs vn, beta * k and the decay, fp32 + [H, K] each.""" + 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, + ) + 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_token_states=True, **kwargs) + try: + num_heads, _, head_dim = with_states.ssm_state_shape + per_token = 2 * 2 * 3 * num_heads * head_dim * 4 + assert with_states._mamba_state_bytes_per_slot() == plain_bytes + per_token + finally: + with_states.shutdown() + + @skip_no_cuda def test_v2_kda_state_index_setup_relocates_generation_history(): class StateCache: @@ -3885,6 +4217,68 @@ def dim_contiguous_conv_cache() -> torch.Tensor: assert mgr.prev_num_accepted_tokens[2] == 5 +def test_v2_kda_token_states_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.local_num_mamba_layers = 2 + mgr.conv_state_shape = [6, 4] + mgr.conv_section_dims = [2, 2, 2] + mgr._request_id_to_state_index = {101: 1, 202: 3} + mgr.prev_num_accepted_tokens = torch.full((4,), 5, dtype=torch.int32) + mgr.all_conv_states = [ + torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4), + torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4) + 1000, + ] + + def dim_contiguous_conv_cache() -> torch.Tensor: + return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) + + mgr.kda_conv_q = dim_contiguous_conv_cache() + mgr.kda_conv_k = dim_contiguous_conv_cache() + mgr.kda_conv_v = dim_contiguous_conv_cache() + mgr.kda_qkg_cache = None + mgr.kda_v_cache = None + mgr.kda_beta_cache = None + mgr.kda_state_tok = torch.full((2, 4, 3, 2, 1, 2), 7.0) + + 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 ( + (mgr.kda_conv_q, 0), + (mgr.kda_conv_k, 2), + (mgr.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 (mgr.kda_state_tok == 7.0).all() + + +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_qkg_cache=t, kda_v_cache=t, kda_beta_cache=t, **conv_caches + ) + with_records = PythonMambaCacheManager.SpeculativeState( + 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 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.