From 5bb3aa8b5abf147cd2b41cd080f4875ea9f7a22d 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 1/8] [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           |  19 ++-
 .../_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(+), 15 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 57a3140595ce..2a5e36b6e0cb 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
@@ -2690,15 +2690,26 @@ def _build_base_config(
                     * (self.max_batch_size - 1)
                 )
 
-                # CUDA graph generation warmup uses one request at max_seq_len and
-                # enough minimal decode requests to fill max_batch_size.
                 min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
+                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 318400132d5f..199577ba7874 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3269,19 +3269,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 f6db6870ab5a..0780415a9939 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -6007,7 +6007,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 a437e30ed4c3..ec64b10f590b 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
@@ -34,6 +34,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
@@ -778,6 +780,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(
@@ -796,6 +799,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 49abcc6251461f2ed4d47b158396a8802d7c58e0 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 2/8] [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 ec64b10f590b..5889be6d8a9e 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
@@ -773,14 +773,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(
@@ -790,7 +795,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),
             ]
@@ -862,7 +867,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 10fef55d68cefcd40b8a1214e452f4cf1c621779 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 3/8] [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 2a5e36b6e0cb..f85e74732d6d 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
@@ -2692,18 +2692,20 @@ def _build_base_config(
 
                 min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
                 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 5889be6d8a9e..7e5d644f175c 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
@@ -774,13 +774,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 b595aec36ea6ee191e25e3fdb210aecee2f712a4 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 4/8] [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 f99fae5397c6be30296982cc974f069c64df517c 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 5/8] [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           | 122 ++++++++++++++----
 .../_torch/pyexecutor/model_engine.py         |  67 +++++-----
 .../kv_cache/test_kv_cache_manager_v2.py      |  40 ++++--
 3 files changed, 157 insertions(+), 72 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 f85e74732d6d..9643dcef9bd8 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
@@ -2691,32 +2691,33 @@ def _build_base_config(
                 )
 
                 min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens
-                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)]
-                        * (self.max_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
@@ -3136,6 +3137,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 199577ba7874..ec63b4b9cc69 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3270,40 +3270,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)
@@ -3320,6 +3290,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 7e5d644f175c..c08d2036a051 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
@@ -153,6 +153,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
@@ -776,7 +777,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),
@@ -804,10 +805,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)]),
     ]
 
@@ -830,11 +830,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")
@@ -845,6 +857,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(
@@ -860,6 +873,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)
@@ -868,16 +882,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
@@ -894,11 +910,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()
@@ -1189,7 +1205,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 4c87a8d3ef04e901fef0733270c6b38d7a89431f 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 6/8] [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 9643dcef9bd8..52939b2d6aeb 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
@@ -3137,75 +3137,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 ec63b4b9cc69..c0b72ca69f3a 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3192,6 +3192,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
@@ -3261,67 +3328,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 c08d2036a051..fc2fbfb0c70b 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
@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import array
+from collections.abc import Iterator
 from dataclasses import dataclass, field, replace
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
@@ -775,18 +776,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,
@@ -794,7 +788,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(
@@ -804,7 +797,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),
@@ -812,52 +805,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(
@@ -872,28 +838,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
@@ -910,16 +873,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 9c3c53b7f39d799ba695888f41d1f15f782c2b55 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 7/8] [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 c0b72ca69f3a..96de92940d6d 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -3192,73 +3192,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
@@ -3328,6 +3261,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 40466ac6c044a3ff7e61f40369d5e57eaa4fb030 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 8/8] [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 fc2fbfb0c70b..2e510c810997 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
@@ -873,8 +873,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.
@@ -893,9 +895,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: