[nvbugs/6035425][fix] Fix host side memory doubling when V2 auto provisioning is used

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -46,7 +46,8 @@
 from .py_executor import PyExecutor
 from .resource_manager import (KVCacheManager, KVCacheManagerV2,
                                PeftCacheManager, ResourceManager,
-                               ResourceManagerType)
+                               ResourceManagerType,
+                               compute_default_v2_host_quota)
 from .sampler import (EarlyStopSampler, EarlyStopWithMMResult, TorchSampler,
                       TRTLLMSampler)
 from .scheduler import (BindCapacityScheduler, BindMicroBatchScheduler,
@@ -899,6 +900,15 @@ def _split_kv_cache_budget_for_draft(self) -> Optional[KvCacheConfig]:
         budget for both target and draft combined.  This method splits both
         budgets proportionally based on their per-token KV cache sizes.
 
+        For V2, when host_cache_size is not configured, KVCacheManagerV2
+        auto-provisions a host tier in its constructor based on the GPU
+        quota.  If we left host_cache_size unset here, both the target and
+        the draft manager would each auto-provision a full host tier --
+        roughly doubling host memory usage.  To prevent that, we
+        pre-compute what the auto-provision would yield from the combined
+        GPU budget and split it explicitly so each manager sees its own
+        share.
+
         Returns a cloned KvCacheConfig for the draft, or None if no split is
         needed.  Also modifies self._kv_cache_config in-place for the target.
         """
@@ -943,6 +953,17 @@ def _split_kv_cache_budget_for_draft(self) -> Optional[KvCacheConfig]:
         draft_kv_cache_config.max_gpu_total_bytes = draft_budget
 
         host_budget = self._kv_cache_config.host_cache_size
+        # When V2 would otherwise auto-provision a host tier per manager,
+        # materialise the would-be combined host budget up front so it gets
+        # split here instead of duplicated across target and draft.
+        if ((host_budget is None or host_budget <= 0)
+                and issubclass(self._kv_cache_manager_cls, KVCacheManagerV2)):
+            host_budget = compute_default_v2_host_quota(total_budget)
+            logger.info(
+                f"V2 host cache size not set; pre-splitting auto-provisioned "
+                f"host budget {host_budget / GB:.2f} GiB between target and "
+                f"draft to avoid duplicate allocation.")
+
         if host_budget is not None and host_budget > 0:
             draft_ratio = draft_budget / total_budget
             draft_host_budget = int(host_budget * draft_ratio)

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -2018,6 +2018,30 @@ def reset_reuse_state(self):
         self.impl.reset_reuse_state()
 
 
+def compute_default_v2_host_quota(gpu_quota: int) -> int:
+    """Default host cache quota when the user did not set ``host_cache_size``.
+
+    KVCacheManagerV2's MAX_UTILIZATION scheduler relies on suspend/resume which
+    requires a host tier to be present.  When the user has not configured one
+    we provision a tier matching the GPU quota, capped at half of the
+    currently-available host memory to avoid allocation failures.
+
+    Centralised here so callers that need to pre-split the host budget (e.g.
+    one-model spec dec creating both a target and a draft manager) can compute
+    the same quota that ``KVCacheManagerV2.__init__`` would have used, then
+    divide it explicitly between the two managers.
+    """
+    try:
+        mem_available = os.sysconf('SC_PAGE_SIZE') * os.sysconf(
+            'SC_AVPHYS_PAGES')
+    except (ValueError, OSError):
+        mem_available = float('inf')
+    host_quota = min(gpu_quota, int(mem_available * 0.5))
+    if host_quota <= 0:
+        host_quota = gpu_quota
+    return host_quota
+
+
 class KVCacheManagerV2(BaseResourceManager):
 
     def __init__(
@@ -2201,23 +2225,7 @@ def append_to_kv_heads_per_layer(num_kv_heads_per_layer: List[int],
         if kv_cache_config.host_cache_size is not None and kv_cache_config.host_cache_size > 0:
             host_quota = kv_cache_config.host_cache_size
         else:
-            # The V2 MAX_UTILIZATION scheduler relies on suspend/resume to
-            # evict and later restore KV cache pages.  Without a host tier,
-            # suspended pages have nowhere to be offloaded and resume()
-            # always fails, causing a scheduling deadlock where no
-            # generation request can ever make progress.
-            #
-            # Automatically provision a host tier matching the GPU quota so
-            # suspend/resume works out of the box.  Cap at available host
-            # memory to avoid allocation failures.
-            try:
-                mem_available = os.sysconf('SC_PAGE_SIZE') * os.sysconf(
-                    'SC_AVPHYS_PAGES')
-            except (ValueError, OSError):
-                mem_available = float('inf')
-            host_quota = min(quota, int(mem_available * 0.5))
-            if host_quota <= 0:
-                host_quota = quota
+            host_quota = compute_default_v2_host_quota(quota)
         if host_quota > 0:
             cache_tiers.append(HostCacheTierConfig(quota=host_quota))
             logger.info(

tests/unittest/_torch/executor/test_kv_cache_budget_split.py

@@ -23,6 +23,8 @@
 import pytest
 
 from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator
+from tensorrt_llm._torch.pyexecutor.resource_manager import (KVCacheManager,
+                                                             KVCacheManagerV2)
 from tensorrt_llm.llmapi.llm_args import KvCacheConfig
 
 GB = 1 << 30
@@ -33,12 +35,19 @@ def _make_creator(
     host_cache_size=None,
     total_kv_per_token: int = 100,
     target_kv_per_token: int = 80,
+    kv_cache_manager_cls=KVCacheManager,
 ) -> KvCacheCreator:
     """Build a minimal KvCacheCreator wired for _split_kv_cache_budget_for_draft.
 
     Per-token costs are exposed via the new ``CacheCost`` shape; the manager
     mock returns a raw int so we also exercise ``_per_manager_cache_cost``'s
     ``CacheCost.from_raw`` wrapping.
+
+    ``kv_cache_manager_cls`` must be a real class so the ``issubclass`` check
+    inside ``_split_kv_cache_budget_for_draft`` (used to detect V2 and
+    pre-split the auto-provisioned host quota) works.  Tests default to the
+    V1 manager so they don't trip the V2 host auto-provision branch unless
+    they explicitly opt in.
     """
     c = object.__new__(KvCacheCreator)
 
@@ -51,9 +60,10 @@ def _make_creator(
     c._mapping = Mock()
     c._model_engine = Mock()
 
-    c._kv_cache_manager_cls = Mock()
-    c._kv_cache_manager_cls.get_cache_size_per_token = Mock(return_value=target_kv_per_token)
+    c._kv_cache_manager_cls = kv_cache_manager_cls
 
+    c._per_manager_cache_cost = Mock(
+        return_value=CacheCost(slope=target_kv_per_token))
     c._get_kv_size_per_token = Mock(return_value=CacheCost(slope=total_kv_per_token))
 
     return c
@@ -126,26 +136,30 @@ def test_budgets_sum_to_original(self):
         ) == total_gpu
         assert (c._kv_cache_config.host_cache_size + draft_config.host_cache_size) == total_host
 
-    def test_no_host_cache_leaves_none(self):
+    def test_v1_no_host_cache_leaves_none(self):
         c = _make_creator(
             max_gpu_total_bytes=10 * GB,
             host_cache_size=None,
             total_kv_per_token=100,
             target_kv_per_token=80,
+            kv_cache_manager_cls=KVCacheManager,
         )
 
         draft_config = c._split_kv_cache_budget_for_draft()
 
         assert draft_config is not None
+        # V1 does not auto-provision a host tier, so an unset
+        # host_cache_size must remain unset for both target and draft.
         assert c._kv_cache_config.host_cache_size is None
         assert draft_config.host_cache_size is None
 
-    def test_zero_host_cache_unchanged(self):
+    def test_v1_zero_host_cache_unchanged(self):
         c = _make_creator(
             max_gpu_total_bytes=10 * GB,
             host_cache_size=0,
             total_kv_per_token=100,
             target_kv_per_token=80,
+            kv_cache_manager_cls=KVCacheManager,
         )
 
         draft_config = c._split_kv_cache_budget_for_draft()
@@ -186,3 +200,125 @@ def test_various_ratios(self, target_frac):
             c._kv_cache_config.max_gpu_total_bytes + draft_config.max_gpu_total_bytes
         ) == total_gpu
         assert (c._kv_cache_config.host_cache_size + draft_config.host_cache_size) == total_host
+
+
+class TestV2AutoProvisionHostSplit:
+    """Tests for the V2-specific behaviour where the host tier is auto-
+    provisioned by ``KVCacheManagerV2.__init__`` when the user did not set
+    ``host_cache_size``.
+
+    Before the fix, both target and draft managers each independently
+    auto-provisioned a full host tier, doubling host memory.  The split
+    method now materialises the would-be combined host quota up front
+    (via ``compute_default_v2_host_quota``) and divides it explicitly.
+    """
+
+    def test_v2_unset_host_is_pre_split(self, monkeypatch):
+        """V2 + host_cache_size unset must end up with both target and
+        draft host budgets set (so the V2 constructor will not
+        auto-provision again)."""
+        from tensorrt_llm._torch.pyexecutor import _util
+
+        # Pin the would-be auto-provision result so the test is deterministic
+        # regardless of the runner's available memory.
+        monkeypatch.setattr(
+            _util, "compute_default_v2_host_quota", lambda gpu_quota: gpu_quota
+        )
+
+        total_gpu = 10 * GB
+        c = _make_creator(
+            max_gpu_total_bytes=total_gpu,
+            host_cache_size=None,
+            total_kv_per_token=100,
+            target_kv_per_token=80,
+            kv_cache_manager_cls=KVCacheManagerV2,
+        )
+
+        draft_config = c._split_kv_cache_budget_for_draft()
+
+        assert draft_config is not None
+        # Both budgets must now be explicitly set (and positive) so the V2
+        # constructor's auto-provision branch will not fire.
+        assert c._kv_cache_config.host_cache_size is not None
+        assert c._kv_cache_config.host_cache_size > 0
+        assert draft_config.host_cache_size is not None
+        assert draft_config.host_cache_size > 0
+
+    def test_v2_unset_host_sum_matches_auto_provision(self, monkeypatch):
+        """The sum of target+draft host budgets must equal what V2 would
+        have auto-provisioned for a single manager from the combined GPU
+        quota -- otherwise the fix has either over- or under-allocated
+        relative to the V2 baseline."""
+        from tensorrt_llm._torch.pyexecutor import _util
+
+        sentinel_host = 7 * GB
+
+        def fake_default(gpu_quota):
+            assert gpu_quota == 10 * GB  # must be called with combined budget
+            return sentinel_host
+
+        monkeypatch.setattr(_util, "compute_default_v2_host_quota", fake_default)
+
+        c = _make_creator(
+            max_gpu_total_bytes=10 * GB,
+            host_cache_size=None,
+            total_kv_per_token=100,
+            target_kv_per_token=80,
+            kv_cache_manager_cls=KVCacheManagerV2,
+        )
+
+        draft_config = c._split_kv_cache_budget_for_draft()
+
+        target_host = c._kv_cache_config.host_cache_size
+        draft_host = draft_config.host_cache_size
+        # Sum must exactly equal the pre-split auto-provision quota.
+        assert target_host + draft_host == sentinel_host
+        # And the split must follow the same 80/20 ratio as the GPU split.
+        # Use the same draft-first int() truncation order as the implementation.
+        expected_draft = int(sentinel_host * 0.2)
+        expected_target = sentinel_host - expected_draft
+        assert target_host == expected_target
+        assert draft_host == expected_draft
+
+    def test_v1_unset_host_left_alone(self):
+        """V1 callers (no V2 auto-provision in the constructor) must not
+        have a host budget conjured up.  This guards against accidentally
+        adding host memory to spec dec runs that opted out."""
+        c = _make_creator(
+            max_gpu_total_bytes=10 * GB,
+            host_cache_size=None,
+            total_kv_per_token=100,
+            target_kv_per_token=80,
+            kv_cache_manager_cls=KVCacheManager,
+        )
+
+        draft_config = c._split_kv_cache_budget_for_draft()
+
+        assert draft_config is not None
+        assert c._kv_cache_config.host_cache_size is None
+        assert draft_config.host_cache_size is None
+
+    def test_v2_explicit_host_still_split(self, monkeypatch):
+        """When the user did set host_cache_size, that value (not the
+        auto-provision default) must drive the split."""
+        from tensorrt_llm._torch.pyexecutor import _util
+
+        # If the auto-provision path is taken by mistake the test will see
+        # a different total.
+        monkeypatch.setattr(
+            _util, "compute_default_v2_host_quota",
+            lambda _: pytest.fail(
+                "auto-provision must not run when host_cache_size is set"),
+        )
+
+        c = _make_creator(
+            max_gpu_total_bytes=10 * GB,
+            host_cache_size=20 * GB,
+            total_kv_per_token=100,
+            target_kv_per_token=80,
+            kv_cache_manager_cls=KVCacheManagerV2,
+        )
+
+        draft_config = c._split_kv_cache_budget_for_draft()
+        assert c._kv_cache_config.host_cache_size == 16 * GB
+        assert draft_config.host_cache_size == 4 * GB