From de551b070529cbb903131c203bb91daabe230f80 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 26 Oct 2026 05:59:44 -0700 Subject: [PATCH] [https://nvbugs/6919850][fix] Cap V2 KV cache estimation host tiers KvCacheCreator discarded host_cache_size during estimation, so V2 probes used automatic host quotas without the configured upper bound. Preserve the V2 host budget and cap the automatic probe quota before fleet-min synchronization, rather than reserving the full explicit host tier. Split the cap across target and draft probes. If cache costs cannot split it, retain the target cap and give the draft a zero host cap. Keep normal explicit quotas and V1 host/V2 disk estimation behavior unchanged. Add regression coverage for cap propagation, target/draft sharing, zero and automatic tiers, and device, available-memory and memlock bounds. Fixes https://nvbugs/6919850 Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 38 +++++----- .../kv_cache/kv_cache_manager_v2.py | 11 ++- .../kv_cache/test_kv_cache_budget_split.py | 71 +++++++++++++++++-- .../kv_cache/test_kvcm2_integration.py | 59 +++++++++++++++ 4 files changed, 155 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 7593f2ec6f4c..17e9bf8d0720 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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 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 54b9deae7fc3..bff6b32fe622 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 @@ -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, " 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 bc0f63abe394..478370d0c378 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 @@ -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. diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index addb326f0a91..b4ecc4984a26 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/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(