From fa5a84a07e79a05664051c00c8fa7fccc093e804 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 9 Sep 2026 03:52:19 +0000
Subject: [PATCH 1/3] [None][fix] Fix KVCM V2 PARD cache accounting

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |  51 +++++---
 .../kv_cache/kv_cache_manager_v2.py           | 114 +++++++++++-------
 .../kv_cache/test_kv_cache_budget_split.py    |  31 +++--
 .../kv_cache/test_kv_cache_estimation.py      |  57 ++++++---
 .../kv_cache/test_kv_cache_manager_v2.py      |  32 ++++-
 5 files changed, 194 insertions(+), 91 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index bc69272e0736..4b3b3dd92e4a 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -762,7 +762,7 @@ def _per_manager_cache_cost(self,
                 tokens_per_block=self._tokens_per_block,
                 max_seq_len=self._max_seq_len,
                 max_batch_size=self._max_batch_size,
-                max_num_tokens=self._max_num_tokens if is_draft else 0,
+                max_num_tokens=getattr(self, "_max_num_tokens", 0),
                 kv_cache_config=kv_cache_config,
                 spec_config=self._speculative_config,
                 is_draft=is_draft,
@@ -797,14 +797,29 @@ def _get_kv_size_per_token(self,
             use_separate_draft_kv_cache=use_separate_draft_kv_cache)
         if self._is_encoder_decoder():
             total += CacheCost.from_raw(self._get_cross_kv_size_per_token())
+        draft_cost = self._get_draft_cache_cost(
+            kv_cache_config,
+            use_separate_draft_kv_cache=use_separate_draft_kv_cache,
+        )
+        if draft_cost is not None:
+            total += draft_cost
+        return total
+
+    def _get_draft_cache_cost(
+        self,
+        kv_cache_config: KvCacheConfig,
+        *,
+        use_separate_draft_kv_cache: bool,
+    ) -> Optional[CacheCost]:
+        """Return the draft manager's standalone cache cost, if it has one."""
         if self._draft_model_engine is not None:
             draft_model_config = self._draft_model_engine.model.model_config
             draft_kv_cache_manager_cls = self._get_model_kv_cache_manager_cls(
                 self._draft_model_engine, kv_cache_config)
-            total += self._per_manager_cache_cost(draft_kv_cache_manager_cls,
-                                                  draft_model_config,
-                                                  kv_cache_config)
-        elif use_separate_draft_kv_cache:
+            return self._per_manager_cache_cost(draft_kv_cache_manager_cls,
+                                                draft_model_config,
+                                                kv_cache_config)
+        if use_separate_draft_kv_cache:
             # One-model draft with separate KV cache layout.
             # Pass num_layers explicitly since the HF config may report a
             # different layer count than what is actually used at runtime
@@ -822,20 +837,19 @@ def _get_kv_size_per_token(self,
                     effective_draft_config,
                     draft_kv_cache_config,
                     is_disagg=self._is_disagg)
-                total += self._per_manager_cache_cost(
-                    draft_kv_cache_manager_cls,
-                    effective_draft_config,
-                    draft_kv_cache_config,
-                    is_draft=True)
+                return self._per_manager_cache_cost(draft_kv_cache_manager_cls,
+                                                    effective_draft_config,
+                                                    draft_kv_cache_config,
+                                                    is_draft=True)
             elif self._mapping.is_last_pp_rank():
                 # EAGLE3/MTP: draft layers only on last PP rank
-                total += self._per_manager_cache_cost(
+                return self._per_manager_cache_cost(
                     self._kv_cache_manager_cls,
                     effective_draft_config,
                     draft_kv_cache_config,
                     num_layers=self._get_num_draft_layers(),
                     is_draft=True)
-        return total
+        return None
 
     def _cal_max_memory(self, peak_memory, total_gpu_memory, fraction,
                         allocated_bytes: int) -> int:
@@ -1679,7 +1693,6 @@ def _get_target_and_draft_cache_costs(
         """Per-manager KV cache costs for target and draft layers."""
         target_kv_cache_config = (kv_cache_config if kv_cache_config is not None
                                   else self._kv_cache_config)
-        total_kv = self._get_kv_size_per_token(target_kv_cache_config)
         use_separate_draft_kv_cache = (
             self._should_create_separate_draft_kv_cache())
         target_kv = self._per_manager_cache_cost(
@@ -1687,8 +1700,14 @@ def _get_target_and_draft_cache_costs(
             self._model_engine.model.model_config,
             target_kv_cache_config,
             use_separate_draft_kv_cache=use_separate_draft_kv_cache)
-        draft_kv = CacheCost(slope=total_kv.slope - target_kv.slope,
-                             intercept=total_kv.intercept - target_kv.intercept)
+        # Estimate the draft component directly so its independently modelled
+        # affine intercept is preserved exactly.
+        draft_kv = self._get_draft_cache_cost(
+            target_kv_cache_config,
+            use_separate_draft_kv_cache=use_separate_draft_kv_cache,
+        )
+        if draft_kv is None:
+            return None
         costs = (target_kv, draft_kv)
         if any(cost.slope < 0 or cost.intercept < 0 or (
                 cost.slope == 0 and cost.intercept == 0) for cost in costs):
@@ -1790,7 +1809,7 @@ def _split_kv_cache_budget_for_draft(
                 raise ValueError(
                     f"KV cache GPU budget ({total_budget / GB:.2f} GiB) is "
                     f"insufficient after the combined fixed cost "
-                    f"({intercept_total / GB:.2f} GiB, e.g. mamba SSM state) "
+                    f"({intercept_total / GB:.2f} GiB, e.g. SWA or mamba state) "
                     f"for target+draft. Increase free_gpu_memory_fraction or "
                     f"max_gpu_total_bytes, or reduce max_batch_size (the fixed "
                     f"cost scales with batch size).")
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 97d117448bb9..a7212d64140a 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
@@ -120,6 +120,11 @@
     "iter_partial_reused_blocks",
     "iter_missed_blocks",
 )
+
+# Every generation step reserves the golden/base token in addition to any
+# speculative draft tokens. Keep this shared with the capacity-growth paths so
+# cache estimation and runtime allocation cannot drift apart.
+BASE_GENERATION_TOKEN_COUNT = 1
 KV_CACHE_ITERATION_STATS_POOL_GROUP_FIELDS = tuple(
     field_name
     for field_name in KV_CACHE_ITERATION_STATS_DELTA_FIELDS
@@ -331,7 +336,7 @@ def _estimate_swa_cache_size(
     *,
     context: bool,
     scratch: bool,
-    generation_capacity_headroom: Optional[int] = None,
+    generation_capacity_headroom: int = BASE_GENERATION_TOKEN_COUNT,
 ) -> tuple[int, int]:
     tokens_per_block = int(tokens_per_block)
     size_per_token = 0
@@ -339,15 +344,12 @@ def _estimate_swa_cache_size(
     scratch_keys = set()
     for layer_size, window_size in zip(layer_sizes, attention_windows):
         if window_size is not None and window_size > 0:
-            if generation_capacity_headroom is None:
-                window_blocks = math.ceil(window_size / tokens_per_block)
-            else:
-                # Match DFlash's retained boundary page and capacity reserved
-                # ahead of committed history for the next draft step.
-                window_blocks = (
-                    math.ceil((window_size + generation_capacity_headroom - 1) / tokens_per_block)
-                    + 1
-                )
+            # Match AttnLifeCycle.get_stale_range(): the live interval contains
+            # window_size + generation_capacity_headroom - 1 tokens. Across all
+            # page offsets, that interval touches at most the count below.
+            window_blocks = (
+                math.ceil((window_size + generation_capacity_headroom - 2) / tokens_per_block) + 1
+            )
             window_tokens = window_blocks * tokens_per_block
             if not context:
                 size_per_request += window_tokens * layer_size
@@ -363,16 +365,30 @@ def _estimate_swa_cache_size(
     return size_per_token, size_per_request
 
 
-def _get_dflash_generation_kv_capacity_headroom(spec_config) -> Optional[int]:
-    """DFlash KV capacity reserved ahead of committed history."""
-    from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode
+def _get_kv_reserve_draft_tokens(spec_config, *, is_draft: bool) -> int:
+    """Return the same speculative KV reserve used by runtime resize paths."""
+    if spec_config is None:
+        return 0
+    reserve = spec_config.max_total_draft_tokens
+    if (
+        is_draft
+        and getattr(spec_config, "use_dynamic_tree", False)
+        and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0
+    ):
+        draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len
+        reserve = max(reserve, draft_loop_tokens)
+    return reserve
 
-    if spec_config is None or spec_config.spec_dec_mode != SpeculativeDecodingMode.DFLASH:
-        return None
 
+def _get_generation_kv_capacity_headroom(spec_config, *, is_draft: bool) -> int:
+    """Maximum capacity lead over history used by generation KV allocation."""
     from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens
 
-    return get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step
+    if spec_config is None:
+        return BASE_GENERATION_TOKEN_COUNT
+    reserve = _get_kv_reserve_draft_tokens(spec_config, is_draft=is_draft)
+    dynamic_reserve = reserve - spec_config.max_total_draft_tokens
+    return get_num_extra_kv_tokens(spec_config) + spec_config.tokens_per_gen_step + dynamic_reserve
 
 
 def _get_single_swa_pool_slot_bytes(
@@ -1029,15 +1045,12 @@ def __init__(
         if not self._supports_reuse_match_backoff:
             self.reuse_match_backoff = 0
         # Mirror V1's KV reserve sizing (see V1 __init__ for rationale).
-        self._kv_reserve_draft_tokens = self.max_total_draft_tokens
-        if (
-            self.is_draft
-            and spec_config is not None
-            and getattr(spec_config, "use_dynamic_tree", False)
-            and getattr(spec_config, "dynamic_tree_max_topK", 0) > 0
-        ):
-            draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len
-            self._kv_reserve_draft_tokens = max(self.max_total_draft_tokens, draft_loop_tokens)
+        self._kv_reserve_draft_tokens = _get_kv_reserve_draft_tokens(
+            spec_config, is_draft=self.is_draft
+        )
+        self._generation_kv_capacity_headroom = _get_generation_kv_capacity_headroom(
+            spec_config, is_draft=self.is_draft
+        )
 
         self.event_buffer_max_size = kv_cache_config.event_buffer_max_size
         self.enable_stats = enable_stats
@@ -1417,11 +1430,8 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]:
             self.max_seq_len = int(max_num_tokens)
 
         # Pad max_blocks_per_seq to next multiple of 4 (copy_block_offsets kernel).
-        # Account for max single-sequence capacity = seq_len + extra KV tokens +
-        # _kv_reserve_draft_tokens (see __init__) + 1 base decode token.
-        max_seq_capacity = (
-            self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1
-        )
+        # Include the same maximum generation lead used by allocation and sizing.
+        max_seq_capacity = self.max_seq_len + self._generation_kv_capacity_headroom
         self.max_blocks_per_seq = (
             max_seq_capacity + self._ledger_tokens_per_block - 1
         ) // self._ledger_tokens_per_block
@@ -1721,7 +1731,12 @@ def _get_max_tokens_from_quota_impl(self, quota: int) -> float:
             generation_swa_size_per_token,
             generation_swa_size_per_request,
         ) = _estimate_swa_cache_size(
-            layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False
+            layer_sizes,
+            attention_windows,
+            self.tokens_per_block,
+            context=False,
+            scratch=False,
+            generation_capacity_headroom=self._generation_kv_capacity_headroom,
         )
         size_per_batch = self.max_batch_size * generation_swa_size_per_request
         if quota < size_per_batch:
@@ -1767,7 +1782,12 @@ def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int:
             generation_swa_size_per_token,
             generation_swa_size_per_request,
         ) = _estimate_swa_cache_size(
-            layer_sizes, attention_windows, self.tokens_per_block, context=False, scratch=False
+            layer_sizes,
+            attention_windows,
+            self.tokens_per_block,
+            context=False,
+            scratch=False,
+            generation_capacity_headroom=self._generation_kv_capacity_headroom,
         )
         context_tokens = min(max_tokens, self.max_num_tokens)
         generation_tokens = max_tokens - context_tokens
@@ -2546,7 +2566,7 @@ def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int:
 
         Grows *current_capacity* by 1 + draft tokens.
         """
-        return current_capacity + 1 + self._generation_draft_slots(req)
+        return current_capacity + BASE_GENERATION_TOKEN_COUNT + self._generation_draft_slots(req)
 
     def _generation_draft_slots(self, req: LlmRequest) -> int:
         """Physical draft width reserved for one iteration. Dynamic-tree draft pools
@@ -4165,9 +4185,9 @@ def get_cache_size_per_token(
         full_attn_size_per_token = _estimate_full_attn_size_per_token(
             layer_sizes, attention_windows
         )
-        dflash_headroom = _get_dflash_generation_kv_capacity_headroom(spec_config)
-        is_dflash_draft = is_draft and dflash_headroom is not None
-        generation_capacity_headroom = dflash_headroom if is_dflash_draft else None
+        generation_capacity_headroom = _get_generation_kv_capacity_headroom(
+            spec_config, is_draft=is_draft
+        )
         swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size(
             layer_sizes,
             attention_windows,
@@ -4176,15 +4196,13 @@ def get_cache_size_per_token(
             scratch=False,
             generation_capacity_headroom=generation_capacity_headroom,
         )
-        context_swa_size_per_token = 0
-        if is_dflash_draft:
-            context_swa_size_per_token, _ = _estimate_swa_cache_size(
-                layer_sizes,
-                attention_windows,
-                tokens_per_block,
-                context=True,
-                scratch=False,
-            )
+        context_swa_size_per_token, _ = _estimate_swa_cache_size(
+            layer_sizes,
+            attention_windows,
+            tokens_per_block,
+            context=True,
+            scratch=bool(kwargs.get("enable_swa_scratch_reuse", False)),
+        )
         fixed_cost = (
             swa_size_per_request * max_batch_size + context_swa_size_per_token * max_num_tokens
         )
@@ -4192,7 +4210,11 @@ def get_cache_size_per_token(
         bytes_per_slot = _get_single_swa_pool_slot_bytes(
             layer_sizes, attention_windows, tokens_per_block
         )
-        if is_dflash_draft and bytes_per_slot is not None:
+        if is_draft and fixed_cost > 0 and bytes_per_slot is not None:
+            # The affine intercept is the configured quota needed to preserve
+            # the fixed usable capacity at V2's resume watermark. Keep the
+            # allocator's page rounding inside the manager estimator rather
+            # than extending the generic CacheCost model with pool geometry.
             required_slots = math.ceil(fixed_cost / bytes_per_slot)
             resume_util = float(np.float32(kv_cache_config.max_util_for_resume))
             fixed_cost = math.ceil(required_slots / resume_util) * bytes_per_slot
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py
index e2575f414ea2..3632b6e89854 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py
@@ -14,9 +14,11 @@
 # limitations under the License.
 """Tests for KV cache budget splitting between target and draft managers."""
 
+import math
 from types import SimpleNamespace
 from unittest.mock import Mock
 
+import numpy as np
 import pytest
 
 from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator
@@ -47,8 +49,8 @@ def _make_creator(
     """Minimal KvCacheCreator for budget-split helpers.
 
     ``*_intercept`` model the affine fixed cost (e.g. mamba SSM state) that a
-    manager pays per batch regardless of token count. The draft cost is derived
-    as ``total - target`` for both slope and intercept.
+    manager pays per batch regardless of token count. The draft mock receives
+    the component-wise ``total - target`` values directly.
     """
     c = object.__new__(KvCacheCreator)
 
@@ -78,6 +80,12 @@ def _make_creator(
         )
     )
     c._should_create_separate_draft_kv_cache = Mock(return_value=True)
+    c._get_draft_cache_cost = Mock(
+        return_value=CacheCost(
+            slope=total_kv_per_token - target_kv_per_token,
+            intercept=total_kv_intercept - target_kv_intercept,
+        )
+    )
 
     return c
 
@@ -136,7 +144,14 @@ def get_cache_size_per_token(model_config, *args, **kwargs):
         creator._mapping = Mock(enable_attention_dp=False, tp_size=1)
         creator._mapping.pp_layers.return_value = [0]
         creator._mapping.is_last_pp_rank.return_value = True
-        creator._speculative_config = SimpleNamespace(spec_dec_mode=mode)
+        creator._speculative_config = SimpleNamespace(
+            spec_dec_mode=mode,
+            max_draft_len=1,
+            max_total_draft_tokens=0,
+            tokens_per_gen_step=1,
+            use_dynamic_tree=False,
+            _use_shared_kv_cache=False,
+        )
         creator._model_engine = SimpleNamespace(
             model=SimpleNamespace(model_config=target_model_config)
         )
@@ -153,9 +168,13 @@ def get_cache_size_per_token(model_config, *args, **kwargs):
         )
 
         # The draft layer stores 64 bytes/token in a fixed 512-token window.
+        # Generation retains one additional 64-token boundary block and the
+        # 128-token context budget consumes two more blocks.
         # Leaking the target's 16K window would instead count it as 64 bytes/token.
         cost = creator._get_kv_size_per_token()
-        assert cost == CacheCost(slope=10, intercept=512 * 64)
+        usable_slots = 11
+        configured_slots = math.ceil(usable_slots / float(np.float32(0.95)))
+        assert cost == CacheCost(slope=10, intercept=configured_slots * 64 * 64)
         assert len(draft_kv_configs) == 1
         draft_kv_config = draft_kv_configs[0]
         assert draft_kv_config.max_attention_window == [512]
@@ -269,9 +288,7 @@ def test_returns_none_when_draft_kv_zero(self):
 
     def test_fixed_only_draft_uses_manager_estimated_quota(self):
         total_gpu = 10 * GB
-        slot_bytes = 327_680
-        configured_slots = 2_561
-        configured_bytes = configured_slots * slot_bytes
+        configured_bytes = 2_561 * 327_680
         c = _make_creator(
             max_gpu_total_bytes=total_gpu,
             total_kv_per_token=80,
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
index 6bad069afaa4..3d7379006138 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
@@ -15,6 +15,7 @@
 from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
+import numpy as np
 import pytest
 import torch
 
@@ -649,12 +650,21 @@ def get_num_attention_layers(self):
     )
 
     # Per layer: K+V * kv_heads * head_dim * bf16 bytes = 2 * 2 * 8 * 2.
-    expected = CacheCost(slope=64, intercept=3 * 2 * 2048 * 64)
+    expected = CacheCost(slope=64, intercept=3 * 2 * (2048 + 64) * 64)
     assert no_scratch_size_per_token == expected
     assert scratch_size_per_token == expected
 
 
-def test_v2_dflash_draft_cost_covers_context_and_generation_slots():
+@pytest.mark.parametrize(
+    ("spec_dec_mode", "tokens_per_gen_step", "headroom", "generation_blocks"),
+    [
+        pytest.param(SpeculativeDecodingMode.DFLASH, 5, 8, 17, id="dflash"),
+        pytest.param(SpeculativeDecodingMode.PARD, 8, 11, 18, id="pard"),
+    ],
+)
+def test_v2_external_draft_cost_covers_context_and_generation_slots(
+    spec_dec_mode, tokens_per_gen_step, headroom, generation_blocks
+):
     class FakeDraftModelConfig:
         quant_config = None
         pretrained_config = SimpleNamespace(
@@ -667,9 +677,12 @@ def get_num_attention_layers(self):
             return 1
 
     spec_config = SimpleNamespace(
-        spec_dec_mode=SpeculativeDecodingMode.DFLASH,
+        spec_dec_mode=spec_dec_mode,
         max_draft_len=4,
-        tokens_per_gen_step=5,
+        max_total_draft_tokens=4,
+        tokens_per_gen_step=tokens_per_gen_step,
+        use_dynamic_tree=False,
+        _use_shared_kv_cache=False,
     )
     mapping = Mock(enable_attention_dp=False, tp_size=1)
     mapping.pp_layers.return_value = [0]
@@ -685,7 +698,7 @@ def get_num_attention_layers(self):
             max_seq_len=4096,
             max_batch_size=max_batch_size,
             max_num_tokens=max_num_tokens,
-            kv_cache_config=KvCacheConfig(max_attention_window=[512]),
+            kv_cache_config=KvCacheConfig(max_attention_window=[504]),
             spec_config=spec_config,
             is_draft=True,
         )
@@ -693,16 +706,14 @@ def get_num_attention_layers(self):
 
     # One layer stores K+V * 2 KV heads * 8 head dim * BF16 = 64 B/token.
     slot_bytes = tokens_per_block * 64
-    # Runtime generation capacity can lead history by max_draft_len - 1
-    # (get_num_extra_kv_tokens) plus tokens_per_gen_step: 3 + 5 = 8.
-    generation_blocks_per_request = math.ceil((512 + 8 - 1) / tokens_per_block) + 1
-    assert generation_blocks_per_request == 18
+    # Runtime generation capacity leads history by max_draft_len - 1 plus
+    # tokens_per_gen_step. For PARD K=4 that is 3 + 2K = 11 tokens.
+    generation_blocks_per_request = math.ceil((504 + headroom - 2) / tokens_per_block) + 1
+    assert generation_blocks_per_request == generation_blocks
     context_slots = max_num_tokens // tokens_per_block
     expected_usable_slots = max_batch_size * generation_blocks_per_request + context_slots
-    assert expected_usable_slots == 2_432
-    # float32(0.95) requires 2561 configured slots for 2432 slots to remain
-    # resumable. The estimator owns this manager-specific quota normalization.
-    expected_configured_slots = 2_561
+    resume_util = float(np.float32(0.95))
+    expected_configured_slots = math.ceil(expected_usable_slots / resume_util)
     assert cost == CacheCost(slope=0, intercept=expected_configured_slots * slot_bytes)
 
 
@@ -759,7 +770,10 @@ class UnsupportedKVCacheManagerV2(KVCacheManagerV2):
         tokens_per_block=64,
         max_seq_len=4096,
         max_batch_size=3,
-        kv_cache_config=KvCacheConfig(max_attention_window=[64]),
+        kv_cache_config=KvCacheConfig(
+            max_attention_window=[64],
+            max_util_for_resume=1.0,
+        ),
         spec_config=spec_config,
     )
 
@@ -797,18 +811,20 @@ class UnsupportedKVCacheManagerV2(KVCacheManagerV2):
                     "kv_cache_config": KvCacheConfig(
                         enable_block_reuse=False,
                         max_attention_window=[64],
+                        max_util_for_resume=1.0,
                     )
                 }
             ),
         )
     )
 
-    # W=64 occupies one page; one-model draft reuse retains W+D=65 and
-    # therefore charges two pages in single and separate KVCM layouts.
-    assert no_draft == CacheCost(slope=0, intercept=3 * 64 * 64)
+    # W=64 plus the base generation token can retain two boundary pages.
+    # One-model draft reuse extends retention to W+D=65; its two-token
+    # generation step can therefore cross into a third page.
+    assert no_draft == CacheCost(slope=0, intercept=3 * 128 * 64)
     assert unsupported == no_draft
     assert block_reuse_disabled == no_draft
-    assert single_kvcm == target == draft == CacheCost(slope=0, intercept=3 * 128 * 64)
+    assert single_kvcm == target == draft == CacheCost(slope=0, intercept=3 * 192 * 64)
 
 
 def test_creator_uses_v2_affine_cache_cost():
@@ -839,6 +855,7 @@ def test_v2_quota_from_max_tokens_models_context_swa_scratch():
     manager.tokens_per_block = 64
     manager.max_batch_size = 4
     manager.max_num_tokens = 1000
+    manager._generation_kv_capacity_headroom = 1
     manager.get_layer_bytes_per_token = lambda local_layer_idx, data_role: [10, 10, 20][
         local_layer_idx
     ]
@@ -847,12 +864,12 @@ def test_v2_quota_from_max_tokens_models_context_swa_scratch():
 
     manager.enable_swa_scratch_reuse = False
     no_scratch_quota = manager._get_quota_from_max_tokens(max_tokens)
-    assert no_scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 20 + 4 * 2 * 128 * 10)
+    assert no_scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 20 + 4 * 2 * 192 * 10)
     assert manager._get_max_tokens_from_quota(no_scratch_quota) == max_tokens
 
     manager.enable_swa_scratch_reuse = True
     scratch_quota = manager._get_quota_from_max_tokens(max_tokens)
-    assert scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 10 + 4 * 2 * 128 * 10)
+    assert scratch_quota == (max_tokens * 20 + manager.max_num_tokens * 10 + 4 * 2 * 192 * 10)
     assert manager._get_max_tokens_from_quota(scratch_quota) == max_tokens
 
 
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 1212a39e8dc3..baa79ad394f7 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
@@ -42,6 +42,7 @@
     Eagle3DecodingConfig,
     KvCacheConfig,
     MTPDecodingConfig,
