[https://nvbugs/6438658][fix] Fix KV cache estimation capacity

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py

@@ -14,16 +14,13 @@
 # limitations under the License.
 
 from collections import defaultdict
+from dataclasses import replace
 from typing import Dict, List, Optional, Tuple
 
 import torch
 
 from tensorrt_llm._torch.pyexecutor import llm_request
-from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import (
-    GPU_LEVEL,
-    BlockReusePolicy,
-    KVCacheManagerV2,
-)
+from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import GPU_LEVEL, KVCacheManagerV2
 from tensorrt_llm._utils import (
     TensorWrapper,
     convert_to_torch_tensor,
@@ -42,13 +39,10 @@
     BatchDesc,
     BufferConfig,
     DataRole,
-    GpuCacheTierConfig,
-    HostCacheTierConfig,
     KVCacheDesc,
     LayerId,
     PageIndexMode,
     ScratchDesc,
-    SwaScratchReuseConfig,
 )
 from tensorrt_llm.runtime.kv_cache_manager_v2 import KVCacheManagerConfig as KVCacheManagerConfigPy
 from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX
@@ -254,9 +248,10 @@ def __init__(
 
         self._init_indexer_dtype(sparse_attn_config)
 
-        # _build_cache_config() needs them to build constraints
+        # _build_cache_config() needs these values to build the typical step.
         self._max_input_len = max_input_len
         self._max_num_tokens = max_num_tokens
+        self._avg_seq_len = kv_cache_config.avg_seq_len
 
         # General initialization
         super().__init__(
@@ -768,16 +763,9 @@ def _get_max_tokens_from_quota(self, quota: int) -> float:
             return float("inf")
         return self._max_num_tokens + (quota - context_limit_quota) / generation_size_per_token
 
-    def _build_cache_config(
-        self,
-        kv_cache_config: KvCacheConfig,
-        *,
-        tokens_per_block: int,
-        vocab_size: int | None,
-        cache_tiers: List[GpuCacheTierConfig | HostCacheTierConfig],
-    ) -> KVCacheManagerConfigPy:
+    def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy:
         """
-        Create the cache manager config for DeepSeek-V4.
+        Add DeepSeek-V4 layers and a typical step to the cache config.
         """
         layers: List[AttentionLayerConfig] = []
         layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {}
@@ -864,19 +852,14 @@ def _add_layer(
         # number of layers in the KVCacheManagerPy
         self._num_manager_layers = len(layers)
 
-        # Build constraints and typical_step for better pool ratio.
+        # Build a typical step for the DeepSeek-V4 pool ratio.
         max_batch_size = self.max_batch_size
         max_seq_len = self.max_seq_len
         max_num_tokens = self._max_num_tokens
         max_draft_len = self._max_draft_len
         typical_step = None
-        constraints = []
-        if kv_cache_config.pool_ratio is None:
-            typical_seq_len = (
-                kv_cache_config.avg_seq_len
-                if kv_cache_config.avg_seq_len is not None
-                else max_seq_len
-            )
+        if config.initial_pool_ratio is None:
+            typical_seq_len = self._avg_seq_len if self._avg_seq_len is not None else max_seq_len
             if typical_seq_len > max_seq_len:
                 raise ValueError(
                     f"kv_cache_config.avg_seq_len ({typical_seq_len}) must be less than or "
@@ -903,50 +886,10 @@ def _add_layer(
                 * (max_batch_size - 1),
             )
 
-            # Constraint 1: cuda graph generation warmup — one decode request that has
-            # accumulated to the tail of max_seq_len. Using history_length=max_seq_len-1
-            # (instead of 0) lets SWA / SSM pools collapse to their windowed working set,
-            # while full-cache pools still need max_seq_len/tokens_per_block blocks
-            # because they don't age.
-            constraints.append(
-                BatchDesc([KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1)])
-            )
-
-            # Constraint 2: general / chunked-prefill warmup — one fresh context request
-            # at max_num_tokens (the per-iteration token budget).
-            if max_num_tokens is not None:
-                constraints.append(
-                    BatchDesc(
-                        [
-                            KVCacheDesc(
-                                capacity=max_num_tokens + self.num_extra_kv_tokens, history_length=0
-                            )
-                        ]
-                    )
-                )
-
-        scratch_reuse_config = None
-        if self.enable_swa_scratch_reuse:
-            # Context requests will allocate num_extra_kv_tokens tokens for spec decoding.
-            # Cache manager should not take them into account when calculating scratch range.
-            # Therefore set max_rewind_len to num_extra_kv_tokens.
-            scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)
-
-        return KVCacheManagerConfigPy(
-            tokens_per_block=tokens_per_block,
-            cache_tiers=cache_tiers,
-            max_util_for_resume=kv_cache_config.max_util_for_resume,
-            enable_partial_reuse=kv_cache_config.enable_partial_reuse,
-            swa_scratch_reuse=scratch_reuse_config,
-            commit_min_snapshot=(
-                kv_cache_config.enable_block_reuse
-                and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE
-            ),
+        return replace(
+            config,
             layers=layers,
             typical_step=typical_step,
-            constraints=constraints,
-            enable_stats=self.enable_stats,
-            initial_pool_ratio=kv_cache_config.pool_ratio,
         )
 
     def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None:

tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py

@@ -168,7 +168,7 @@ def __init__(
             disable_index_value_layer_ids = list(sparse_layer_ids)
 
         # Must be set BEFORE super().__init__ — the base
-        # ``_build_cache_config`` invokes ``_extra_buffers_per_layer``
+        # ``_build_base_config`` invokes ``_extra_buffers_per_layer``
         # which reads these attributes.
         self.sparse_layer_ids = sorted(int(i) for i in sparse_layer_ids)
         self.disable_index_value_layer_ids = set(int(i) for i in disable_index_value_layer_ids)
@@ -198,7 +198,7 @@ def _extra_buffers_per_layer(self, *, tokens_per_block):
 
         ``size`` is bytes per **block**: ``1 * sparse_index_dim *
         elem_bytes * tokens_per_block``. Keyed by **local** layer id —
-        the base ``_build_cache_config`` iterates local ids, so keying
+        the base ``_build_base_config`` iterates local ids, so keying
         by global ids would silently skip registration on non-trivial
         PP ranks.
         """

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -17,6 +17,7 @@
 import os
 import sys
 from collections import OrderedDict, defaultdict
+from dataclasses import replace
 from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -45,6 +46,7 @@
     GPU_LEVEL,
     AttentionLayerConfig,
     AttnLifeCycle,
+    BatchDesc,
     BufferConfig,
     CacheLevel,
     CacheTierConfig,
@@ -53,6 +55,7 @@
     DiskCacheTierConfig,
     GpuCacheTierConfig,
     HostCacheTierConfig,
+    KVCacheDesc,
     KVCacheEventManager,
     KVCacheIterationStatsDelta,
     LayerId,
@@ -121,8 +124,8 @@ class Role:
     # Sparse-attention per-layer index-K cache (MiniMax-M3 and similar
     # sparse-block-selection backends). Registered as a native V2
     # BufferConfig on sparse layers via the extra_buffers_per_layer hook on
-    # _build_cache_config, so allocation, free, slot reuse, and prefix
-    # reuse share the lifecycle of the main K/V buffers for the same layer.
+    # _build_base_config, so allocation, free, slot reuse, and prefix reuse
+    # share the lifecycle of the main K/V buffers for the same layer.
     INDEX_KEY = DataRole("index_key")
     ALL = DataRole("all")
 
@@ -797,6 +800,7 @@ def append_to_kv_heads_per_layer(
 
         self.is_vswa = len(set(self.max_attention_window_vec)) > 1
 
+        max_util_for_resume = kv_cache_config.max_util_for_resume
         quota = sys.maxsize
         if (
             kv_cache_config.max_gpu_total_bytes is not None
@@ -808,7 +812,7 @@ def append_to_kv_heads_per_layer(
             quota_from_max_tokens = int(
                 math.ceil(
                     self._get_quota_from_max_tokens(kv_cache_config.max_tokens)
-                    / kv_cache_config.max_util_for_resume
+                    / max_util_for_resume
                 )
             )
             quota = min(quota, quota_from_max_tokens)
@@ -822,13 +826,13 @@ def append_to_kv_heads_per_layer(
             "Quota not set. Check kv_cache_config.max_tokens or kv_cache_config.max_gpu_total_bytes"
         )
 
-        # Sync KV cache token capacity across ranks so all ranks allocate
-        # the same number of tokens and the scheduler produces identical
-        # batches.  Normalize to token count before the allreduce because
-        # bytes_per_token varies across PP ranks (different local layers).
+        # Sync resumable token capacity across ranks so the scheduler produces
+        # identical batches. Normalize to tokens because cache costs vary
+        # across PP ranks, including fixed per-rank costs.
         if mapping.world_size > 1:
             dist = Distributed.get(mapping)
-            max_tokens = self._get_max_tokens_from_quota(quota)
+            resumable_quota = int(quota * max_util_for_resume)
+            max_tokens = self._get_max_tokens_from_quota(resumable_quota)
             max_tokens = dist.allreduce(max_tokens, op=ReduceOp.MIN)
             # inf max_tokens means all layers are SWA and every rank quota can
             # fit all SWA fixed cache.
@@ -837,7 +841,10 @@ def append_to_kv_heads_per_layer(
                 # token↔quota round-trip is not identity when SWA layers
                 # dominate (full_attn_size_per_token==0), so clamp to guard
                 # against a bogus inflation (nvbugs/6418103).
-                quota = min(quota, self._get_quota_from_max_tokens(max_tokens))
+                synced_quota = int(
+                    math.ceil(self._get_quota_from_max_tokens(max_tokens) / max_util_for_resume)
+                )
+                quota = min(quota, synced_quota)
 
         logger.info(f"KV cache manager v2 device quota set to {quota / (1 << 30)}GiB")
 
@@ -891,12 +898,12 @@ def append_to_kv_heads_per_layer(
 
         self.vocab_size = vocab_size
 
-        config = self._build_cache_config(
+        config = self._build_base_config(
             kv_cache_config,
             tokens_per_block=tokens_per_block,
-            vocab_size=vocab_size,
             cache_tiers=cache_tiers,
         )
+        config = self._build_cache_config(config)
 
         self.kv_cache_manager_py_config = config
 
@@ -910,12 +917,7 @@ def append_to_kv_heads_per_layer(
                     "Retrying without host cache tier."
                 )
                 cache_tiers_gpu_only = [t for t in cache_tiers if isinstance(t, GpuCacheTierConfig)]
-                config = self._build_cache_config(
-                    kv_cache_config,
-                    tokens_per_block=tokens_per_block,
-                    vocab_size=vocab_size,
-                    cache_tiers=cache_tiers_gpu_only,
-                )
+                config = replace(config, cache_tiers=cache_tiers_gpu_only)
                 cache_tiers = cache_tiers_gpu_only
                 self.kv_cache_manager_py_config = config
                 self.impl = KVCacheManagerPy(config, event_manager=self.event_manager)
@@ -1489,17 +1491,52 @@ def _copy_batch_block_offsets_per_layer(
             non_blocking=True,
         )
 
-    def _build_cache_config(
+    def _build_base_config(
         self,
         kv_cache_config: KvCacheConfig,
         *,
         tokens_per_block: int,
-        vocab_size: int | None,
         cache_tiers: List[CacheTierConfig],
     ) -> KVCacheManagerConfigPy:
-        # Kept in the virtual method contract for cache-manager subclasses.
-        # The generic C++ config no longer stores the vocabulary size.
-        del vocab_size
+        """Build the general cache configuration used by most models.
+
+        Models that need a custom configuration should subclass
+        :class:`KVCacheManagerV2` and override :meth:`_build_cache_config` to
+        update the necessary fields.
+        """
+        scratch_reuse_config = None
+        if self.enable_swa_scratch_reuse:
+            # Context requests allocate num_extra_kv_tokens for spec decoding.
+            # They should not count toward the scratch range.
+            scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)
+
+        constraints = []
+        if kv_cache_config.pool_ratio is None:
+            # CUDA graph generation warmup uses one request at max_seq_len and
+            # enough minimal decode requests to fill max_batch_size.
+            min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
+            constraints.append(
+                BatchDesc(
+                    [KVCacheDesc(capacity=self.max_seq_len, history_length=self.max_seq_len - 1)]
+                    + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)]
+                    * (self.max_batch_size - 1)
+                )
+            )
+
+            # General and chunked-prefill warmup uses one fresh context request
+            # at the per-iteration token budget.
+            if self.max_num_tokens is not None:
+                constraints.append(
+                    BatchDesc(
+                        [
+                            KVCacheDesc(
+                                capacity=self.max_num_tokens + self.num_extra_kv_tokens,
+                                history_length=0,
+                            )
+                        ]
+                    )
+                )
+
         buffer_type = [Role.KEY]
         if self.kv_cache_type != CacheTypeCpp.SELFKONLY:
             buffer_type.append(Role.VALUE)
@@ -1512,12 +1549,6 @@ def _build_cache_config(
             if self.kv_cache_type != CacheTypeCpp.SELFKONLY:
                 buffer_type.append(Role.VALUE_BLOCK_SCALE)
 
-        scratch_reuse_config = None
-        if self.enable_swa_scratch_reuse:
-            # Context requests allocate num_extra_kv_tokens for spec decoding.
-            # They should not count toward the scratch range.
-            scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)
-
         # Subclasses (e.g. MiniMax-M3 sparse cache) can register additional
         # per-layer BufferConfig entries — for example a sparse index-K
         # buffer — without overriding the K/V/NVFP4 scale wiring above.
@@ -1559,6 +1590,9 @@ def _build_cache_config(
         return KVCacheManagerConfigPy(
             tokens_per_block=tokens_per_block,
             cache_tiers=cache_tiers,
+            layers=layer_configs,
+            typical_step=None,
+            constraints=constraints,
             max_util_for_resume=kv_cache_config.max_util_for_resume,
             enable_partial_reuse=kv_cache_config.enable_partial_reuse,
             enable_stats=self.enable_stats,
@@ -1568,9 +1602,12 @@ def _build_cache_config(
                 and self.block_reuse_policy != BlockReusePolicy.ALL_REUSABLE
             ),
             initial_pool_ratio=kv_cache_config.pool_ratio,
-            layers=layer_configs,
         )
 
+    def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy:
+        """Customize the general cache config for a specialized cache manager."""
+        return config
+
     def _extra_buffers_per_layer(
         self, *, tokens_per_block: int
     ) -> Optional[dict[int, List[BufferConfig]]]:
@@ -1582,7 +1619,7 @@ def _extra_buffers_per_layer(
         a sparse index-K buffer for each sparse local layer. Each
         ``BufferConfig.size`` is interpreted as bytes per block (i.e.,
         ``bytes_per_token * tokens_per_block``), matching the standard
-        buffers built in :meth:`_build_cache_config`. The block storage
+        buffers built in :meth:`_build_base_config`. The block storage
         groups buffers by lifecycle and size with an opaque role key, so
         new roles do not require C++ changes.
         """

tensorrt_llm/llmapi/llm_args.py

@@ -3528,7 +3528,7 @@ class KvCacheConfig(StrictBaseModel, PybindMirror):
 
     max_util_for_resume: float = Field(
         default=0.95,
-        ge=0,
+        gt=0,
         le=1,
         status="prototype",
         description=
@@ -3677,14 +3677,6 @@ def validate_max_attention_window(cls, v: Optional[List[int]]):
                 )
         return v
 
-    @field_validator('max_util_for_resume')
-    @classmethod
-    def validate_max_util_for_resume(cls, v: float):
-        if not 0 <= v <= 1:
-            raise ValueError(
-                "kv_cache_config.max_util_for_resume must be between 0 and 1")
-        return v
-
     @field_validator('pool_ratio')
     @classmethod
     def validate_pool_ratio(cls, v: Optional[List[float]]):