[https://nvbugs/6919850][fix] Cap V2 KV cache estimation host tiers

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -3001,23 +3001,21 @@ def _needs_gpu_kv_cache_budget_split(
             max_seq_len, kv_cache_config)
         return uses_vswa_kv_cache_layout(draft_windows)
 
-    @classmethod
-    def _drop_explicit_offload_tier_budgets(
-            cls, kv_cache_config: Optional[KvCacheConfig]
+    def _prepare_estimation_kv_cache_config(
+            self, kv_cache_config: Optional[KvCacheConfig]
     ) -> Optional[KvCacheConfig]:
-        """Return a copy of the config with explicit offload budgets unset.
+        """Keep V2 host budgets as caps on automatically sized probe tiers.
 
-        Host sizing then falls to the V2 auto host tier policy, which matches
-        the tier to the manager's own device quota; V1 builds no secondary pool.
-        The disk tier is V2 only and has no auto policy, so dropping its budget
-        leaves the tier out.
+        V1 host budgets and V2 disk budgets are unset so estimation does not
+        reserve their full explicit offload tiers.
         """
         if kv_cache_config is None:
             return kv_cache_config
         dropped = {
             attr: None
-            for attr in cls._OFFLOAD_TIER_BUDGET_ATTRS
-            if getattr(kv_cache_config, attr)
+            for attr in self._OFFLOAD_TIER_BUDGET_ATTRS
+            if getattr(kv_cache_config, attr) and (
+                attr != "host_cache_size" or not self._is_kv_cache_manager_v2)
         }
         if not dropped:
             return kv_cache_config
@@ -3026,7 +3024,11 @@ def _drop_explicit_offload_tier_budgets(
     def build_managers(self,
                        resources: Dict,
                        estimating_kv_cache: bool = False) -> None:
-        """Construct KV caches for model and draft model (if applicable)."""
+        """Construct KV caches for model and draft model (if applicable).
+
+        When estimation cannot split a V2 host cap using the cache costs,
+        retain the cap on the target and give the draft no host tier.
+        """
         if self._skip_est:
             self.configure_kv_cache_capacity()
         original_max_seq_len = self._max_seq_len
@@ -3049,14 +3051,10 @@ def build_managers(self,
                                       self_kv_cache_config.enable_block_reuse
                                       and self._joint_reuse_supported())
 
-        # Estimation managers are throwaway probes whose pools only hold dummy
-        # requests, so an explicit offload tier would reserve capacity the probe
-        # cannot fill. Encoder-decoder runs skip estimation, so dropping the
-        # cross budgets is defensive.
         if estimating_kv_cache:
-            self_kv_cache_config = self._drop_explicit_offload_tier_budgets(
+            self_kv_cache_config = self._prepare_estimation_kv_cache_config(
                 self_kv_cache_config)
-            cross_kv_cache_config = self._drop_explicit_offload_tier_budgets(
+            cross_kv_cache_config = self._prepare_estimation_kv_cache_config(
                 cross_kv_cache_config)
 
         # Split combined KV cache budgets before creating managers.
@@ -3078,6 +3076,12 @@ def build_managers(self,
                     self._split_kv_cache_budget_for_draft(
                         budget_attr, self_kv_cache_config,
                         draft_kv_cache_config))
+            if (estimating_kv_cache and self._is_kv_cache_manager_v2
+                    and draft_kv_cache_config is None
+                    and self_kv_cache_config.host_cache_size is not None
+                    and self_kv_cache_config.host_cache_size >= 0):
+                draft_kv_cache_config = self_kv_cache_config.model_copy(
+                    update={"host_cache_size": 0})
 
         # Estimation managers are throwaway probes; the real build that
         # follows runs the check against the same engine and scheduler

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -1919,7 +1919,8 @@ def append_to_kv_heads_per_layer(
         connector_registers_pages = (
             kv_connector_manager is not None and not kv_connector_manager.capacity_only
         )
-        if connector_registers_pages and kv_cache_config.host_cache_size is None:
+        host_cache_size = kv_cache_config.host_cache_size
+        if connector_registers_pages and host_cache_size is None:
             # A connector holds device addresses for its pages, and evicting a
             # page to another tier reassigns its GPU slot underneath it. The
             # automatic host tier only exists to give MAX_UTILIZATION somewhere
@@ -1932,8 +1933,10 @@ def append_to_kv_heads_per_layer(
                 "KV cache manager v2 host tier disabled: a KV connector is attached "
                 "and registers GPU page addresses that tier migration would invalidate."
             )
-        elif 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
+        elif (
+            not self.is_estimating_kv_cache and host_cache_size is not None and host_cache_size >= 0
+        ):
+            host_quota = host_cache_size
         else:
             # The V2 MAX_UTILIZATION scheduler relies on suspend/resume to
             # evict and later restore KV cache pages. Without a secondary
@@ -1969,6 +1972,8 @@ def append_to_kv_heads_per_layer(
                 memlock_limit,
                 shares_device_memory=is_device_integrated(),
             )
+            if self.is_estimating_kv_cache and host_cache_size is not None and host_cache_size >= 0:
+                host_quota = min(host_quota, host_cache_size)
             logger.info(
                 f"KV cache manager v2 auto host tier sizing: "
                 f"{local_ranks} co-located rank(s) on this node, "

tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py

@@ -15,6 +15,7 @@
 """Tests for KV cache budget splitting between target and draft managers."""
 
 import math
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import Mock
 
@@ -991,15 +992,21 @@ def test_offload_budget_is_split_between_target_and_draft(self, budget_attr, tmp
         assert draft_share < total_offload
 
     @pytest.mark.parametrize("budget_attr", OFFLOAD_TIER_BUDGET_ATTRS)
-    def test_estimation_managers_drop_explicit_offload_tiers(self, budget_attr, tmp_path):
-        """An explicit offload tier would reserve capacity a probe cannot fill."""
+    @pytest.mark.parametrize("is_v2", [False, True])
+    def test_estimation_managers_drop_explicit_offload_tiers(self, budget_attr, is_v2, tmp_path):
+        """V2 host budgets cap automatic tiers; other explicit tiers are omitted."""
         c = self._make_build_creator(20 * GB, total_gpu=10 * GB, disk_path=str(tmp_path))
+        c._is_kv_cache_manager_v2 = is_v2
 
         c.build_managers({}, estimating_kv_cache=True)
 
         target_config, draft_config = self._configs_passed_to_managers(c)
-        assert getattr(target_config, budget_attr) is None
-        assert getattr(draft_config, budget_attr) is None
+        if is_v2 and budget_attr == "host_cache_size":
+            assert target_config.host_cache_size == 16 * GB
+            assert draft_config.host_cache_size == 4 * GB
+        else:
+            assert getattr(target_config, budget_attr) is None
+            assert getattr(draft_config, budget_attr) is None
 
     def test_gpu_budget_split_stays_skipped_during_estimation(self, tmp_path):
         """Estimation sizes GPU pools from max_tokens, so it keeps the budget whole."""
@@ -1011,6 +1018,62 @@ def test_gpu_budget_split_stays_skipped_during_estimation(self, tmp_path):
         assert target_config.max_gpu_total_bytes == 10 * GB
         c._needs_gpu_kv_cache_budget_split.assert_not_called()
 
+    @pytest.mark.parametrize("is_v2", [False, True])
+    @pytest.mark.parametrize("separate_draft", [False, True])
+    @pytest.mark.parametrize("host_cache_size", [None, 0, GB, 20 * GB])
+    def test_estimation_preserves_v2_host_budget_caps(
+        self,
+        is_v2: bool,
+        separate_draft: bool,
+        host_cache_size: int | None,
+        tmp_path: Path,
+    ) -> None:
+        creator = self._make_build_creator(20 * GB, total_gpu=10 * GB, disk_path=str(tmp_path))
+        creator._is_kv_cache_manager_v2 = is_v2
+        creator._should_create_separate_draft_kv_cache.return_value = separate_draft
+        creator._kv_cache_config.host_cache_size = host_cache_size
+
+        creator.build_managers({}, estimating_kv_cache=True)
+
+        target_config = creator._create_kv_cache_manager.call_args.kwargs[
+            "kv_cache_config_override"
+        ]
+        expected_target = host_cache_size
+        expected_draft = host_cache_size
+        if host_cache_size and host_cache_size > 0:
+            if is_v2 and separate_draft:
+                expected_draft = host_cache_size // 5
+                expected_target = host_cache_size - expected_draft
+            elif not is_v2:
+                expected_target = expected_draft = None
+        assert target_config.host_cache_size == expected_target
+        assert target_config.disk_cache_size is None
+        if separate_draft:
+            draft_config = creator._create_one_model_draft_kv_cache_manager.call_args.kwargs[
+                "kv_cache_config_override"
+            ]
+            assert draft_config.host_cache_size == expected_draft
+            assert draft_config.disk_cache_size is None
+        else:
+            creator._create_one_model_draft_kv_cache_manager.assert_not_called()
+        assert creator._kv_cache_config.host_cache_size == host_cache_size
+        assert creator._kv_cache_config.disk_cache_size == 20 * GB
+
+    @pytest.mark.parametrize("host_cache_size", [0, GB, 20 * GB])
+    def test_estimation_unknown_costs_do_not_duplicate_host_cap(
+        self, host_cache_size: int, tmp_path: Path
+    ) -> None:
+        creator = self._make_build_creator(20 * GB, total_gpu=10 * GB, disk_path=str(tmp_path))
+        creator._kv_cache_config.host_cache_size = host_cache_size
+        creator._get_target_and_draft_cache_costs = Mock(return_value=None)
+
+        creator.build_managers({}, estimating_kv_cache=True)
+
+        target_config, draft_config = self._configs_passed_to_managers(creator)
+        assert target_config.host_cache_size == host_cache_size
+        assert draft_config.host_cache_size == 0
+        assert creator._kv_cache_config.host_cache_size == host_cache_size
+
 
 class TestExternalDrafterKvDtype:
     """The draft budget must be charged at the dtype the draft pool is ALLOCATED in.

tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py

@@ -380,6 +380,7 @@ def _make_manager_for_cache_tier_test(
     spec_config=None,
     is_draft: bool = False,
     is_disagg: bool = False,
+    is_estimating_kv_cache: bool = False,
     joint_reuse: bool = False,
     mapping: Mapping | None = None,
     vocab_size: int | None = 16,
@@ -452,6 +453,7 @@ def build_cache_config(
             spec_config=spec_config,
             is_draft=is_draft,
             is_disagg=is_disagg,
+            is_estimating_kv_cache=is_estimating_kv_cache,
             joint_kv_cache_reuse=joint_reuse,
             vocab_size=vocab_size,
             execution_stream=Mock(),
@@ -1351,6 +1353,63 @@ def test_avg_seq_len_must_not_exceed_max_seq_len() -> None:
         )
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("estimating", [False, True])
+@pytest.mark.parametrize("host_cache_size", [None, -1, 0, 1 << 20, 128 << 20])
+@pytest.mark.parametrize(
+    "mem_available,local_ranks,memlock_limit,auto_quota",
+    [
+        (1 << 30, 1, -1, 16 << 20),
+        (16 << 20, 2, -1, 4 << 20),
+        (1 << 30, 1, 8 << 20, int((8 << 20) * 0.8)),
+    ],
+)
+def test_estimation_host_tier_respects_explicit_and_auto_bounds(
+    estimating: bool,
+    host_cache_size: int | None,
+    mem_available: int,
+    local_ranks: int,
+    memlock_limit: int,
+    auto_quota: int,
+) -> None:
+    module = "tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2"
+    original_sysconf = os.sysconf
+
+    def sysconf(name: str | int) -> int:
+        if name == "SC_PAGE_SIZE":
+            return 4096
+        if name == "SC_AVPHYS_PAGES":
+            return mem_available // 4096
+        return original_sysconf(name)
+
+    with (
+        patch(f"{module}.os.sysconf", side_effect=sysconf),
+        patch("resource.getrlimit", return_value=(memlock_limit, memlock_limit)),
+        patch(f"{module}.is_device_integrated", return_value=False),
+        patch(
+            f"{module}.Distributed.get", return_value=SimpleNamespace(local_world_size=local_ranks)
+        ),
+        patch(f"{module}._sync_host_tier_quota", side_effect=lambda quota, mapping: quota) as sync,
+    ):
+        manager, impl_constructor = _make_manager_for_cache_tier_test(
+            KvCacheConfig(max_gpu_total_bytes=16 << 20, host_cache_size=host_cache_size),
+            [Mock()],
+            is_estimating_kv_cache=estimating,
+        )
+
+    has_explicit_cap = host_cache_size is not None and host_cache_size >= 0
+    expected_quota = host_cache_size if has_explicit_cap else auto_quota
+    if estimating and has_explicit_cap:
+        expected_quota = min(auto_quota, host_cache_size)
+    cache_tiers = impl_constructor.call_args.args[0].cache_tiers
+    host_quotas = [tier.quota for tier in cache_tiers if isinstance(tier, HostCacheTierConfig)]
+    assert host_quotas == ([expected_quota] if expected_quota > 0 else [])
+    if estimating or not has_explicit_cap:
+        sync.assert_called_once_with(expected_quota, manager.mapping)
+    else:
+        sync.assert_not_called()
+
+
 def test_disk_secondary_tier_enables_eviction(tmp_path) -> None:
     impl = Mock()
     manager, impl_constructor = _make_manager_for_cache_tier_test(