From ef0361a001bdccf67345eb417cbe4c164ec279e4 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Tue, 15 Sep 2026 10:25:21 +0000
Subject: [PATCH 01/11] [None][fix] Respect KVCM V2 initialization and warmup
 budgets

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../kv_cache/kv_cache_manager_v2.py           |  17 ++-
 .../_torch/pyexecutor/model_engine.py         |  33 +++--
 .../defs/accuracy/test_llm_api_pytorch.py     |   3 +-
 .../test_llm_api_pytorch_multimodal.py        |   2 +-
 .../kv_cache/test_kv_cache_manager_v2.py      | 129 ++++++++++++++++++
 tests/unittest/grpc/smg/test_smg.py           |   2 +-
 6 files changed, 173 insertions(+), 13 deletions(-)

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 0774d0affc9f..83673d270432 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
@@ -2743,12 +2743,25 @@ def _build_base_config(
                 min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
                 # Model one request at max_seq_len plus minimal decode requests
                 # to fill constraint_batch_size.
+                generation_capacity = self.max_seq_len
+                if type(self) is KVCacheManagerV2 and all(
+                    window is None for window in self.max_attention_window_vec
+                ):
+                    # Full attention has one lifecycle, so its pool already receives
+                    # the available quota. A max_seq_len floor would silently grow
+                    # that quota when the model's context limit cannot fit in memory.
+                    # Keep the minimum batch floor; CUDA graph warmup queries the
+                    # allocated pool after reserving its short decode requests to
+                    # determine how long the remaining request can actually be.
+                    generation_capacity = min_decode_capacity
+                # Other layouts need the full-length constraint to distribute
+                # capacity across their distinct attention/recurrent pools.
                 constraints.append(
                     BatchDesc(
                         [
                             KVCacheDesc(
-                                capacity=self.max_seq_len,
-                                history_length=self.max_seq_len - 1,
+                                capacity=generation_capacity,
+                                history_length=generation_capacity - 1,
                             )
                         ]
                         + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)]
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 0e8cf9a4d356..adca872a952d 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3315,19 +3315,36 @@ def free_warmup_requests() -> None:
         # Use max_draft_loop_tokens for capacity estimation to account
         # for the actual KV reservation per request.
         _kv_draft = self.max_draft_loop_tokens
-        available_tokens = kv_cache_manager.get_num_available_tokens(
-            token_num_upper_bound=max_seq_len,
-            batch_size=batch_size,
-            max_num_draft_tokens=_kv_draft)
 
-        # Also consider draft KV cache capacity when it exists
-        if draft_kv_cache_manager is not None:
-            draft_available_tokens = draft_kv_cache_manager.get_num_available_tokens(
-                batch_size=batch_size,
+        def get_available_tokens(manager):
+            capacity_batch_size = batch_size
+            if type(manager) is KVCacheManagerV2 and all(
+                    window is None
+                    for window in manager.max_attention_window_vec):
+                # The capacity query reserves one page for each other sequence.
+                # Full attention has one lifecycle, so use the actual occupied
+                # pages, including multi-page draft dummies and the guard page.
+                capacity_batch_size = 1 + sum(
+                    int(cache.num_blocks)
+                    for cache in manager.kv_cache_map.values())
+            return manager.get_num_available_tokens(
                 token_num_upper_bound=max_seq_len,
+                batch_size=capacity_batch_size,
                 max_num_draft_tokens=_kv_draft)
+
+        available_tokens = get_available_tokens(kv_cache_manager)
+
+        # Also consider draft KV cache capacity when it exists
+        if draft_kv_cache_manager is not None:
+            draft_available_tokens = get_available_tokens(
+                draft_kv_cache_manager)
             available_tokens = min(available_tokens, draft_available_tokens)
 
+        if isinstance(kv_cache_manager, KVCacheManagerV2):
+            # V2's generation dummy reserves one more token after allocating
+            # its input. Leave room for that token in both target and draft KV.
+            available_tokens -= 1
+
         token_num = max(
             ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1,
             min(
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 34688c88513e..0f4a7902a0e4 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -6028,7 +6028,8 @@ class TestSeedOss_36B(LlmapiAccuracyTestHarness):
     @pytest.mark.timeout(14400)
     @pytest.mark.skip_less_device_memory(140000)
     def test_auto_dtype(self):
-        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8)
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8,
+                                        use_kv_cache_manager_v2=True)
         chat_template_kwargs = dict(thinking_budget=-1)
 
         with LLM(self.MODEL_PATH, kv_cache_config=kv_cache_config) as llm:
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py b/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
index 15f437d4f8b0..02591600143d 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
@@ -639,7 +639,7 @@ class TestMistralSmall24B(LlmapiAccuracyTestHarness):
         ids=["forced_chunked_prefill"],
     )
     def test_auto_dtype(self, max_num_tokens):
-        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75)
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75, use_kv_cache_manager_v2=True)
         with LLM(
             self.MODEL_PATH,
             kv_cache_config=kv_cache_config,
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index 55c917fd6f9e..531b813c704d 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -35,6 +35,8 @@
     _update_kv_cache_draft_token_location,
 )
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState
+from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine
+from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm.bindings import DataType, SamplingConfig
 from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE
@@ -835,6 +837,7 @@ def test_avg_seq_len_builds_warmup_constraints() -> None:
         max_seq_len=1024,
         max_num_tokens=2048,
         max_draft_len=2,
+        max_attention_window_vec=[None, 256],
     )
 
     assert config.typical_step == BatchDesc(
@@ -853,6 +856,132 @@ def test_avg_seq_len_builds_warmup_constraints() -> None:
     ]
 
 
+@pytest.fixture(params=[0, 16], ids=["decode", "speculative"])
+def _budget_warmup_spec_config(request: pytest.FixtureRequest):
+    return (
+        Eagle3DecodingConfig(max_draft_len=request.param, speculative_model="dummy")
+        if request.param
+        else None
+    )
+
+
+@pytest.fixture(params=[False, True], ids=["no_guard", "guard"])
+def _budget_warmup_guard_page(request: pytest.FixtureRequest, monkeypatch):
+    if request.param:
+        monkeypatch.setenv("TRTLLM_KV_GUARD_PAGE", "1")
+    else:
+        monkeypatch.delenv("TRTLLM_KV_GUARD_PAGE", raising=False)
+    return request.param
+
+
+@pytest.fixture(params=["bytes", "tokens"])
+def _budget_limited_full_attention_manager(
+    request: pytest.FixtureRequest,
+    _budget_warmup_spec_config,
+    _budget_warmup_guard_page,
+):
+    if not torch.cuda.is_available():
+        pytest.skip("requires CUDA")
+    init_cuda_once()
+    budget = {"max_gpu_total_bytes": 17 << 20} if request.param == "bytes" else {"max_tokens": 8192}
+    config = KvCacheConfig(
+        use_kv_cache_manager_v2=True,
+        enable_block_reuse=False,
+        host_cache_size=0,
+        avg_seq_len=32768,
+        **budget,
+    )
+    manager = KVCacheManagerV2(
+        config,
+        CacheType.SELF,
+        num_layers=2,
+        num_kv_heads=2,
+        head_dim=128,
+        tokens_per_block=32,
+        max_seq_len=131072,
+        max_batch_size=8,
+        max_num_tokens=2048,
+        mapping=Mapping(),
+        dtype=DataType.HALF,
+        spec_config=_budget_warmup_spec_config,
+    )
+    try:
+        assert len(manager.kv_cache_map) == int(_budget_warmup_guard_page)
+        yield manager
+    finally:
+        manager.shutdown()
+
+
+def test_full_attention_warmup_respects_allocated_budget(
+    _budget_limited_full_attention_manager: KVCacheManagerV2,
+) -> None:
+    manager = _budget_limited_full_attention_manager
+    requested_quota = manager.kv_cache_manager_py_config.cache_tiers[0].quota
+    allocated_bytes = manager.impl.get_quota(kv_cache_v2_module.GPU_LEVEL)
+    # These small quotas round up to a 2 MiB GPU allocation grain. The model's
+    # full context would require 256 MiB, far beyond either configured budget.
+    assert 0 < allocated_bytes <= requested_quota + (2 << 20)
+    assert manager.max_num_tokens < manager.max_seq_len < 131072
+
+    requests = manager.add_dummy_requests(
+        [0], token_nums=[manager.max_num_tokens // 2], is_gen=False
+    )
+    assert requests is not None
+    try:
+        cache = manager.kv_cache_map[requests[0].py_request_id]
+        assert cache.resize(manager.max_num_tokens, history_length=0)
+        assert cache.capacity == manager.max_num_tokens
+        cache.suspend()
+        assert cache.resume(torch.cuda.current_stream().cuda_stream)
+    finally:
+        for request in requests:
+            manager.free_resources(request)
+
+
+def test_full_attention_budget_supports_cuda_graph_warmup(
+    _budget_limited_full_attention_manager: KVCacheManagerV2,
+    _budget_warmup_spec_config,
+) -> None:
+    manager = _budget_limited_full_attention_manager
+    # Exercise the real graph request builder: it allocates the short requests,
+    # queries the remaining capacity, then grows the longest generation request.
+    engine = SimpleNamespace()
+    engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
+    engine.spec_config = _budget_warmup_spec_config
+    engine.max_beam_width = 1
+    engine.max_draft_loop_tokens = manager.max_draft_len
+    engine.max_seq_len = 131072
+    engine.use_mrope = False
+    engine.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1
+    engine._get_draft_kv_cache_manager = lambda _: None
+    engine._is_encoder_decoder_model = lambda: False
+    engine.model = SimpleNamespace(
+        model_config=SimpleNamespace(pretrained_config=SimpleNamespace())
+    )
+    resources = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager})
+
+    batch = PyTorchModelEngine._create_cuda_graph_warmup_request(
+        engine, resources, batch_size=manager.max_batch_size, draft_len=manager.max_draft_len
+    )
+    assert batch is not None
+    requests = list(batch.generation_requests)
+    try:
+        assert len(requests) == manager.max_batch_size
+        longest_request = requests[0]
+        cache = manager.kv_cache_map[longest_request.py_request_id]
+        assert cache.capacity > manager.max_num_tokens
+        manager.free_resources(longest_request)
+        requests.remove(longest_request)
+        # Once the longest request finishes, another generation request can
+        # grow across page boundaries into the released capacity.
+        cache = manager.kv_cache_map[requests[0].py_request_id]
+        assert cache.capacity < manager.max_num_tokens
+        assert cache.resize(manager.max_num_tokens, history_length=cache.history_length + 1)
+    finally:
+        for request in requests:
+            manager.free_resources(request)
+
+
 def test_avg_seq_len_updates_typical_step() -> None:
     config = _make_cache_config_for_test(
         KvCacheConfig(avg_seq_len=256),
diff --git a/tests/unittest/grpc/smg/test_smg.py b/tests/unittest/grpc/smg/test_smg.py
index bbf6f5a31717..8cd1909517ef 100644
--- a/tests/unittest/grpc/smg/test_smg.py
+++ b/tests/unittest/grpc/smg/test_smg.py
@@ -842,7 +842,7 @@ def grpc_vlm_service():
     model_path = get_model_path(vlm_model_name)
     llm = LLM(
         model=model_path,
-        kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.6),
+        kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.6, use_kv_cache_manager_v2=True),
         load_format="dummy",
     )
     tokenizer = llm.tokenizer

