[None][feat] KDA verify: optional per-token states in the V2 hybrid cache manager

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 starts from the state after the
+                # accepted drafts (instead of replaying them) asks for the
+                # state after every draft 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 \

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,14 @@ 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 state after every
+        draft of the last verify round (``kda_state_tok``), so that a verify
+        kernel can start the next round from the state after the accepted
+        drafts instead of replaying them."""
+        return False
+
     @abstractmethod
     def get_conv_states(self, layer_idx: int) -> torch.Tensor:
         """Return conv states for specific layer.
@@ -438,6 +449,14 @@ class SpeculativeState(State):
         kda_qkg_cache: torch.Tensor | None = None
         kda_v_cache: torch.Tensor | None = None
         kda_beta_cache: torch.Tensor | None = None
+        # Optional: the state after every draft of the last verify round,
+        # [slots, num_spec, H, V, K] fp32 (the SSM pool holds the state
+        # after the round's first, non-draft token). A verify kernel starts
+        # the next round from entry n - 1 when the round accepted n > 0
+        # drafts (prev_num_accepted_tokens) instead of replaying them from
+        # the caches above; a slot reset for a new request has n = 0, so a
+        # previous owner's entries 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).
@@ -3024,6 +3043,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 +3060,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 state after every
+        # verify token (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 +3443,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 +3477,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
@@ -3504,8 +3532,19 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor:
              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.
+            self.kda_state_tok = torch.zeros(
+                (self.local_num_mamba_layers, cache_size, num_spec,
+                 *self.ssm_state_shape),
+                dtype=torch.float32,
+                device=device,
+            )
+        per_token = (", with per-token states"
+                     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 +3837,18 @@ def _get_pool_roles(self,
     def _max_resident_sequences(self) -> int:
         return self.max_batch_size * self.mapping.pp_size
 
+    def _kda_token_state_bytes_per_slot(self) -> int:
+        """Bytes of ``kda_state_tok`` per SSM slot: an fp32 SSM state per
+        draft of every local Mamba layer, or 0 without per-token states."""
+        if not getattr(self, "_kda_token_states", False):
+            return 0
+        return (self.local_num_mamba_layers * self._kda_replay_num_spec *
+                math.prod(self.ssm_state_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,14 +3944,23 @@ 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, plus the token
+        states of the further slots the live floor keeps
+        (``_live_floor_token_state_slots``)."""
+        extra_slots = self._live_floor_token_state_slots(
+            self.kv_cache_config.max_util_for_resume) - state_slots
+        return (state_slots * self._mamba_state_bytes_per_slot() +
+                max(0, extra_slots) * self._kda_token_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()
         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()
+        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 +3996,76 @@ 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 states 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 states' 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 token states 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 +4074,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 states' 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 per-token KDA states 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 +4127,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 +4157,7 @@ def _build_cache_config(
             ]
         return replace(
             config,
+            cache_tiers=cache_tiers,
             layers=layers,
             typical_step=typical_step,
             constraints=constraints,
@@ -4285,16 +4427,20 @@ 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 = (
+        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,
-        )
+        ]
+        assert all(replay_buffer is not None
+                   for replay_buffer in replay_buffers)
+        # The per-token verify states belong to the slot as well.
+        if self.kda_state_tok is not None:
+            replay_buffers.append(self.kda_state_tok)
         for replay_buffer in replay_buffers:
-            assert replay_buffer is not None
             replay_buffer.index_copy_(
                 1,
                 destination_slots,
@@ -4498,6 +4644,8 @@ def mamba_layer_cache(
             "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]
         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 +4929,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()