diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 1095ed03c9f6..f7defa84f8c4 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 2d02ba2a29aa..6e529288b8c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/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( diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py index 4c58321f4987..15fd3c7c630f 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py +++ b/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