diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index e0254ee7783b..ead81739a726 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -729,16 +729,15 @@ def _create_one_model_draft_kv_cache_manager( ) def _split_kv_cache_budget_for_draft(self) -> Optional[KvCacheConfig]: - """Split max_gpu_total_bytes between target and draft KV caches. + """Split KV cache budgets between target and draft KV caches. When using KVCacheManagerV2 with a separate draft KV cache, - max_gpu_total_bytes represents the total budget for both target and - draft combined. This method splits the budget proportionally based - on their per-token KV cache sizes. + max_gpu_total_bytes and host_cache_size each represent the total + budget for both target and draft combined. This method splits both + budgets proportionally based on their per-token KV cache sizes. Returns a cloned KvCacheConfig for the draft, or None if no split is - needed. Also modifies self._kv_cache_config.max_gpu_total_bytes - in-place for the target. + needed. Also modifies self._kv_cache_config in-place for the target. """ total_budget = self._kv_cache_config.max_gpu_total_bytes if total_budget is None or total_budget <= 0: @@ -753,7 +752,9 @@ def _split_kv_cache_budget_for_draft(self) -> Optional[KvCacheConfig]: if total_kv <= 0 or draft_kv <= 0: return None - draft_budget = int(total_budget * draft_kv / total_kv) + draft_ratio = draft_kv / total_kv + + draft_budget = int(total_budget * draft_ratio) target_budget = total_budget - draft_budget logger.info( @@ -765,6 +766,18 @@ def _split_kv_cache_budget_for_draft(self) -> Optional[KvCacheConfig]: draft_kv_cache_config = self._kv_cache_config.model_copy() draft_kv_cache_config.max_gpu_total_bytes = draft_budget + + host_budget = self._kv_cache_config.host_cache_size + if host_budget is not None and host_budget > 0: + draft_host_budget = int(host_budget * draft_ratio) + target_host_budget = host_budget - draft_host_budget + self._kv_cache_config.host_cache_size = target_host_budget + draft_kv_cache_config.host_cache_size = draft_host_budget + logger.info( + f"Splitting KV cache host budget: total={host_budget / GB:.2f} GiB, " + f"target={target_host_budget / GB:.2f} GiB, " + f"draft={draft_host_budget / GB:.2f} GiB") + return draft_kv_cache_config def build_managers(self, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index eb7acd36c126..17f538207b11 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -29,6 +29,7 @@ l0_a10: - unittest/_torch/executor/test_async_transfer_manager.py - unittest/_torch/executor/test_scheduler_serializable_output.py - unittest/_torch/executor/test_kv_cache_estimation.py + - unittest/_torch/executor/test_kv_cache_budget_split.py # NOTE: this is a CPU-only test, but we do not have a dedicated job for this (and therefore no # test list either). - unittest/_torch/models/checkpoints/hf/test_weight_loader.py diff --git a/tests/unittest/_torch/executor/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py new file mode 100644 index 000000000000..5edcfcf23fe5 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_budget_split.py @@ -0,0 +1,182 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for _split_kv_cache_budget_for_draft in KvCacheCreator. + +Verifies that both GPU (max_gpu_total_bytes) and host (host_cache_size) +budgets are split proportionally between target and draft KV cache managers. +""" + +from unittest.mock import Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator +from tensorrt_llm.llmapi.llm_args import KvCacheConfig + +GB = 1 << 30 + + +def _make_creator( + max_gpu_total_bytes: int, + host_cache_size=None, + total_kv_per_token: int = 100, + target_kv_per_token: int = 80, +) -> KvCacheCreator: + """Build a minimal KvCacheCreator wired for _split_kv_cache_budget_for_draft.""" + c = object.__new__(KvCacheCreator) + + c._kv_cache_config = KvCacheConfig( + max_gpu_total_bytes=max_gpu_total_bytes, + host_cache_size=host_cache_size, + ) + c._tokens_per_block = 64 + 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._get_kv_size_per_token = Mock(return_value=total_kv_per_token) + + return c + + +class TestSplitKvCacheBudgetForDraft: + def test_gpu_budget_split_proportionally(self): + total_gpu = 10 * GB + c = _make_creator( + max_gpu_total_bytes=total_gpu, total_kv_per_token=100, target_kv_per_token=80 + ) + + draft_config = c._split_kv_cache_budget_for_draft() + + assert draft_config is not None + assert c._kv_cache_config.max_gpu_total_bytes == 8 * GB + assert draft_config.max_gpu_total_bytes == 2 * GB + + def test_host_budget_split_proportionally(self): + total_gpu = 10 * GB + total_host = 20 * GB + c = _make_creator( + max_gpu_total_bytes=total_gpu, + host_cache_size=total_host, + total_kv_per_token=100, + target_kv_per_token=80, + ) + + draft_config = c._split_kv_cache_budget_for_draft() + + assert draft_config is not None + # GPU: 80% target, 20% draft + assert c._kv_cache_config.max_gpu_total_bytes == 8 * GB + assert draft_config.max_gpu_total_bytes == 2 * GB + # Host: same ratio + assert c._kv_cache_config.host_cache_size == 16 * GB + assert draft_config.host_cache_size == 4 * GB + + def test_host_budget_not_doubled(self): + """Regression: before the fix, both target and draft managers each + received the full host_cache_size, doubling total host memory.""" + total_host = 20 * GB + c = _make_creator( + max_gpu_total_bytes=10 * GB, + host_cache_size=total_host, + total_kv_per_token=100, + target_kv_per_token=80, + ) + + 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 + assert target_host + draft_host == total_host + + def test_budgets_sum_to_original(self): + total_gpu = 15 * GB + total_host = 30 * GB + c = _make_creator( + max_gpu_total_bytes=total_gpu, + host_cache_size=total_host, + total_kv_per_token=1000, + target_kv_per_token=700, + ) + + draft_config = c._split_kv_cache_budget_for_draft() + + assert ( + 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 + + def test_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, + ) + + 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_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, + ) + + draft_config = c._split_kv_cache_budget_for_draft() + + assert draft_config is not None + # host_cache_size=0 should not be split (guard: host_budget > 0) + assert draft_config.host_cache_size == 0 + + def test_returns_none_when_no_gpu_budget(self): + c = _make_creator(max_gpu_total_bytes=0) + + assert c._split_kv_cache_budget_for_draft() is None + + def test_returns_none_when_draft_kv_zero(self): + c = _make_creator( + max_gpu_total_bytes=10 * GB, total_kv_per_token=100, target_kv_per_token=100 + ) + + assert c._split_kv_cache_budget_for_draft() is None + + @pytest.mark.parametrize("target_frac", [0.5, 0.75, 0.9, 0.95]) + def test_various_ratios(self, target_frac): + total_gpu = 10 * GB + total_host = 20 * GB + total_kv = 1000 + target_kv = int(total_kv * target_frac) + + c = _make_creator( + max_gpu_total_bytes=total_gpu, + host_cache_size=total_host, + total_kv_per_token=total_kv, + target_kv_per_token=target_kv, + ) + + draft_config = c._split_kv_cache_budget_for_draft() + + assert ( + 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