+    PARDDecodingConfig,
 )
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
@@ -154,6 +155,7 @@ def _make_manager_for_cache_tier_test(
     is_disagg: bool = False,
     joint_reuse: bool = False,
     mapping: Mapping | None = None,
+    max_seq_len: int = MAX_SEQ_LEN,
 ) -> tuple[KVCacheManagerV2, Mock]:
     impl_constructor = Mock(side_effect=impl_side_effect)
     if mapping is None:
@@ -200,7 +202,7 @@ def build_cache_config(
         patch(f"{module}.KVCacheManagerPy", impl_constructor),
         patch.object(KVCacheManagerV2, "_build_base_config", build_base_config),
         patch.object(KVCacheManagerV2, "_build_cache_config", build_cache_config),
-        patch.object(KVCacheManagerV2, "get_num_available_tokens", return_value=MAX_SEQ_LEN),
+        patch.object(KVCacheManagerV2, "get_num_available_tokens", return_value=max_seq_len),
         patch.object(KVCacheManagerV2, "_prepare_page_table_tensor"),
         patch.object(KVCacheManagerV2, "_log_kv_cache_pool_lifecycle_mapping"),
         patch(f"{module}.get_pp_layers", return_value=([0], 1)),
@@ -212,7 +214,7 @@ def build_cache_config(
             num_kv_heads=1,
             head_dim=1,
             tokens_per_block=TOKENS_PER_BLOCK,
-            max_seq_len=MAX_SEQ_LEN,
+            max_seq_len=max_seq_len,
             max_batch_size=1,
             mapping=mapping,
             dtype=DataType.HALF,
@@ -917,6 +919,32 @@ def resize(capacity, history_length=None):
     assert request.py_request_id not in manager._allocated_draft_lens
 
 
+def test_pard_max_blocks_cover_generation_capacity_headroom() -> None:
+    manager, _ = _make_manager_for_cache_tier_test(
+        KvCacheConfig(
+            enable_block_reuse=False,
+            max_gpu_total_bytes=16 << 20,
+        ),
+        [Mock()],
+        spec_config=PARDDecodingConfig(
+            max_draft_len=4,
+            speculative_model="draft-model",
+        ),
+        is_draft=True,
+        max_seq_len=24,
+    )
+
+    assert manager._generation_kv_capacity_headroom == 11
+    required_capacity = manager.max_seq_len + manager._generation_kv_capacity_headroom
+    required_blocks = (
+        required_capacity + manager._ledger_tokens_per_block - 1
+    ) // manager._ledger_tokens_per_block
+    expected_blocks = ((required_blocks + 3) // 4) * 4
+
+    assert expected_blocks == 12
+    assert manager.max_blocks_per_seq == expected_blocks
+
+
 def _revert_context_request(request_id: int) -> SimpleNamespace:
     return SimpleNamespace(
         py_request_id=request_id,

From 1de090cb1e26943920483548175d829f4c5a4c12 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 9 Sep 2026 06:53:04 +0000
Subject: [PATCH 2/3] [None][refactor] Share V2 KV cache capacity accounting

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 .../kv_cache/kv_cache_manager_v2.py           | 107 +++++++++---------
 .../kv_cache/test_kv_cache_estimation.py      |  78 +++++++++++++
 2 files changed, 132 insertions(+), 53 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 a7212d64140a..e2aa74131c5f 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
@@ -380,6 +380,39 @@ def _get_kv_reserve_draft_tokens(spec_config, *, is_draft: bool) -> int:
     return reserve
 
 
+def _estimate_cache_size_components(
+    layer_sizes: Sequence[int],
+    attention_windows: Sequence[int | None],
+    tokens_per_block: int,
+    *,
+    scratch: bool,
+    generation_capacity_headroom: int,
+) -> tuple[int, int, int]:
+    """Return context/generation bytes per token and generation bytes per request.
+
+    Static profiling and runtime quota conversion must charge the same SWA
+    retention pages and context scratch space. Resume-watermark normalization
+    is separate from these usable-capacity costs.
+    """
+    full_attn_size = _estimate_full_attn_size_per_token(layer_sizes, attention_windows)
+    context_swa_size, _ = _estimate_swa_cache_size(
+        layer_sizes, attention_windows, tokens_per_block, context=True, scratch=scratch
+    )
+    generation_swa_size, generation_swa_per_request = _estimate_swa_cache_size(
+        layer_sizes,
+        attention_windows,
+        tokens_per_block,
+        context=False,
+        scratch=False,
+        generation_capacity_headroom=generation_capacity_headroom,
+    )
+    return (
+        full_attn_size + context_swa_size,
+        full_attn_size + generation_swa_size,
+        generation_swa_per_request,
+    )
+
+
 def _get_generation_kv_capacity_headroom(spec_config, *, is_draft: bool) -> int:
     """Maximum capacity lead over history used by generation KV allocation."""
     from tensorrt_llm._torch.speculative import get_num_extra_kv_tokens
@@ -1717,38 +1750,26 @@ def _get_max_tokens_from_quota(self, quota: int) -> float:
 
     def _get_max_tokens_from_quota_impl(self, quota: int) -> float:
         layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components()
-        full_attn_size_per_token = _estimate_full_attn_size_per_token(
-            layer_sizes, attention_windows
-        )
-        context_swa_size_per_token, _ = _estimate_swa_cache_size(
-            layer_sizes,
-            attention_windows,
-            self.tokens_per_block,
-            context=True,
-            scratch=self.enable_swa_scratch_reuse,
-        )
         (
-            generation_swa_size_per_token,
+            context_size_per_token,
+            generation_size_per_token,
             generation_swa_size_per_request,
-        ) = _estimate_swa_cache_size(
+        ) = _estimate_cache_size_components(
             layer_sizes,
             attention_windows,
             self.tokens_per_block,
-            context=False,
-            scratch=False,
+            scratch=self.enable_swa_scratch_reuse,
             generation_capacity_headroom=self._generation_kv_capacity_headroom,
         )
         size_per_batch = self.max_batch_size * generation_swa_size_per_request
         if quota < size_per_batch:
             return 0
-        context_size_per_token = full_attn_size_per_token + context_swa_size_per_token
         context_limit_quota = self.max_num_tokens * context_size_per_token + size_per_batch
         if quota <= context_limit_quota:
             if context_size_per_token <= 0:
                 return float("inf")
             return (quota - size_per_batch) / context_size_per_token
 
-        generation_size_per_token = full_attn_size_per_token + generation_swa_size_per_token
         if generation_size_per_token <= 0:
             return float("inf")
         return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token
@@ -1765,39 +1786,24 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int:
 
     def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int:
         layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components()
-        full_attn_size_per_token = _estimate_full_attn_size_per_token(
-            layer_sizes, attention_windows
-        )
         (
-            context_swa_size_per_token,
-            _,
-        ) = _estimate_swa_cache_size(
-            layer_sizes,
-            attention_windows,
-            self.tokens_per_block,
-            context=True,
-            scratch=self.enable_swa_scratch_reuse,
-        )
-        (
-            generation_swa_size_per_token,
+            context_size_per_token,
+            generation_size_per_token,
             generation_swa_size_per_request,
-        ) = _estimate_swa_cache_size(
+        ) = _estimate_cache_size_components(
             layer_sizes,
             attention_windows,
             self.tokens_per_block,
-            context=False,
-            scratch=False,
+            scratch=self.enable_swa_scratch_reuse,
             generation_capacity_headroom=self._generation_kv_capacity_headroom,
         )
         context_tokens = min(max_tokens, self.max_num_tokens)
         generation_tokens = max_tokens - context_tokens
-        generation_quota = (
-            max_tokens * full_attn_size_per_token
-            + generation_tokens * generation_swa_size_per_token
+        return int(
+            context_tokens * context_size_per_token
+            + generation_tokens * generation_size_per_token
             + self.max_batch_size * generation_swa_size_per_request
         )
-        context_extra_quota = context_tokens * context_swa_size_per_token
-        return int(generation_quota + context_extra_quota)
 
     def _get_event_num_blocks_per_cache_level(
         self,
@@ -4182,31 +4188,26 @@ def get_cache_size_per_token(
                 backoff,
                 max_seq_len,
             )
-        full_attn_size_per_token = _estimate_full_attn_size_per_token(
-            layer_sizes, attention_windows
-        )
         generation_capacity_headroom = _get_generation_kv_capacity_headroom(
             spec_config, is_draft=is_draft
         )
-        swa_size_per_token, swa_size_per_request = _estimate_swa_cache_size(
-            layer_sizes,
-            attention_windows,
-            tokens_per_block,
-            context=False,
-            scratch=False,
-            generation_capacity_headroom=generation_capacity_headroom,
-        )
-        context_swa_size_per_token, _ = _estimate_swa_cache_size(
+        (
+            context_size_per_token,
+            cache_size_per_token,
+            swa_size_per_request,
+        ) = _estimate_cache_size_components(
             layer_sizes,
             attention_windows,
             tokens_per_block,
-            context=True,
             scratch=bool(kwargs.get("enable_swa_scratch_reuse", False)),
+            generation_capacity_headroom=generation_capacity_headroom,
         )
+        # The affine slope covers all tokens; context additionally retains SWA
+        # pages for the current token batch beyond the generation windows.
         fixed_cost = (
-            swa_size_per_request * max_batch_size + context_swa_size_per_token * max_num_tokens
+            swa_size_per_request * max_batch_size
+            + (context_size_per_token - cache_size_per_token) * max_num_tokens
         )
-        cache_size_per_token = full_attn_size_per_token + swa_size_per_token
         bytes_per_slot = _get_single_swa_pool_slot_bytes(
             layer_sizes, attention_windows, tokens_per_block
         )
diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
index 3d7379006138..62e69d4c02e0 100644
--- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
+++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py
@@ -27,9 +27,11 @@
 from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode
 from tensorrt_llm.inputs.multimodal import MultimodalParams
 from tensorrt_llm.llmapi.llm_args import (
+    DFlashDecodingConfig,
     KvCacheConfig,
     MTPDecodingConfig,
     MultimodalConfig,
+    PARDDecodingConfig,
     TorchLlmArgs,
 )
 from tensorrt_llm.mapping import Mapping
@@ -873,6 +875,82 @@ def test_v2_quota_from_max_tokens_models_context_swa_scratch():
     assert manager._get_max_tokens_from_quota(scratch_quota) == max_tokens
 
 
+@pytest.mark.parametrize(
+    ("spec_config", "headroom", "generation_blocks"),
+    [
+        pytest.param(None, 1, 17, id="normal"),
+        pytest.param(
+            DFlashDecodingConfig(max_draft_len=4, speculative_model="draft"),
+            8,
+            17,
+            id="dflash",
+        ),
+        pytest.param(
+            PARDDecodingConfig(max_draft_len=4, speculative_model="draft"),
+            11,
+            18,
+            id="pard",
+        ),
+    ],
+)
+@pytest.mark.parametrize("scratch", [False, True])
+@pytest.mark.parametrize("is_draft", [False, True])
+@pytest.mark.parametrize("max_tokens", [64, 4096, 8192])
+def test_v2_static_and_runtime_cache_costs_agree(
+    spec_config: DFlashDecodingConfig | PARDDecodingConfig | None,
+    headroom: int,
+    generation_blocks: int,
+    scratch: bool,
+    is_draft: bool,
+    max_tokens: int,
+) -> None:
+    class FakeModelConfig:
+        quant_config = None
+        pretrained_config = SimpleNamespace(
+            hidden_size=32, num_attention_heads=4, num_key_value_heads=2
+        )
+
+        def get_num_attention_layers(self) -> int:
+            return 3
+
+    manager = object.__new__(KVCacheManagerV2)
+    manager._has_cp_helix = False
+    manager.num_local_layers = 3
+    manager.max_attention_window_vec = [504, 504, None]
+    manager.tokens_per_block = 32
+    manager.max_batch_size = 3
+    manager.max_num_tokens = 4096
+    manager._generation_kv_capacity_headroom = headroom
+    manager.enable_swa_scratch_reuse = scratch
+    manager.get_layer_bytes_per_token = lambda local_layer_idx, data_role: 64
+
+    cost = CacheCost.from_raw(
+        KVCacheManagerV2.get_cache_size_per_token(
+            FakeModelConfig(),
+            Mapping(),
+            tokens_per_block=32,
+            max_seq_len=16384,
+            max_batch_size=3,
+            max_num_tokens=4096,
+            kv_cache_config=KvCacheConfig(
+                enable_block_reuse=False, max_attention_window=[504, 504, 16384]
+            ),
+            spec_config=spec_config,
+            is_draft=is_draft,
+            enable_swa_scratch_reuse=scratch,
+        )
+    )
+    # Each layer stores 64 B/token. The two SWA layers can share context
+    # scratch, but both retain their own generation pages. W=504 puts PARD's
+    # 11-token headroom across a page boundary that DFlash's 8 tokens fit under.
+    context_swa_bytes = 64 if scratch else 128
+    generation_bytes = 3 * 2 * generation_blocks * 32 * 64
+    expected_quota = max_tokens * 64 + min(max_tokens, 4096) * context_swa_bytes + generation_bytes
+    assert cost == CacheCost(slope=64, intercept=4096 * context_swa_bytes + generation_bytes)
+    assert manager._get_quota_from_max_tokens(max_tokens) == expected_quota
+    assert manager._get_max_tokens_from_quota(expected_quota) == max_tokens
+
+
 # ---------------------------------------------------------------------------
 # KVCacheManagerV2 clamp_max_seq_len_for_mem float-to-int cast regression
 # ---------------------------------------------------------------------------

From b8169da570e3aab1b8267b47b52aae9b06c00233 Mon Sep 17 00:00:00 2001
From: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
Date: Wed, 9 Sep 2026 07:40:57 +0000
Subject: [PATCH 3/3] [None][test] Run PARD accuracy cases with KV cache
 manager V2

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
---
 tests/integration/defs/accuracy/test_llm_api_pytorch.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 99689df243fe..ece9bd911a50 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -616,8 +616,10 @@ def test_pard(self, overlap_scheduler):
                                               enable_padding=True),
         )
         kv_cache_config = KvCacheConfig(
-            enable_block_reuse=True, free_gpu_memory_fraction=0.8
-        )  # both one-model and two-model supports this feature
+            enable_block_reuse=True,
+            free_gpu_memory_fraction=0.8,
+            use_kv_cache_manager_v2=True,
+        )
 
         pard_model_dir = f"{llm_models_root()}/PARD-Llama-3.2-1B"
         target_model_dir = f"{llm_models_root()}/llama-3.1-model/Llama-3.1-8B-Instruct"