Bound decode slot constraint by token budget

Signed-off-by: Kevin-Li-2025 <2242139@qq.com>

Author: Kevin-Li-2025

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -1052,6 +1052,7 @@ def _build_cache_config(
             constraints=[
                 self._build_concurrent_decode_constraint(
                     max_batch_size=self.max_batch_size,
+                    max_tokens=kv_cache_config.max_tokens,
                     tokens_per_block=tokens_per_block,
                 )
             ],
@@ -1080,10 +1081,12 @@ def _extra_buffers_per_layer(
 
     @staticmethod
     def _build_concurrent_decode_constraint(
-        *, max_batch_size: int, tokens_per_block: int
+        *, max_batch_size: int, max_tokens: Optional[int], tokens_per_block: int
     ) -> BatchDesc:
         assert max_batch_size > 0
         assert tokens_per_block > 0
+        if max_tokens is not None:
+            max_batch_size = max(1, min(max_batch_size, max_tokens // tokens_per_block))
         return BatchDesc(
             [
                 KVCacheDesc(capacity=tokens_per_block, history_length=tokens_per_block - 1)

tests/unittest/_torch/executor/test_per_layer_head_dim.py

@@ -98,6 +98,27 @@ def test_build_cache_config_reserves_concurrent_decode_slots(self):
         self.assertEqual(config.typical_step.kv_caches[0].capacity, 2049)
         self.assertEqual(config.typical_step.kv_caches[0].history_length, 2048)
 
+    def test_build_cache_config_bounds_concurrent_decode_slots_by_max_tokens(self):
+        mgr = KVCacheManagerV2.__new__(KVCacheManagerV2)
+        mgr.kv_cache_type = CacheType.SELF
+        mgr.dtype = DataType.HALF
+        mgr.kv_factor = 2
+        mgr.max_batch_size = 4
+        mgr.max_attention_window_vec = [4, None]
+        mgr.num_local_layers = 2
+        mgr.pp_layers = [0, 1]
+        mgr.num_kv_heads_per_layer = [1, 1]
+        mgr.head_dim_per_layer = [8, 8]
+
+        config = mgr._build_cache_config(
+            KvCacheConfigV2(max_tokens=16, enable_block_reuse=False),
+            tokens_per_block=8,
+            vocab_size=32000,
+            cache_tiers=[GpuCacheTierConfig(quota=1 << 20)],
+        )
+
+        self.assertEqual(len(config.constraints[0].kv_caches), 2)
+
     def test_per_layer_head_dim_wrong_length(self):
         """Test that mismatched list length raises assertion."""
         with self.assertRaises(AssertionError):