[https://nvbugs/6035425][fix] Fix host memory usage regression with spec dec

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,

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

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