From 10f389c319c6eac912a226efc12db2d55e503862 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Tue, 15 Sep 2026 10:26:40 +0000
Subject: [PATCH 02/11] [None][test] Preserve full and mixed attention
 constraint coverage

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../executor/kv_cache/test_kv_cache_manager_v2.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index 531b813c704d..25b2f0d3bd7e 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -830,14 +830,19 @@ def test_default_uses_allocator_fallback() -> None:
     assert config.constraints == []
 
 
-def test_avg_seq_len_builds_warmup_constraints() -> None:
+@pytest.mark.parametrize(
+    "attention_windows,generation_capacity",
+    [([None, None], 3), ([None, 256], 1024)],
+    ids=["full_attention", "mixed_attention"],
+)
+def test_avg_seq_len_builds_warmup_constraints(attention_windows, generation_capacity) -> None:
     config = _make_cache_config_for_test(
-        KvCacheConfig(host_cache_size=0, avg_seq_len=1024),
+        KvCacheConfig(use_kv_cache_manager_v2=True, host_cache_size=0, avg_seq_len=1024),
         max_batch_size=3,
         max_seq_len=1024,
         max_num_tokens=2048,
         max_draft_len=2,
-        max_attention_window_vec=[None, 256],
+        max_attention_window_vec=attention_windows,
     )
 
     assert config.typical_step == BatchDesc(
@@ -847,7 +852,7 @@ def test_avg_seq_len_builds_warmup_constraints() -> None:
     assert config.constraints == [
         BatchDesc(
             [
-                KVCacheDesc(capacity=1024, history_length=1023),
+                KVCacheDesc(capacity=generation_capacity, history_length=generation_capacity - 1),
                 KVCacheDesc(capacity=3, history_length=0),
                 KVCacheDesc(capacity=3, history_length=0),
             ]
@@ -919,7 +924,7 @@ def test_full_attention_warmup_respects_allocated_budget(
     requested_quota = manager.kv_cache_manager_py_config.cache_tiers[0].quota
     allocated_bytes = manager.impl.get_quota(kv_cache_v2_module.GPU_LEVEL)
     # These small quotas round up to a 2 MiB GPU allocation grain. The model's
-    # full context would require 256 MiB, far beyond either configured budget.
+    # full context requires at least 256 MiB, far beyond either configured budget.
     assert 0 < allocated_bytes <= requested_quota + (2 << 20)
     assert manager.max_num_tokens < manager.max_seq_len < 131072
 

From 7a241330e4d6bd80ab75cccb037a129e7ebf1915 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Tue, 15 Sep 2026 10:57:59 +0000
Subject: [PATCH 03/11] [None][fix] Limit initialization floor changes to self
 attention

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../pyexecutor/kv_cache/kv_cache_manager_v2.py   | 12 +++++++-----
 .../kv_cache/test_kv_cache_manager_v2.py         | 16 ++++++++++++----
 2 files changed, 19 insertions(+), 9 deletions(-)

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 83673d270432..7a1d71f8a3bc 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
@@ -2744,18 +2744,20 @@ def _build_base_config(
                 # Model one request at max_seq_len plus minimal decode requests
                 # to fill constraint_batch_size.
                 generation_capacity = self.max_seq_len
-                if type(self) is KVCacheManagerV2 and all(
-                    window is None for window in self.max_attention_window_vec
+                if (
+                    type(self) is KVCacheManagerV2
+                    and self.kv_cache_type == CacheTypeCpp.SELF
+                    and all(window is None for window in self.max_attention_window_vec)
                 ):
-                    # Full attention has one lifecycle, so its pool already receives
+                    # Full self attention has one lifecycle, so its pool receives
                     # the available quota. A max_seq_len floor would silently grow
                     # that quota when the model's context limit cannot fit in memory.
                     # Keep the minimum batch floor; CUDA graph warmup queries the
                     # allocated pool after reserving its short decode requests to
                     # determine how long the remaining request can actually be.
                     generation_capacity = min_decode_capacity
-                # Other layouts need the full-length constraint to distribute
-                # capacity across their distinct attention/recurrent pools.
+                # Preserve the full-length constraint for other cache types and
+                # layouts. Cross attention sizes encoder warmup independently.
                 constraints.append(
                     BatchDesc(
                         [
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index 25b2f0d3bd7e..a37043870cf6 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -831,13 +831,21 @@ def test_default_uses_allocator_fallback() -> None:
 
 
 @pytest.mark.parametrize(
-    "attention_windows,generation_capacity",
-    [([None, None], 3), ([None, 256], 1024)],
-    ids=["full_attention", "mixed_attention"],
+    "kv_cache_type,attention_windows,generation_capacity",
+    [
+        (CacheType.SELF, [None, None], 3),
+        (CacheType.SELF, [None, 256], 1024),
+        (CacheType.CROSS, [None, None], 1024),
+        (CacheType.SELFKONLY, [None, None], 1024),
+    ],
+    ids=["full_attention", "mixed_attention", "cross_attention", "key_only"],
 )
-def test_avg_seq_len_builds_warmup_constraints(attention_windows, generation_capacity) -> None:
+def test_avg_seq_len_builds_warmup_constraints(
+    kv_cache_type, attention_windows, generation_capacity
+) -> None:
     config = _make_cache_config_for_test(
         KvCacheConfig(use_kv_cache_manager_v2=True, host_cache_size=0, avg_seq_len=1024),
+        kv_cache_type=kv_cache_type,
         max_batch_size=3,
         max_seq_len=1024,
         max_num_tokens=2048,

From d39cfb01ea6864aa0a950ec7286781a7b43e3c7d Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Tue, 15 Sep 2026 07:30:19 -0700
Subject: [PATCH 04/11] [None][test] Cover multimodal examples with KVCM V2
 budget fix

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py    | 4 ++++
 .../llmapi/apps/_test_trtllm_serve_multimodal_example.py      | 1 +
 2 files changed, 5 insertions(+)

diff --git a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
index 03019f68f0fd..b45d2b9d1716 100644
--- a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
+++ b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import io
 import os
 import tempfile
@@ -43,6 +46,7 @@ def temp_extra_llm_api_options_file(request):
             "kv_cache_config": {
                 "enable_block_reuse": False,
                 "free_gpu_memory_fraction": 0.6,
+                "use_kv_cache_manager_v2": True,
             },
         }
 
diff --git a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
index de7aa43fa301..c5a03e59ef17 100644
--- a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
+++ b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py
@@ -40,6 +40,7 @@ def temp_extra_llm_api_options_file(request):
             "kv_cache_config": {
                 "enable_block_reuse": False,
                 "free_gpu_memory_fraction": 0.6,
+                "use_kv_cache_manager_v2": True,
             },
             "max_num_tokens": 16384,  # for pytorch backend
             # NOTE: This is for video support.

From 72bfd754a560905aa4c727c3af7dc7ee1627db3f Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 16 Sep 2026 05:45:53 +0000
Subject: [PATCH 05/11] [None][fix] Estimate V2 cache constraints and query
 warmup capacity

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../kv_cache/kv_cache_manager_v2.py           | 140 ++++++++++++------
 .../_torch/pyexecutor/model_engine.py         |  67 +++++----
 .../kv_cache/test_kv_cache_manager_v2.py      |  40 +++--
 3 files changed, 157 insertions(+), 90 deletions(-)

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 7a1d71f8a3bc..d0204d7a0416 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
@@ -2724,51 +2724,34 @@ def _build_base_config(
                     * (generation_request_capacity - 1)
                 )
 
-                # CUDA graph generation warmup uses one request at max_seq_len and
-                # enough minimal decode requests to fill the resident capacity.
-                if (
-                    self.max_cuda_graph_batch_size is not None
-                    and self.max_cuda_graph_batch_size > 0
-                    and self.is_estimating_kv_cache
-                    and all(window is None for window in self.max_attention_window_vec)
-                ):
-                    # Estimation graph warmup needs the smaller of the resident
-                    # capacity and the largest captured CUDA graph batch.
-                    constraint_batch_size = min(
-                        generation_request_capacity, self.max_cuda_graph_batch_size
-                    )
-                else:
-                    constraint_batch_size = generation_request_capacity
-                constraint_batch_size = max(1, constraint_batch_size)
                 min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
-                # Model one request at max_seq_len plus minimal decode requests
-                # to fill constraint_batch_size.
-                generation_capacity = self.max_seq_len
-                if (
-                    type(self) is KVCacheManagerV2
-                    and self.kv_cache_type == CacheTypeCpp.SELF
-                    and all(window is None for window in self.max_attention_window_vec)
-                ):
-                    # Full self attention has one lifecycle, so its pool receives
-                    # the available quota. A max_seq_len floor would silently grow
-                    # that quota when the model's context limit cannot fit in memory.
-                    # Keep the minimum batch floor; CUDA graph warmup queries the
-                    # allocated pool after reserving its short decode requests to
-                    # determine how long the remaining request can actually be.
-                    generation_capacity = min_decode_capacity
-                # Preserve the full-length constraint for other cache types and
-                # layouts. Cross attention sizes encoder warmup independently.
-                constraints.append(
-                    BatchDesc(
-                        [
-                            KVCacheDesc(
-                                capacity=generation_capacity,
-                                history_length=generation_capacity - 1,
-                            )
-                        ]
-                        + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)]
-                        * (constraint_batch_size - 1)
-                    )
+                gpu_quota = next(
+                    tier.quota for tier in cache_tiers if isinstance(tier, GpuCacheTierConfig)
+                )
+                # Native minimum slot counts are divided by the resume watermark.
+                # Normalize the quota before estimating a feasible long request.
+                estimate = self._get_max_tokens_from_quota(
+                    int(gpu_quota * kv_cache_config.max_util_for_resume)
+                )
+                generation_capacity = int(min(self.max_seq_len, max(min_decode_capacity, estimate)))
+                # These are independent workloads. Graph warmup shortens its long
+                # request after allocating the short requests; requiring both at
+                # this estimated length would count their memory twice.
+                constraints.extend(
+                    [
+                        BatchDesc(
+                            [
+                                KVCacheDesc(
+                                    capacity=generation_capacity,
+                                    history_length=generation_capacity - 1,
+                                )
+                            ]
+                        ),
+                        BatchDesc(
+                            [KVCacheDesc(capacity=min_decode_capacity, history_length=0)]
+                            * self.max_batch_size
+                        ),
+                    ]
                 )
 
                 # General and chunked-prefill warmup uses one fresh context request
@@ -3188,6 +3171,75 @@ def get_num_available_tokens(
             clamped = min(clamped, self._gpu_max_tokens - extra_tokens)
         return clamped
 
+    def get_warmup_token_capacity(
+        self,
+        *,
+        token_num_upper_bound: int,
+        max_num_draft_tokens: int = 0,
+        draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None,
+    ) -> int:
+        """Return an input length that fits one additional generation dummy.
+
+        Other dummies must already be resident. This query is for exclusive
+        warmup, not concurrent request admission, and never grows the pool.
+        The descriptors match both resize calls in ``add_dummy_requests``:
+        the target keeps generation history, while a coupled draft cache
+        materializes the entire input. All sequence/position limits must be
+        applied to the upper bound before calling this method.
+        """
+        managers = [self]
+        if draft_kv_cache_manager is not None:
+            managers.append(draft_kv_cache_manager)
+        available_slots = []
+        overhead = self.num_extra_kv_tokens + max_num_draft_tokens + 1
+        upper = token_num_upper_bound
+        for manager in managers:
+            statistics = manager._get_storage_statistics(GPU_LEVEL)
+            max_util = manager.kv_cache_manager_py_config.max_util_for_resume
+            if any(stat.total and stat.unavailable / stat.total > max_util for stat in statistics):
+                return 0
+            available_slots.append([stat.available for stat in statistics])
+            if manager._gpu_max_tokens is not None:
+                upper = min(upper, manager._gpu_max_tokens - overhead)
+        minimum_tokens = 2 if self._has_cp_helix else 1
+        if upper < minimum_tokens:
+            return 0
+
+        def fits(tokens: int) -> bool:
+            for index, (manager, available) in enumerate(zip(managers, available_slots)):
+                materialize_history = index != 0
+                descriptor = KVCacheDesc(
+                    capacity=tokens + overhead,
+                    history_length=0 if materialize_history else tokens - 1,
+                )
+                needed = _introspection.compute_slots_for_batch(
+                    manager.impl,
+                    BatchDesc([descriptor]),
+                    manager._ledger_tokens_per_block,
+                    manager.kv_cache_manager_py_config.swa_scratch_reuse
+                    if materialize_history
+                    else None,
+                )
+                if any(required > free for required, free in zip(needed, available)):
+                    return False
+            return True
+
+        if fits(upper):
+            return upper
+        if not fits(minimum_tokens):
+            return 0
+        # SWA retention can oscillate by one slot within a page. Search one
+        # common page phase, checking target and draft at the same input length.
+        page = math.lcm(*(manager._ledger_tokens_per_block for manager in managers))
+        lo, hi = 0, upper // page + 1
+        while hi - lo > 1:
+            mid = (lo + hi) // 2
+            if fits(mid * page):
+                lo = mid
+            else:
+                hi = mid
+        return max(minimum_tokens, lo * page)
+
     def get_num_free_blocks(self) -> int:
         # NOTE This method is used to get the number of blocks in the primary pool not the FREE blocks.
         # However, since we only use this function when the kv cache manager is empty, so it is safe to do so.
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index adca872a952d..fbf457239887 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3316,40 +3316,10 @@ def free_warmup_requests() -> None:
         # for the actual KV reservation per request.
         _kv_draft = self.max_draft_loop_tokens
 
-        def get_available_tokens(manager):
-            capacity_batch_size = batch_size
-            if type(manager) is KVCacheManagerV2 and all(
-                    window is None
-                    for window in manager.max_attention_window_vec):
-                # The capacity query reserves one page for each other sequence.
-                # Full attention has one lifecycle, so use the actual occupied
-                # pages, including multi-page draft dummies and the guard page.
-                capacity_batch_size = 1 + sum(
-                    int(cache.num_blocks)
-                    for cache in manager.kv_cache_map.values())
-            return manager.get_num_available_tokens(
-                token_num_upper_bound=max_seq_len,
-                batch_size=capacity_batch_size,
-                max_num_draft_tokens=_kv_draft)
-
-        available_tokens = get_available_tokens(kv_cache_manager)
-
-        # Also consider draft KV cache capacity when it exists
-        if draft_kv_cache_manager is not None:
-            draft_available_tokens = get_available_tokens(
-                draft_kv_cache_manager)
-            available_tokens = min(available_tokens, draft_available_tokens)
-
-        if isinstance(kv_cache_manager, KVCacheManagerV2):
-            # V2's generation dummy reserves one more token after allocating
-            # its input. Leave room for that token in both target and draft KV.
-            available_tokens -= 1
-
-        token_num = max(
-            ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1,
-            min(
-                available_tokens, max_seq_len - 1 -
-                get_num_extra_kv_tokens(self.spec_config) - _kv_draft))
+        # Apply every static limit before the V2 query. Changing its result
+        # afterward can increase SWA page demand at a different page phase.
+        token_num = max_seq_len - 1 - get_num_extra_kv_tokens(
+            self.spec_config) - _kv_draft
         model_config = self.model.model_config.pretrained_config
         max_position_embeddings = getattr(model_config,
                                           'max_position_embeddings', None)
@@ -3366,6 +3336,35 @@ def get_available_tokens(manager):
         if max_position_embeddings is not None:
             token_num = min(token_num, max_position_embeddings - _kv_draft)
 
+        if isinstance(kv_cache_manager, KVCacheManagerV2):
+            assert draft_kv_cache_manager is None or isinstance(
+                draft_kv_cache_manager, KVCacheManagerV2)
+            token_num = kv_cache_manager.get_warmup_token_capacity(
+                token_num_upper_bound=token_num,
+                max_num_draft_tokens=_kv_draft,
+                draft_kv_cache_manager=draft_kv_cache_manager)
+            minimum_tokens = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1
+            if token_num < minimum_tokens:
+                free_warmup_requests()
+                return None
+        else:
+            available_tokens = kv_cache_manager.get_num_available_tokens(
+                token_num_upper_bound=max_seq_len,
+                batch_size=batch_size,
+                max_num_draft_tokens=_kv_draft)
+            if draft_kv_cache_manager is not None:
+                available_tokens = min(
+                    available_tokens,
+                    draft_kv_cache_manager.get_num_available_tokens(
+                        token_num_upper_bound=max_seq_len,
+                        batch_size=batch_size,
+                        max_num_draft_tokens=_kv_draft))
+            token_num = max(
+                ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1,
+                min(token_num, available_tokens))
+            if max_position_embeddings is not None:
+                token_num = min(token_num, max_position_embeddings - _kv_draft)
+
         token_num = int(
             token_num)  # Ensure int for range() in add_dummy_requests
 
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index a37043870cf6..937df41e8ab1 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -158,6 +158,7 @@ def _make_cache_config_for_test(
     cache_manager.enable_joint_kv_cache_reuse = False
     cache_manager.reuse_match_backoff = 0
     cache_manager.get_layer_bytes_per_token = lambda **_: 128
+    cache_manager._get_max_tokens_from_quota = lambda _: max_seq_len
     # Mirrors __init__: without helix the ledger block equals the physical
     # page (the helper re-enacts construction for partial instances).
     cache_manager._ledger_tokens_per_block = 128
@@ -833,7 +834,7 @@ def test_default_uses_allocator_fallback() -> None:
 @pytest.mark.parametrize(
     "kv_cache_type,attention_windows,generation_capacity",
     [
-        (CacheType.SELF, [None, None], 3),
+        (CacheType.SELF, [None, None], 1024),
         (CacheType.SELF, [None, 256], 1024),
         (CacheType.CROSS, [None, None], 1024),
         (CacheType.SELFKONLY, [None, None], 1024),
@@ -861,10 +862,9 @@ def test_avg_seq_len_builds_warmup_constraints(
         BatchDesc(
             [
                 KVCacheDesc(capacity=generation_capacity, history_length=generation_capacity - 1),
-                KVCacheDesc(capacity=3, history_length=0),
-                KVCacheDesc(capacity=3, history_length=0),
             ]
         ),
+        BatchDesc([KVCacheDesc(capacity=3, history_length=0)] * 3),
         BatchDesc([KVCacheDesc(capacity=2048, history_length=0)]),
     ]
 
@@ -887,11 +887,23 @@ def _budget_warmup_guard_page(request: pytest.FixtureRequest, monkeypatch):
     return request.param
 
 
+@pytest.fixture(params=[False, True], ids=["final", "estimation"])
+def _budget_warmup_phase(request: pytest.FixtureRequest):
+    return request.param
+
+
+@pytest.fixture(params=[None, [256, 256], [256, 131072]], ids=["full", "swa", "mixed"])
+def _budget_warmup_windows(request: pytest.FixtureRequest):
+    return request.param
+
+
 @pytest.fixture(params=["bytes", "tokens"])
-def _budget_limited_full_attention_manager(
+def _budget_limited_attention_manager(
     request: pytest.FixtureRequest,
     _budget_warmup_spec_config,
     _budget_warmup_guard_page,
+    _budget_warmup_phase,
+    _budget_warmup_windows,
 ):
     if not torch.cuda.is_available():
         pytest.skip("requires CUDA")
@@ -902,6 +914,7 @@ def _budget_limited_full_attention_manager(
         enable_block_reuse=False,
         host_cache_size=0,
         avg_seq_len=32768,
+        max_attention_window=_budget_warmup_windows,
         **budget,
     )
     manager = KVCacheManagerV2(
@@ -917,6 +930,7 @@ def _budget_limited_full_attention_manager(
         mapping=Mapping(),
         dtype=DataType.HALF,
         spec_config=_budget_warmup_spec_config,
+        is_estimating_kv_cache=_budget_warmup_phase,
     )
     try:
         assert len(manager.kv_cache_map) == int(_budget_warmup_guard_page)
@@ -925,16 +939,18 @@ def _budget_limited_full_attention_manager(
         manager.shutdown()
 
 
-def test_full_attention_warmup_respects_allocated_budget(
-    _budget_limited_full_attention_manager: KVCacheManagerV2,
+def test_attention_warmup_respects_allocated_budget(
+    _budget_limited_attention_manager: KVCacheManagerV2,
 ) -> None:
-    manager = _budget_limited_full_attention_manager
+    manager = _budget_limited_attention_manager
     requested_quota = manager.kv_cache_manager_py_config.cache_tiers[0].quota
     allocated_bytes = manager.impl.get_quota(kv_cache_v2_module.GPU_LEVEL)
     # These small quotas round up to a 2 MiB GPU allocation grain. The model's
     # full context requires at least 256 MiB, far beyond either configured budget.
     assert 0 < allocated_bytes <= requested_quota + (2 << 20)
-    assert manager.max_num_tokens < manager.max_seq_len < 131072
+    assert manager.max_num_tokens < manager.max_seq_len <= 131072
+    if any(window is None for window in manager.max_attention_window_vec):
+        assert manager.max_seq_len < 131072
 
     requests = manager.add_dummy_requests(
         [0], token_nums=[manager.max_num_tokens // 2], is_gen=False
@@ -951,11 +967,11 @@ def test_full_attention_warmup_respects_allocated_budget(
             manager.free_resources(request)
 
 
-def test_full_attention_budget_supports_cuda_graph_warmup(
-    _budget_limited_full_attention_manager: KVCacheManagerV2,
+def test_attention_budget_supports_cuda_graph_warmup(
+    _budget_limited_attention_manager: KVCacheManagerV2,
     _budget_warmup_spec_config,
 ) -> None:
-    manager = _budget_limited_full_attention_manager
+    manager = _budget_limited_attention_manager
     # Exercise the real graph request builder: it allocates the short requests,
     # queries the remaining capacity, then grows the longest generation request.
     engine = SimpleNamespace()
@@ -1246,7 +1262,7 @@ def test_extra_tokens_are_in_context_capacity() -> None:
     )
 
     assert config.typical_step == BatchDesc([KVCacheDesc(capacity=258, history_length=0)])
-    assert config.constraints[1] == BatchDesc([KVCacheDesc(capacity=258, history_length=0)])
+    assert config.constraints[2] == BatchDesc([KVCacheDesc(capacity=258, history_length=0)])
 
 
 def test_try_commit_blocks_commits_partial_block_at_context_end() -> None:

From 8995290f38dfe23d13f988b4b7ac91dae6f3e1e0 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 16 Sep 2026 00:41:59 -0700
Subject: [PATCH 06/11] [None][fix] Keep V2 budget fix focused on existing CI
 failures

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../kv_cache/kv_cache_manager_v2.py           |  69 ----------
 .../_torch/pyexecutor/model_engine.py         | 128 +++++++++---------
 .../kv_cache/test_kv_cache_manager_v2.py      |  93 ++++---------
 3 files changed, 94 insertions(+), 196 deletions(-)

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 d0204d7a0416..535a9228544a 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
@@ -3171,75 +3171,6 @@ def get_num_available_tokens(
             clamped = min(clamped, self._gpu_max_tokens - extra_tokens)
         return clamped
 
-    def get_warmup_token_capacity(
-        self,
-        *,
-        token_num_upper_bound: int,
-        max_num_draft_tokens: int = 0,
-        draft_kv_cache_manager: Optional["KVCacheManagerV2"] = None,
-    ) -> int:
-        """Return an input length that fits one additional generation dummy.
-
-        Other dummies must already be resident. This query is for exclusive
-        warmup, not concurrent request admission, and never grows the pool.
-        The descriptors match both resize calls in ``add_dummy_requests``:
-        the target keeps generation history, while a coupled draft cache
-        materializes the entire input. All sequence/position limits must be
-        applied to the upper bound before calling this method.
-        """
-        managers = [self]
-        if draft_kv_cache_manager is not None:
-            managers.append(draft_kv_cache_manager)
-        available_slots = []
-        overhead = self.num_extra_kv_tokens + max_num_draft_tokens + 1
-        upper = token_num_upper_bound
-        for manager in managers:
-            statistics = manager._get_storage_statistics(GPU_LEVEL)
-            max_util = manager.kv_cache_manager_py_config.max_util_for_resume
-            if any(stat.total and stat.unavailable / stat.total > max_util for stat in statistics):
-                return 0
-            available_slots.append([stat.available for stat in statistics])
-            if manager._gpu_max_tokens is not None:
-                upper = min(upper, manager._gpu_max_tokens - overhead)
-        minimum_tokens = 2 if self._has_cp_helix else 1
-        if upper < minimum_tokens:
-            return 0
-
-        def fits(tokens: int) -> bool:
-            for index, (manager, available) in enumerate(zip(managers, available_slots)):
-                materialize_history = index != 0
-                descriptor = KVCacheDesc(
-                    capacity=tokens + overhead,
-                    history_length=0 if materialize_history else tokens - 1,
-                )
-                needed = _introspection.compute_slots_for_batch(
-                    manager.impl,
-                    BatchDesc([descriptor]),
-                    manager._ledger_tokens_per_block,
-                    manager.kv_cache_manager_py_config.swa_scratch_reuse
-                    if materialize_history
-                    else None,
-                )
-                if any(required > free for required, free in zip(needed, available)):
-                    return False
-            return True
-
-        if fits(upper):
-            return upper
-        if not fits(minimum_tokens):
-            return 0
-        # SWA retention can oscillate by one slot within a page. Search one
-        # common page phase, checking target and draft at the same input length.
-        page = math.lcm(*(manager._ledger_tokens_per_block for manager in managers))
-        lo, hi = 0, upper // page + 1
-        while hi - lo > 1:
-            mid = (lo + hi) // 2
-            if fits(mid * page):
-                lo = mid
-            else:
-                hi = mid
-        return max(minimum_tokens, lo * page)
-
     def get_num_free_blocks(self) -> int:
         # NOTE This method is used to get the number of blocks in the primary pool not the FREE blocks.
         # However, since we only use this function when the kv cache manager is empty, so it is safe to do so.
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index fbf457239887..d09726342932 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3238,6 +3238,73 @@ def _create_cuda_graph_warmup_request(
         if num_mixed_contexts >= batch_size:
             return None
 
+        # Add one dummy request with the maximum possible sequence length.
+        max_seq_len = min(
+            self.max_seq_len if max_seq_len is None else max_seq_len,
+            kv_cache_manager.max_seq_len)
+
+        # Use max_draft_loop_tokens for capacity estimation to account
+        # for the actual KV reservation per request.
+        _kv_draft = self.max_draft_loop_tokens
+
+        # Determine the input bound before allocating the warmup batch.
+        token_num = max_seq_len - 1 - get_num_extra_kv_tokens(
+            self.spec_config) - _kv_draft
+        model_config = self.model.model_config.pretrained_config
+        max_position_embeddings = getattr(model_config,
+                                          'max_position_embeddings', None)
+        if is_enc_dec:
+            # For enc-dec models the engine max_seq_len covers the encoder
+            # sequence, which may exceed the decoder's position table (e.g.
+            # Whisper: 1500 encoder positions vs max_target_positions=448).
+            decoder_position_limit = getattr(model_config,
+                                             'max_target_positions', None)
+            if decoder_position_limit is not None:
+                max_position_embeddings = (
+                    decoder_position_limit if max_position_embeddings is None
+                    else min(max_position_embeddings, decoder_position_limit))
+        if max_position_embeddings is not None:
+            token_num = min(token_num, max_position_embeddings - _kv_draft)
+
+        if isinstance(kv_cache_manager, KVCacheManagerV2):
+            # V2 adds one generation token beyond the draft/extra reservation.
+            # Include that token in the existing query, then return input length.
+            available_tokens = kv_cache_manager.get_num_available_tokens(
+                token_num_upper_bound=token_num + 1,
+                batch_size=batch_size,
+                max_num_draft_tokens=_kv_draft)
+            if draft_kv_cache_manager is not None:
+                available_tokens = min(
+                    available_tokens,
+                    draft_kv_cache_manager.get_num_available_tokens(
+                        token_num_upper_bound=token_num + 1,
+                        batch_size=batch_size,
+                        max_num_draft_tokens=_kv_draft))
+            token_num = min(token_num, available_tokens - 1)
+            minimum_tokens = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1
+            if token_num < minimum_tokens:
+                return None
+        else:
+            available_tokens = kv_cache_manager.get_num_available_tokens(
+                token_num_upper_bound=max_seq_len,
+                batch_size=batch_size,
+                max_num_draft_tokens=_kv_draft)
+            if draft_kv_cache_manager is not None:
+                available_tokens = min(
+                    available_tokens,
+                    draft_kv_cache_manager.get_num_available_tokens(
+                        token_num_upper_bound=max_seq_len,
+                        batch_size=batch_size,
+                        max_num_draft_tokens=_kv_draft))
+            token_num = max(
+                ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1,
+                min(token_num, available_tokens))
+            if max_position_embeddings is not None:
+                token_num = min(token_num, max_position_embeddings - _kv_draft)
+
+        token_num = int(
+            token_num)  # Ensure int for range() in add_dummy_requests
+
         # Add (batch_size - 1) dummy requests with the minimal sequence
         # length. Mixed capture must create its context rows as real context
         # requests; converting generation dummies afterward leaves their
@@ -3307,67 +3374,6 @@ def free_warmup_requests() -> None:
                 if draft_kv_cache_manager is not None:
                     draft_kv_cache_manager.free_resources(r)
 
-        # Add one dummy request with the maximum possible sequence length.
-        max_seq_len = min(
-            self.max_seq_len if max_seq_len is None else max_seq_len,
-            kv_cache_manager.max_seq_len)
-
-        # Use max_draft_loop_tokens for capacity estimation to account
-        # for the actual KV reservation per request.
-        _kv_draft = self.max_draft_loop_tokens
-
-        # Apply every static limit before the V2 query. Changing its result
-        # afterward can increase SWA page demand at a different page phase.
-        token_num = max_seq_len - 1 - get_num_extra_kv_tokens(
-            self.spec_config) - _kv_draft
-        model_config = self.model.model_config.pretrained_config
-        max_position_embeddings = getattr(model_config,
-                                          'max_position_embeddings', None)
-        if is_enc_dec:
-            # For enc-dec models the engine max_seq_len covers the encoder
-            # sequence, which may exceed the decoder's position table (e.g.
-            # Whisper: 1500 encoder positions vs max_target_positions=448).
-            decoder_position_limit = getattr(model_config,
-                                             'max_target_positions', None)
-            if decoder_position_limit is not None:
-                max_position_embeddings = (
-                    decoder_position_limit if max_position_embeddings is None
-                    else min(max_position_embeddings, decoder_position_limit))
-        if max_position_embeddings is not None:
-            token_num = min(token_num, max_position_embeddings - _kv_draft)
-
-        if isinstance(kv_cache_manager, KVCacheManagerV2):
-            assert draft_kv_cache_manager is None or isinstance(
-                draft_kv_cache_manager, KVCacheManagerV2)
-            token_num = kv_cache_manager.get_warmup_token_capacity(
-                token_num_upper_bound=token_num,
-                max_num_draft_tokens=_kv_draft,
-                draft_kv_cache_manager=draft_kv_cache_manager)
-            minimum_tokens = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1
-            if token_num < minimum_tokens:
-                free_warmup_requests()
-                return None
-        else:
-            available_tokens = kv_cache_manager.get_num_available_tokens(
-                token_num_upper_bound=max_seq_len,
-                batch_size=batch_size,
-                max_num_draft_tokens=_kv_draft)
-            if draft_kv_cache_manager is not None:
-                available_tokens = min(
-                    available_tokens,
-                    draft_kv_cache_manager.get_num_available_tokens(
-                        token_num_upper_bound=max_seq_len,
-                        batch_size=batch_size,
-                        max_num_draft_tokens=_kv_draft))
-            token_num = max(
-                ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1,
-                min(token_num, available_tokens))
-            if max_position_embeddings is not None:
-                token_num = min(token_num, max_position_embeddings - _kv_draft)
-
-        token_num = int(
-            token_num)  # Ensure int for range() in add_dummy_requests
-
         max_seq_len_request = kv_cache_manager.add_dummy_requests(
             request_ids=[batch_size - 1],
             token_nums=[token_num],
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index 937df41e8ab1..b8dbb856610c 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -15,6 +15,7 @@
 
 import array
 import os
+from collections.abc import Iterator
 from dataclasses import dataclass, field, replace
 from types import SimpleNamespace
 from unittest.mock import Mock, call, patch
@@ -832,18 +833,11 @@ def test_default_uses_allocator_fallback() -> None:
 
 
 @pytest.mark.parametrize(
-    "kv_cache_type,attention_windows,generation_capacity",
-    [
-        (CacheType.SELF, [None, None], 1024),
-        (CacheType.SELF, [None, 256], 1024),
-        (CacheType.CROSS, [None, None], 1024),
-        (CacheType.SELFKONLY, [None, None], 1024),
-    ],
-    ids=["full_attention", "mixed_attention", "cross_attention", "key_only"],
+    "kv_cache_type",
+    [CacheType.SELF, CacheType.SELFKONLY],
+    ids=["full_attention", "key_only"],
 )
-def test_avg_seq_len_builds_warmup_constraints(
-    kv_cache_type, attention_windows, generation_capacity
-) -> None:
+def test_avg_seq_len_builds_warmup_constraints(kv_cache_type: CacheType) -> None:
     config = _make_cache_config_for_test(
         KvCacheConfig(use_kv_cache_manager_v2=True, host_cache_size=0, avg_seq_len=1024),
         kv_cache_type=kv_cache_type,
@@ -851,7 +845,6 @@ def test_avg_seq_len_builds_warmup_constraints(
         max_seq_len=1024,
         max_num_tokens=2048,
         max_draft_len=2,
-        max_attention_window_vec=attention_windows,
     )
 
     assert config.typical_step == BatchDesc(
@@ -861,7 +854,7 @@ def test_avg_seq_len_builds_warmup_constraints(
     assert config.constraints == [
         BatchDesc(
             [
-                KVCacheDesc(capacity=generation_capacity, history_length=generation_capacity - 1),
+                KVCacheDesc(capacity=1024, history_length=1023),
             ]
         ),
         BatchDesc([KVCacheDesc(capacity=3, history_length=0)] * 3),
@@ -869,52 +862,25 @@ def test_avg_seq_len_builds_warmup_constraints(
     ]
 
 
-@pytest.fixture(params=[0, 16], ids=["decode", "speculative"])
-def _budget_warmup_spec_config(request: pytest.FixtureRequest):
-    return (
-        Eagle3DecodingConfig(max_draft_len=request.param, speculative_model="dummy")
-        if request.param
-        else None
-    )
-
-
-@pytest.fixture(params=[False, True], ids=["no_guard", "guard"])
-def _budget_warmup_guard_page(request: pytest.FixtureRequest, monkeypatch):
-    if request.param:
-        monkeypatch.setenv("TRTLLM_KV_GUARD_PAGE", "1")
-    else:
-        monkeypatch.delenv("TRTLLM_KV_GUARD_PAGE", raising=False)
-    return request.param
-
-
-@pytest.fixture(params=[False, True], ids=["final", "estimation"])
-def _budget_warmup_phase(request: pytest.FixtureRequest):
-    return request.param
-
-
-@pytest.fixture(params=[None, [256, 256], [256, 131072]], ids=["full", "swa", "mixed"])
-def _budget_warmup_windows(request: pytest.FixtureRequest):
-    return request.param
-
-
-@pytest.fixture(params=["bytes", "tokens"])
-def _budget_limited_attention_manager(
+@pytest.fixture(
+    params=[("bytes", False), ("bytes", True), ("tokens", False), ("tokens", True)],
+    ids=["bytes-final", "bytes-estimation", "tokens-final", "tokens-estimation"],
+)
+def _budget_limited_full_attention_manager(
     request: pytest.FixtureRequest,
-    _budget_warmup_spec_config,
-    _budget_warmup_guard_page,
-    _budget_warmup_phase,
-    _budget_warmup_windows,
-):
+    monkeypatch: pytest.MonkeyPatch,
+) -> Iterator[KVCacheManagerV2]:
     if not torch.cuda.is_available():
         pytest.skip("requires CUDA")
     init_cuda_once()
-    budget = {"max_gpu_total_bytes": 17 << 20} if request.param == "bytes" else {"max_tokens": 8192}
+    monkeypatch.delenv("TRTLLM_KV_GUARD_PAGE", raising=False)
+    budget_type, is_estimating = request.param
+    budget = {"max_gpu_total_bytes": 17 << 20} if budget_type == "bytes" else {"max_tokens": 8192}
     config = KvCacheConfig(
         use_kv_cache_manager_v2=True,
         enable_block_reuse=False,
         host_cache_size=0,
         avg_seq_len=32768,
-        max_attention_window=_budget_warmup_windows,
         **budget,
     )
     manager = KVCacheManagerV2(
@@ -929,28 +895,25 @@ def _budget_limited_attention_manager(
         max_num_tokens=2048,
         mapping=Mapping(),
         dtype=DataType.HALF,
-        spec_config=_budget_warmup_spec_config,
-        is_estimating_kv_cache=_budget_warmup_phase,
+        is_estimating_kv_cache=is_estimating,
     )
     try:
-        assert len(manager.kv_cache_map) == int(_budget_warmup_guard_page)
+        assert not manager.kv_cache_map
         yield manager
     finally:
         manager.shutdown()
 
 
-def test_attention_warmup_respects_allocated_budget(
-    _budget_limited_attention_manager: KVCacheManagerV2,
+def test_full_attention_warmup_respects_allocated_budget(
+    _budget_limited_full_attention_manager: KVCacheManagerV2,
 ) -> None:
-    manager = _budget_limited_attention_manager
+    manager = _budget_limited_full_attention_manager
     requested_quota = manager.kv_cache_manager_py_config.cache_tiers[0].quota
     allocated_bytes = manager.impl.get_quota(kv_cache_v2_module.GPU_LEVEL)
     # These small quotas round up to a 2 MiB GPU allocation grain. The model's
     # full context requires at least 256 MiB, far beyond either configured budget.
     assert 0 < allocated_bytes <= requested_quota + (2 << 20)
-    assert manager.max_num_tokens < manager.max_seq_len <= 131072
-    if any(window is None for window in manager.max_attention_window_vec):
-        assert manager.max_seq_len < 131072
+    assert manager.max_num_tokens < manager.max_seq_len < 131072
 
     requests = manager.add_dummy_requests(
         [0], token_nums=[manager.max_num_tokens // 2], is_gen=False
@@ -967,16 +930,14 @@ def test_attention_warmup_respects_allocated_budget(
             manager.free_resources(request)
 
 
-def test_attention_budget_supports_cuda_graph_warmup(
-    _budget_limited_attention_manager: KVCacheManagerV2,
-    _budget_warmup_spec_config,
+def test_full_attention_budget_supports_cuda_graph_warmup(
+    _budget_limited_full_attention_manager: KVCacheManagerV2,
 ) -> None:
-    manager = _budget_limited_attention_manager
-    # Exercise the real graph request builder: it allocates the short requests,
-    # queries the remaining capacity, then grows the longest generation request.
+    manager = _budget_limited_full_attention_manager
+    # Exercise the real graph request builder against the allocated pool budget.
     engine = SimpleNamespace()
     engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
-    engine.spec_config = _budget_warmup_spec_config
+    engine.spec_config = None
     engine.max_beam_width = 1
     engine.max_draft_loop_tokens = manager.max_draft_len
     engine.max_seq_len = 131072

From 51c69ff3ac7588bb9ca136709ea3c303adf9079e Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 16 Sep 2026 00:59:08 -0700
Subject: [PATCH 07/11] [None][fix] Restore warmup capacity query after short
 request allocation

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         | 120 ++++++++----------
 1 file changed, 53 insertions(+), 67 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index d09726342932..0ad222feeb49 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3238,73 +3238,6 @@ def _create_cuda_graph_warmup_request(
         if num_mixed_contexts >= batch_size:
             return None
 
-        # Add one dummy request with the maximum possible sequence length.
-        max_seq_len = min(
-            self.max_seq_len if max_seq_len is None else max_seq_len,
-            kv_cache_manager.max_seq_len)
-
-        # Use max_draft_loop_tokens for capacity estimation to account
-        # for the actual KV reservation per request.
-        _kv_draft = self.max_draft_loop_tokens
-
-        # Determine the input bound before allocating the warmup batch.
-        token_num = max_seq_len - 1 - get_num_extra_kv_tokens(
-            self.spec_config) - _kv_draft
-        model_config = self.model.model_config.pretrained_config
-        max_position_embeddings = getattr(model_config,
-                                          'max_position_embeddings', None)
-        if is_enc_dec:
-            # For enc-dec models the engine max_seq_len covers the encoder
-            # sequence, which may exceed the decoder's position table (e.g.
-            # Whisper: 1500 encoder positions vs max_target_positions=448).
-            decoder_position_limit = getattr(model_config,
-                                             'max_target_positions', None)
-            if decoder_position_limit is not None:
-                max_position_embeddings = (
-                    decoder_position_limit if max_position_embeddings is None
-                    else min(max_position_embeddings, decoder_position_limit))
-        if max_position_embeddings is not None:
-            token_num = min(token_num, max_position_embeddings - _kv_draft)
-
-        if isinstance(kv_cache_manager, KVCacheManagerV2):
-            # V2 adds one generation token beyond the draft/extra reservation.
-            # Include that token in the existing query, then return input length.
-            available_tokens = kv_cache_manager.get_num_available_tokens(
-                token_num_upper_bound=token_num + 1,
-                batch_size=batch_size,
-                max_num_draft_tokens=_kv_draft)
-            if draft_kv_cache_manager is not None:
-                available_tokens = min(
-                    available_tokens,
-                    draft_kv_cache_manager.get_num_available_tokens(
-                        token_num_upper_bound=token_num + 1,
-                        batch_size=batch_size,
-                        max_num_draft_tokens=_kv_draft))
-            token_num = min(token_num, available_tokens - 1)
-            minimum_tokens = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1
-            if token_num < minimum_tokens:
-                return None
-        else:
-            available_tokens = kv_cache_manager.get_num_available_tokens(
-                token_num_upper_bound=max_seq_len,
-                batch_size=batch_size,
-                max_num_draft_tokens=_kv_draft)
-            if draft_kv_cache_manager is not None:
-                available_tokens = min(
-                    available_tokens,
-                    draft_kv_cache_manager.get_num_available_tokens(
-                        token_num_upper_bound=max_seq_len,
-                        batch_size=batch_size,
-                        max_num_draft_tokens=_kv_draft))
-            token_num = max(
-                ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1,
-                min(token_num, available_tokens))
-            if max_position_embeddings is not None:
-                token_num = min(token_num, max_position_embeddings - _kv_draft)
-
-        token_num = int(
-            token_num)  # Ensure int for range() in add_dummy_requests
-
         # Add (batch_size - 1) dummy requests with the minimal sequence
         # length. Mixed capture must create its context rows as real context
         # requests; converting generation dummies afterward leaves their
@@ -3374,6 +3307,59 @@ def free_warmup_requests() -> None:
                 if draft_kv_cache_manager is not None:
                     draft_kv_cache_manager.free_resources(r)
 
+        # Add one dummy request with the maximum possible sequence length.
+        max_seq_len = min(
+            self.max_seq_len if max_seq_len is None else max_seq_len,
+            kv_cache_manager.max_seq_len)
+
+        # Use max_draft_loop_tokens for capacity estimation to account
+        # for the actual KV reservation per request.
+        _kv_draft = self.max_draft_loop_tokens
+        available_tokens = kv_cache_manager.get_num_available_tokens(
+            token_num_upper_bound=max_seq_len,
+            batch_size=batch_size,
+            max_num_draft_tokens=_kv_draft)
+
+        # Also consider draft KV cache capacity when it exists
+        if draft_kv_cache_manager is not None:
+            draft_available_tokens = draft_kv_cache_manager.get_num_available_tokens(
+                batch_size=batch_size,
+                token_num_upper_bound=max_seq_len,
+                max_num_draft_tokens=_kv_draft)
+            available_tokens = min(available_tokens, draft_available_tokens)
+
+        if isinstance(kv_cache_manager, KVCacheManagerV2):
+            # V2 reserves one generation token beyond the draft/extra tokens.
+            available_tokens -= 1
+            minimum_tokens = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1
+            if available_tokens < minimum_tokens:
+                free_warmup_requests()
+                return None
+
+        token_num = max(
+            ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1,
+            min(
+                available_tokens, max_seq_len - 1 -
+                get_num_extra_kv_tokens(self.spec_config) - _kv_draft))
+        model_config = self.model.model_config.pretrained_config
+        max_position_embeddings = getattr(model_config,
+                                          'max_position_embeddings', None)
+        if is_enc_dec:
+            # For enc-dec models the engine max_seq_len covers the encoder
+            # sequence, which may exceed the decoder's position table (e.g.
+            # Whisper: 1500 encoder positions vs max_target_positions=448).
+            decoder_position_limit = getattr(model_config,
+                                             'max_target_positions', None)
+            if decoder_position_limit is not None:
+                max_position_embeddings = (
+                    decoder_position_limit if max_position_embeddings is None
+                    else min(max_position_embeddings, decoder_position_limit))
+        if max_position_embeddings is not None:
+            token_num = min(token_num, max_position_embeddings - _kv_draft)
+
+        token_num = int(
+            token_num)  # Ensure int for range() in add_dummy_requests
+
         max_seq_len_request = kv_cache_manager.add_dummy_requests(
             request_ids=[batch_size - 1],
             token_nums=[token_num],

From 3be8a7435be22d4a333faaaec1d06232b5ec4d04 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 16 Sep 2026 18:18:46 -0700
Subject: [PATCH 08/11] [None][test] Cover insufficient CUDA graph warmup
 capacity

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../kv_cache/test_kv_cache_manager_v2.py      | 31 +++++++++++++++++--
 1 file changed, 28 insertions(+), 3 deletions(-)

diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index b8dbb856610c..b7398633a9d2 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -930,8 +930,10 @@ def test_full_attention_warmup_respects_allocated_budget(
             manager.free_resources(request)
 
 
+@pytest.mark.parametrize("max_seq_len", [None, 1], ids=["sufficient", "insufficient"])
 def test_full_attention_budget_supports_cuda_graph_warmup(
     _budget_limited_full_attention_manager: KVCacheManagerV2,
+    max_seq_len: int | None,
 ) -> None:
     manager = _budget_limited_full_attention_manager
     # Exercise the real graph request builder against the allocated pool budget.
@@ -950,9 +952,32 @@ def test_full_attention_budget_supports_cuda_graph_warmup(
     )
     resources = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager})
 
-    batch = PyTorchModelEngine._create_cuda_graph_warmup_request(
-        engine, resources, batch_size=manager.max_batch_size, draft_len=manager.max_draft_len
-    )
+    with (
+        patch.object(
+            manager, "get_num_available_tokens", wraps=manager.get_num_available_tokens
+        ) as get_available_tokens,
+        patch.object(manager, "free_resources", wraps=manager.free_resources) as free_resources,
+    ):
+        batch = PyTorchModelEngine._create_cuda_graph_warmup_request(
+            engine,
+            resources,
+            batch_size=manager.max_batch_size,
+            draft_len=manager.max_draft_len,
+            max_seq_len=max_seq_len,
+        )
+        if max_seq_len == 1:
+            # The one-token budget cannot hold both the prompt and V2's extra
+            # generation token. Short requests must be allocated, then freed.
+            get_available_tokens.assert_called_once_with(
+                token_num_upper_bound=1,
+                batch_size=manager.max_batch_size,
+                max_num_draft_tokens=manager.max_draft_len,
+            )
+            assert batch is None
+            assert free_resources.call_count == manager.max_batch_size - 1
+            assert not manager.kv_cache_map
+            return
+
     assert batch is not None
     requests = list(batch.generation_requests)
     try:

From 4767601f1aac5339251ec824273bad1d6f2e1e10 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sun, 20 Sep 2026 13:56:09 +0000
Subject: [PATCH 09/11] [None][fix] Keep maximum-length warmup constraints
 specific to DeepSeek V4

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../sparse/deepseek_v4/cache_manager.py       | 34 ++++++++++-
 .../kv_cache/kv_cache_manager_v2.py           | 46 ---------------
 .../kv_cache/mamba_cache_manager.py           |  5 +-
 .../test_deepseek_v4_cache_manager.py         | 57 +++++++++++++++++--
 .../kv_cache/test_kv_cache_manager_v2.py      | 22 ++++---
 5 files changed, 96 insertions(+), 68 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py
index a367efdd011b..862f1d7452f8 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py
@@ -44,8 +44,10 @@
 from tensorrt_llm.runtime import ModelConfig
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     AttentionLayerConfig,
+    BatchDesc,
     BufferConfig,
     DataRole,
+    KVCacheDesc,
     LayerId,
     PageIndexMode,
     ScratchDesc,
@@ -1028,7 +1030,7 @@ def _get_max_tokens_from_quota(self, quota: int) -> float:
 
     def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy:
         """
-        Add DeepSeek-V4 layers to the cache config.
+        Add DeepSeek-V4 layers and warmup constraints to the cache config.
         """
         layers: List[AttentionLayerConfig] = []
         layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {}
@@ -1173,9 +1175,39 @@ def _add_layer(
         # number of layers in the KVCacheManagerPy
         self._num_manager_layers = len(layers)
 
+        constraints = []
+        if config.initial_pool_ratio is None:
+            # DeepSeek-V4's windowed and compressed pools must support both
+            # the longest decode request and the context warmup workload.
+            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)
+                )
+            )
+            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,
+                            )
+                        ]
+                    )
+                )
+
         return replace(
             config,
             layers=layers,
+            constraints=constraints,
         )
 
     def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None:
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 535a9228544a..559b48327d0d 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
@@ -2695,7 +2695,6 @@ def _build_base_config(
             scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)
 
         typical_step = None
-        constraints = []
         if kv_cache_config.pool_ratio is None:
             typical_seq_len = self._get_typical_seq_len(kv_cache_config)
             if typical_seq_len is not None and typical_seq_len > self.max_seq_len:
@@ -2724,50 +2723,6 @@ def _build_base_config(
                     * (generation_request_capacity - 1)
                 )
 
-                min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
-                gpu_quota = next(
-                    tier.quota for tier in cache_tiers if isinstance(tier, GpuCacheTierConfig)
-                )
-                # Native minimum slot counts are divided by the resume watermark.
-                # Normalize the quota before estimating a feasible long request.
-                estimate = self._get_max_tokens_from_quota(
-                    int(gpu_quota * kv_cache_config.max_util_for_resume)
-                )
-                generation_capacity = int(min(self.max_seq_len, max(min_decode_capacity, estimate)))
-                # These are independent workloads. Graph warmup shortens its long
-                # request after allocating the short requests; requiring both at
-                # this estimated length would count their memory twice.
-                constraints.extend(
-                    [
-                        BatchDesc(
-                            [
-                                KVCacheDesc(
-                                    capacity=generation_capacity,
-                                    history_length=generation_capacity - 1,
-                                )
-                            ]
-                        ),
-                        BatchDesc(
-                            [KVCacheDesc(capacity=min_decode_capacity, history_length=0)]
-                            * self.max_batch_size
-                        ),
-                    ]
-                )
-
-                # 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)
@@ -2823,7 +2778,6 @@ def _build_base_config(
             cache_tiers=cache_tiers,
             layers=layer_configs,
             typical_step=typical_step,
-            constraints=constraints,
             max_util_for_resume=kv_cache_config.max_util_for_resume,
             enable_partial_reuse=kv_cache_config.enable_partial_reuse,
             # Keep the lookahead evidence and its backoff in the same tree match.
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 f2e6cb94c79f..3c7c441fc0aa 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py
@@ -4007,9 +4007,8 @@ def _build_cache_config(
         # 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
+        # of sequence length. Speculative decoding inflates the reserved dummy
+        # slots (CUDA-graph padding), so without this 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.
diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
index fda64c2f673d..c165cd67f47a 100644
--- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
+++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
@@ -332,6 +332,7 @@ def _create_deepseek_v4_cache_manager(
         enable_swa_scratch_reuse: bool = True,
         cold_page_codec_provider: object | None = None,
         host_cache_size: int | None = None,
+        avg_seq_len: int | None = None,
     ) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]:
         """Helper to create a DeepseekV4CacheManager for testing."""
 
@@ -355,6 +356,7 @@ def _create_deepseek_v4_cache_manager(
             event_buffer_max_size=0,
             enable_swa_scratch_reuse=enable_swa_scratch_reuse,
             host_cache_size=host_cache_size,
+            avg_seq_len=avg_seq_len,
         )
 
         # Create mapping (single GPU, no parallelism)
@@ -1203,7 +1205,10 @@ def _assert_cache_equal(
                     msg=f"Mismatch for layer {layer_idx}, attention type {attn_type.name} (scales)",
                 )
 
-    def test_max_num_tokens_is_used_by_base_config(self):
+    @pytest.mark.parametrize("avg_seq_len", [None, 256], ids=["default", "explicit_avg"])
+    def test_warmup_constraints_preserve_deepseek_v4_capacity(
+        self, scratch_reuse_enabled: bool, avg_seq_len: int | None
+    ) -> None:
         max_batch_size = 2
         max_seq_len = 1024
         max_input_len = 127
@@ -1216,14 +1221,54 @@ def test_max_num_tokens_is_used_by_base_config(self):
             compress_ratios=[1, 4],
             dtype=DataType.BF16,
             compressor_dtype=DataType.FLOAT,
+            enable_swa_scratch_reuse=scratch_reuse_enabled,
+            avg_seq_len=avg_seq_len,
         )
 
-        assert cache_manager.kv_cache_manager_py_config.typical_step == BatchDesc(
-            [
-                KVCacheDesc(capacity=max_num_tokens, history_length=0),
-                KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1),
+        requests = []
+        try:
+            config = cache_manager.kv_cache_manager_py_config
+            typical_seq_len = max_seq_len if avg_seq_len is None else avg_seq_len
+            assert config.typical_step == BatchDesc(
+                [
+                    KVCacheDesc(capacity=max_num_tokens, history_length=0),
+                    KVCacheDesc(capacity=typical_seq_len, history_length=typical_seq_len - 1),
+                ]
+            )
+            assert config.constraints == [
+                BatchDesc(
+                    [
+                        KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1),
+                        KVCacheDesc(capacity=1, history_length=0),
+                    ]
+                ),
+                BatchDesc([KVCacheDesc(capacity=max_num_tokens, history_length=0)]),
             ]
-        )
+
+            # Both V4 workloads must fit the native pools after the constraints
+            # move to the specialized manager.
+            generation_requests = cache_manager.add_dummy_requests(
+                request_ids=[0, 1], token_nums=[1, max_seq_len - 1], is_gen=True
+            )
+            assert generation_requests is not None
+            requests.extend(generation_requests)
+            long_cache = cache_manager.kv_cache_map[requests[1].py_request_id]
+            assert long_cache.capacity == max_seq_len
+            for request in requests:
+                cache_manager.free_resources(request)
+            requests.clear()
+
+            context_requests = cache_manager.add_dummy_requests(
+                request_ids=[0], token_nums=[max_num_tokens], is_gen=False
+            )
+            assert context_requests is not None
+            requests.extend(context_requests)
+            context_cache = cache_manager.kv_cache_map[requests[0].py_request_id]
+            assert context_cache.capacity == max_num_tokens
+        finally:
+            for request in requests:
+                cache_manager.free_resources(request)
+            cache_manager.shutdown()
 
     def test_indexer_cache_layout_default(self):
         """DeepSeek-V4 defaults to FP4 indexer K cache on Blackwell+."""
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index b7398633a9d2..1b717da7442b 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -159,7 +159,6 @@ def _make_cache_config_for_test(
     cache_manager.enable_joint_kv_cache_reuse = False
     cache_manager.reuse_match_backoff = 0
     cache_manager.get_layer_bytes_per_token = lambda **_: 128
-    cache_manager._get_max_tokens_from_quota = lambda _: max_seq_len
     # Mirrors __init__: without helix the ledger block equals the physical
     # page (the helper re-enacts construction for partial instances).
     cache_manager._ledger_tokens_per_block = 128
@@ -837,10 +836,16 @@ def test_default_uses_allocator_fallback() -> None:
     [CacheType.SELF, CacheType.SELFKONLY],
     ids=["full_attention", "key_only"],
 )
-def test_avg_seq_len_builds_warmup_constraints(kv_cache_type: CacheType) -> None:
+@pytest.mark.parametrize(
+    "max_attention_window_vec", [[None], [128, None]], ids=["uniform", "mixed"]
+)
+def test_avg_seq_len_does_not_require_max_length_warmup(
+    kv_cache_type: CacheType, max_attention_window_vec: list[int | None]
+) -> None:
     config = _make_cache_config_for_test(
         KvCacheConfig(use_kv_cache_manager_v2=True, host_cache_size=0, avg_seq_len=1024),
         kv_cache_type=kv_cache_type,
+        max_attention_window_vec=max_attention_window_vec,
         max_batch_size=3,
         max_seq_len=1024,
         max_num_tokens=2048,
@@ -851,15 +856,7 @@ def test_avg_seq_len_builds_warmup_constraints(kv_cache_type: CacheType) -> None
         [KVCacheDesc(capacity=2048, history_length=0)]
         + [KVCacheDesc(capacity=1024, history_length=1021)] * 2
     )
-    assert config.constraints == [
-        BatchDesc(
-            [
-                KVCacheDesc(capacity=1024, history_length=1023),
-            ]
-        ),
-        BatchDesc([KVCacheDesc(capacity=3, history_length=0)] * 3),
-        BatchDesc([KVCacheDesc(capacity=2048, history_length=0)]),
-    ]
+    assert config.constraints == []
 
 
 @pytest.fixture(
@@ -913,6 +910,7 @@ def test_full_attention_warmup_respects_allocated_budget(
     # These small quotas round up to a 2 MiB GPU allocation grain. The model's
     # full context requires at least 256 MiB, far beyond either configured budget.
     assert 0 < allocated_bytes <= requested_quota + (2 << 20)
+    assert manager.kv_cache_manager_py_config.constraints == []
     assert manager.max_num_tokens < manager.max_seq_len < 131072
 
     requests = manager.add_dummy_requests(
@@ -1248,7 +1246,7 @@ def test_extra_tokens_are_in_context_capacity() -> None:
     )
 
     assert config.typical_step == BatchDesc([KVCacheDesc(capacity=258, history_length=0)])
-    assert config.constraints[2] == BatchDesc([KVCacheDesc(capacity=258, history_length=0)])
+    assert config.constraints == []
 
 
 def test_try_commit_blocks_commits_partial_block_at_context_end() -> None:

From 28d08c0db1cc93262347da47da257498a4e0976d Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sun, 20 Sep 2026 16:36:35 +0000
Subject: [PATCH 10/11] [None][test] Keep KVCM constraint test changes minimal

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../test_deepseek_v4_cache_manager.py         |  57 +------
 .../kv_cache/test_kv_cache_manager_v2.py      | 155 +-----------------
 2 files changed, 8 insertions(+), 204 deletions(-)

diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
index c165cd67f47a..fda64c2f673d 100644
--- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
+++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py
@@ -332,7 +332,6 @@ def _create_deepseek_v4_cache_manager(
         enable_swa_scratch_reuse: bool = True,
         cold_page_codec_provider: object | None = None,
         host_cache_size: int | None = None,
-        avg_seq_len: int | None = None,
     ) -> Tuple[DeepseekV4CacheManager, DeepSeekV4SparseAttentionConfig]:
         """Helper to create a DeepseekV4CacheManager for testing."""
 
@@ -356,7 +355,6 @@ def _create_deepseek_v4_cache_manager(
             event_buffer_max_size=0,
             enable_swa_scratch_reuse=enable_swa_scratch_reuse,
             host_cache_size=host_cache_size,
-            avg_seq_len=avg_seq_len,
         )
 
         # Create mapping (single GPU, no parallelism)
@@ -1205,10 +1203,7 @@ def _assert_cache_equal(
                     msg=f"Mismatch for layer {layer_idx}, attention type {attn_type.name} (scales)",
                 )
 
-    @pytest.mark.parametrize("avg_seq_len", [None, 256], ids=["default", "explicit_avg"])
-    def test_warmup_constraints_preserve_deepseek_v4_capacity(
-        self, scratch_reuse_enabled: bool, avg_seq_len: int | None
-    ) -> None:
+    def test_max_num_tokens_is_used_by_base_config(self):
         max_batch_size = 2
         max_seq_len = 1024
         max_input_len = 127
@@ -1221,54 +1216,14 @@ def test_warmup_constraints_preserve_deepseek_v4_capacity(
             compress_ratios=[1, 4],
             dtype=DataType.BF16,
             compressor_dtype=DataType.FLOAT,
-            enable_swa_scratch_reuse=scratch_reuse_enabled,
-            avg_seq_len=avg_seq_len,
         )
 
-        requests = []
-        try:
-            config = cache_manager.kv_cache_manager_py_config
-            typical_seq_len = max_seq_len if avg_seq_len is None else avg_seq_len
-            assert config.typical_step == BatchDesc(
-                [
-                    KVCacheDesc(capacity=max_num_tokens, history_length=0),
-                    KVCacheDesc(capacity=typical_seq_len, history_length=typical_seq_len - 1),
-                ]
-            )
-            assert config.constraints == [
-                BatchDesc(
-                    [
-                        KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1),
-                        KVCacheDesc(capacity=1, history_length=0),
-                    ]
-                ),
-                BatchDesc([KVCacheDesc(capacity=max_num_tokens, history_length=0)]),
+        assert cache_manager.kv_cache_manager_py_config.typical_step == BatchDesc(
+            [
+                KVCacheDesc(capacity=max_num_tokens, history_length=0),
+                KVCacheDesc(capacity=max_seq_len, history_length=max_seq_len - 1),
             ]
-
-            # Both V4 workloads must fit the native pools after the constraints
-            # move to the specialized manager.
-            generation_requests = cache_manager.add_dummy_requests(
-                request_ids=[0, 1], token_nums=[1, max_seq_len - 1], is_gen=True
-            )
-            assert generation_requests is not None
-            requests.extend(generation_requests)
-            long_cache = cache_manager.kv_cache_map[requests[1].py_request_id]
-            assert long_cache.capacity == max_seq_len
-            for request in requests:
-                cache_manager.free_resources(request)
-            requests.clear()
-
-            context_requests = cache_manager.add_dummy_requests(
-                request_ids=[0], token_nums=[max_num_tokens], is_gen=False
-            )
-            assert context_requests is not None
-            requests.extend(context_requests)
-            context_cache = cache_manager.kv_cache_map[requests[0].py_request_id]
-            assert context_cache.capacity == max_num_tokens
-        finally:
-            for request in requests:
-                cache_manager.free_resources(request)
-            cache_manager.shutdown()
+        )
 
     def test_indexer_cache_layout_default(self):
         """DeepSeek-V4 defaults to FP4 indexer K cache on Blackwell+."""
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index 1b717da7442b..6006ded382a3 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -15,7 +15,6 @@
 
 import array
 import os
-from collections.abc import Iterator
 from dataclasses import dataclass, field, replace
 from types import SimpleNamespace
 from unittest.mock import Mock, call, patch
@@ -36,8 +35,6 @@
     _update_kv_cache_draft_token_location,
 )
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState
-from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine
-from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm.bindings import DataType, SamplingConfig
 from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE
@@ -831,21 +828,9 @@ def test_default_uses_allocator_fallback() -> None:
     assert config.constraints == []
 
 
-@pytest.mark.parametrize(
-    "kv_cache_type",
-    [CacheType.SELF, CacheType.SELFKONLY],
-    ids=["full_attention", "key_only"],
-)
-@pytest.mark.parametrize(
-    "max_attention_window_vec", [[None], [128, None]], ids=["uniform", "mixed"]
-)
-def test_avg_seq_len_does_not_require_max_length_warmup(
-    kv_cache_type: CacheType, max_attention_window_vec: list[int | None]
-) -> None:
+def test_avg_seq_len_does_not_require_max_length_warmup() -> None:
     config = _make_cache_config_for_test(
-        KvCacheConfig(use_kv_cache_manager_v2=True, host_cache_size=0, avg_seq_len=1024),
-        kv_cache_type=kv_cache_type,
-        max_attention_window_vec=max_attention_window_vec,
+        KvCacheConfig(host_cache_size=0, avg_seq_len=1024),
         max_batch_size=3,
         max_seq_len=1024,
         max_num_tokens=2048,
@@ -859,142 +844,6 @@ def test_avg_seq_len_does_not_require_max_length_warmup(
     assert config.constraints == []
 
 
-@pytest.fixture(
-    params=[("bytes", False), ("bytes", True), ("tokens", False), ("tokens", True)],
-    ids=["bytes-final", "bytes-estimation", "tokens-final", "tokens-estimation"],
-)
-def _budget_limited_full_attention_manager(
-    request: pytest.FixtureRequest,
-    monkeypatch: pytest.MonkeyPatch,
-) -> Iterator[KVCacheManagerV2]:
-    if not torch.cuda.is_available():
-        pytest.skip("requires CUDA")
-    init_cuda_once()
-    monkeypatch.delenv("TRTLLM_KV_GUARD_PAGE", raising=False)
-    budget_type, is_estimating = request.param
-    budget = {"max_gpu_total_bytes": 17 << 20} if budget_type == "bytes" else {"max_tokens": 8192}
-    config = KvCacheConfig(
-        use_kv_cache_manager_v2=True,
-        enable_block_reuse=False,
-        host_cache_size=0,
-        avg_seq_len=32768,
-        **budget,
-    )
-    manager = KVCacheManagerV2(
-        config,
-        CacheType.SELF,
-        num_layers=2,
-        num_kv_heads=2,
-        head_dim=128,
-        tokens_per_block=32,
-        max_seq_len=131072,
-        max_batch_size=8,
-        max_num_tokens=2048,
-        mapping=Mapping(),
-        dtype=DataType.HALF,
-        is_estimating_kv_cache=is_estimating,
-    )
-    try:
-        assert not manager.kv_cache_map
-        yield manager
-    finally:
-        manager.shutdown()
-
-
-def test_full_attention_warmup_respects_allocated_budget(
-    _budget_limited_full_attention_manager: KVCacheManagerV2,
-) -> None:
-    manager = _budget_limited_full_attention_manager
-    requested_quota = manager.kv_cache_manager_py_config.cache_tiers[0].quota
-    allocated_bytes = manager.impl.get_quota(kv_cache_v2_module.GPU_LEVEL)
-    # These small quotas round up to a 2 MiB GPU allocation grain. The model's
-    # full context requires at least 256 MiB, far beyond either configured budget.
-    assert 0 < allocated_bytes <= requested_quota + (2 << 20)
-    assert manager.kv_cache_manager_py_config.constraints == []
-    assert manager.max_num_tokens < manager.max_seq_len < 131072
-
-    requests = manager.add_dummy_requests(
-        [0], token_nums=[manager.max_num_tokens // 2], is_gen=False
-    )
-    assert requests is not None
-    try:
-        cache = manager.kv_cache_map[requests[0].py_request_id]
-        assert cache.resize(manager.max_num_tokens, history_length=0)
-        assert cache.capacity == manager.max_num_tokens
-        cache.suspend()
-        assert cache.resume(torch.cuda.current_stream().cuda_stream)
-    finally:
-        for request in requests:
-            manager.free_resources(request)
-
-
-@pytest.mark.parametrize("max_seq_len", [None, 1], ids=["sufficient", "insufficient"])
-def test_full_attention_budget_supports_cuda_graph_warmup(
-    _budget_limited_full_attention_manager: KVCacheManagerV2,
-    max_seq_len: int | None,
-) -> None:
-    manager = _budget_limited_full_attention_manager
-    # Exercise the real graph request builder against the allocated pool budget.
-    engine = SimpleNamespace()
-    engine.kv_cache_manager_key = ResourceManagerType.KV_CACHE_MANAGER
-    engine.spec_config = None
-    engine.max_beam_width = 1
-    engine.max_draft_loop_tokens = manager.max_draft_len
-    engine.max_seq_len = 131072
-    engine.use_mrope = False
-    engine.get_runtime_tokens_per_gen_step = lambda draft_len: draft_len + 1
-    engine._get_draft_kv_cache_manager = lambda _: None
-    engine._is_encoder_decoder_model = lambda: False
-    engine.model = SimpleNamespace(
-        model_config=SimpleNamespace(pretrained_config=SimpleNamespace())
-    )
-    resources = ResourceManager({ResourceManagerType.KV_CACHE_MANAGER: manager})
-
-    with (
-        patch.object(
-            manager, "get_num_available_tokens", wraps=manager.get_num_available_tokens
-        ) as get_available_tokens,
-        patch.object(manager, "free_resources", wraps=manager.free_resources) as free_resources,
-    ):
-        batch = PyTorchModelEngine._create_cuda_graph_warmup_request(
-            engine,
-            resources,
-            batch_size=manager.max_batch_size,
-            draft_len=manager.max_draft_len,
-            max_seq_len=max_seq_len,
-        )
-        if max_seq_len == 1:
-            # The one-token budget cannot hold both the prompt and V2's extra
-            # generation token. Short requests must be allocated, then freed.
-            get_available_tokens.assert_called_once_with(
-                token_num_upper_bound=1,
-                batch_size=manager.max_batch_size,
-                max_num_draft_tokens=manager.max_draft_len,
-            )
-            assert batch is None
-            assert free_resources.call_count == manager.max_batch_size - 1
-            assert not manager.kv_cache_map
-            return
-
-    assert batch is not None
-    requests = list(batch.generation_requests)
-    try:
-        assert len(requests) == manager.max_batch_size
-        longest_request = requests[0]
-        cache = manager.kv_cache_map[longest_request.py_request_id]
-        assert cache.capacity > manager.max_num_tokens
-        manager.free_resources(longest_request)
-        requests.remove(longest_request)
-        # Once the longest request finishes, another generation request can
-        # grow across page boundaries into the released capacity.
-        cache = manager.kv_cache_map[requests[0].py_request_id]
-        assert cache.capacity < manager.max_num_tokens
-        assert cache.resize(manager.max_num_tokens, history_length=cache.history_length + 1)
-    finally:
-        for request in requests:
-            manager.free_resources(request)
-
-
 def test_avg_seq_len_updates_typical_step() -> None:
     config = _make_cache_config_for_test(
         KvCacheConfig(avg_seq_len=256),

From 8b6febffa5fdd0f851c76cc230939dc7e30330b4 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Sun, 20 Sep 2026 16:46:05 +0000
Subject: [PATCH 11/11] [None][fix] Preserve generic context warmup constraints

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../sparse/deepseek_v4/cache_manager.py        | 18 +++---------------
 .../pyexecutor/kv_cache/kv_cache_manager_v2.py | 16 ++++++++++++++++
 .../pyexecutor/kv_cache/mamba_cache_manager.py |  5 +++--
 .../kv_cache/test_kv_cache_manager_v2.py       |  6 +++---
 4 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py
index 862f1d7452f8..3ee6fb5ca226 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py
@@ -1175,10 +1175,10 @@ def _add_layer(
         # number of layers in the KVCacheManagerPy
         self._num_manager_layers = len(layers)
 
-        constraints = []
+        constraints = list(config.constraints)
         if config.initial_pool_ratio is None:
-            # DeepSeek-V4's windowed and compressed pools must support both
-            # the longest decode request and the context warmup workload.
+            # DeepSeek-V4's windowed and compressed pools must also support
+            # the longest decode request alongside the short decode requests.
             min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
             constraints.append(
                 BatchDesc(
@@ -1192,18 +1192,6 @@ def _add_layer(
                     * (self.max_batch_size - 1)
                 )
             )
-            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,
-                            )
-                        ]
-                    )
-                )
-
         return replace(
             config,
             layers=layers,
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 559b48327d0d..35baf5c50fef 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
@@ -2695,6 +2695,7 @@ def _build_base_config(
             scratch_reuse_config = SwaScratchReuseConfig(max_rewind_len=self.num_extra_kv_tokens)
 
         typical_step = None
+        constraints = []
         if kv_cache_config.pool_ratio is None:
             typical_seq_len = self._get_typical_seq_len(kv_cache_config)
             if typical_seq_len is not None and typical_seq_len > self.max_seq_len:
@@ -2723,6 +2724,20 @@ def _build_base_config(
                     * (generation_request_capacity - 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)
@@ -2778,6 +2793,7 @@ def _build_base_config(
             cache_tiers=cache_tiers,
             layers=layer_configs,
             typical_step=typical_step,
+            constraints=constraints,
             max_util_for_resume=kv_cache_config.max_util_for_resume,
             enable_partial_reuse=kv_cache_config.enable_partial_reuse,
             # Keep the lookahead evidence and its backoff in the same tree match.
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 3c7c441fc0aa..f2e6cb94c79f 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py
@@ -4007,8 +4007,9 @@ def _build_cache_config(
         # 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. Speculative decoding inflates the reserved dummy
-        # slots (CUDA-graph padding), so without this explicit floor the SSM
+        # 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.
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
index 6006ded382a3..80e3d00ba9e2 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py
@@ -828,7 +828,7 @@ def test_default_uses_allocator_fallback() -> None:
     assert config.constraints == []
 
 
-def test_avg_seq_len_does_not_require_max_length_warmup() -> None:
+def test_avg_seq_len_builds_context_warmup_constraint() -> None:
     config = _make_cache_config_for_test(
         KvCacheConfig(host_cache_size=0, avg_seq_len=1024),
         max_batch_size=3,
@@ -841,7 +841,7 @@ def test_avg_seq_len_does_not_require_max_length_warmup() -> None:
         [KVCacheDesc(capacity=2048, history_length=0)]
         + [KVCacheDesc(capacity=1024, history_length=1021)] * 2
     )
-    assert config.constraints == []
+    assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])]
 
 
 def test_avg_seq_len_updates_typical_step() -> None:
@@ -1095,7 +1095,7 @@ def test_extra_tokens_are_in_context_capacity() -> None:
     )
 
     assert config.typical_step == BatchDesc([KVCacheDesc(capacity=258, history_length=0)])
-    assert config.constraints == []
+    assert config.constraints == [BatchDesc([KVCacheDesc(capacity=258, history_length=0)])]
 
 
 def test_try_commit_blocks_commits_partial_block_at_context_end() -> None: