From 38864aa4e17430853356df473c1f515d2a29e097 Mon Sep 17 00:00:00 2001 From: peter941221 Date: Wed, 10 Jun 2026 17:01:45 +0800 Subject: [PATCH 1/3] Fix unified-memory Mamba KV estimation Signed-off-by: peter941221 --- tensorrt_llm/_torch/pyexecutor/_util.py | 9 ++++++++ .../executor/test_kv_cache_estimation.py | 21 +++++++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 1e96f88621b5..9b6ecd999361 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -11,6 +11,7 @@ from tensorrt_llm._torch.models.modeling_utils import \ MODEL_CLASS_VISION_ENCODER_MAPPING from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, + is_device_integrated, str_dtype_to_binding, torch_dtype_to_str) from tensorrt_llm.bindings.executor import DecodingMode @@ -559,6 +560,14 @@ def _get_token_num_for_estimation(self) -> int: free_mem, total_mem = torch.cuda.mem_get_info() max_memory = self._kv_cache_config.free_gpu_memory_fraction * free_mem kv_size_per_token = self._get_kv_size_per_token() + if kv_size_per_token.intercept > 0 and is_device_integrated(): + # On unified-memory GPUs the free counter can already be depressed + # by mmap-backed weights sharing the same physical pool. During the + # estimation dry run, subtracting the recurrent-state fixed cost + # here can collapse the provisional token cap to zero even though + # post-profiling affine sizing succeeds. Keep the final affine + # sizing unchanged; relax only this provisional cap. + kv_size_per_token = CacheCost(slope=kv_size_per_token.slope) max_num_tokens_in_memory = ( kv_size_per_token.tokens_for_budget(max_memory) // self._tokens_per_block * self._tokens_per_block) diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 33eeda0f2478..3e66d33e8a8d 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -172,6 +172,27 @@ def test_regression_without_fix_would_overcount(): assert result != wrong +def test_integrated_gpu_estimation_ignores_affine_intercept(): + """Unified-memory estimation must not clamp to zero just because the + recurrent-state fixed cost exceeds the mem_get_info-derived budget.""" + tpb = 16 + c = _make_creator( + tpb, + [_make_mock_request(63)], + enable_attention_dp=False, + tp_size=1, + ) + + with ( + patch("torch.cuda.mem_get_info", return_value=(100, 100)), + patch.object(KvCacheCreator, "_get_kv_size_per_token", + return_value=CacheCost(slope=1, intercept=128)), + patch("tensorrt_llm._torch.pyexecutor._util.is_device_integrated", + return_value=True), + ): + assert c._get_token_num_for_estimation() == 64 + + # --------------------------------------------------------------------------- # VSWA hybrid attention pool-group scaling (Gemma4 hybrid MMMU Pro hang fix) # --------------------------------------------------------------------------- From ae539cc2134e54cfde6580782e222a9bed566abf Mon Sep 17 00:00:00 2001 From: peter941221 Date: Wed, 24 Jun 2026 16:30:48 +0800 Subject: [PATCH 2/3] [\#15178][style] Apply pre-commit formatting Signed-off-by: peter941221 --- tensorrt_llm/_torch/pyexecutor/_util.py | 4 ++-- .../unittest/_torch/executor/test_kv_cache_estimation.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9b6ecd999361..bc45d545684c 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -11,8 +11,8 @@ from tensorrt_llm._torch.models.modeling_utils import \ MODEL_CLASS_VISION_ENCODER_MAPPING from tensorrt_llm._utils import (confidential_compute_enabled, get_sm_version, - is_device_integrated, - str_dtype_to_binding, torch_dtype_to_str) + is_device_integrated, str_dtype_to_binding, + torch_dtype_to_str) from tensorrt_llm.bindings.executor import DecodingMode # isort: off diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 3e66d33e8a8d..1dacc3fb1d49 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -185,10 +185,10 @@ def test_integrated_gpu_estimation_ignores_affine_intercept(): with ( patch("torch.cuda.mem_get_info", return_value=(100, 100)), - patch.object(KvCacheCreator, "_get_kv_size_per_token", - return_value=CacheCost(slope=1, intercept=128)), - patch("tensorrt_llm._torch.pyexecutor._util.is_device_integrated", - return_value=True), + patch.object( + KvCacheCreator, "_get_kv_size_per_token", return_value=CacheCost(slope=1, intercept=128) + ), + patch("tensorrt_llm._torch.pyexecutor._util.is_device_integrated", return_value=True), ): assert c._get_token_num_for_estimation() == 64 From d2f6f56c275e245396d010ca751a183c07dc4eee Mon Sep 17 00:00:00 2001 From: "Peter Chen J." Date: Tue, 22 Sep 2026 17:36:14 +0800 Subject: [PATCH 3/3] [#15178][style] Apply ruff-format to KV estimation test Signed-off-by: Peter Chen J. --- .../_torch/executor/kv_cache/test_kv_cache_estimation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 89eddddc209c..149c1d677928 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -350,6 +350,8 @@ def test_regression_without_fix_would_overcount(): wrong = tp * 3 * tpb # 768 (all duplicates summed) assert result == correct assert result != wrong + + def test_integrated_gpu_estimation_ignores_affine_intercept(): """Unified-memory estimation must not clamp to zero just because the recurrent-state fixed cost exceeds the mem_get_info-derived budget.""" @@ -376,8 +378,6 @@ def test_integrated_gpu_estimation_ignores_affine_intercept(): assert c._get_token_num_for_estimation() == 64 - - @pytest.mark.parametrize( ("model_cls", "encoder_cache_max_bytes", "expected_reserve"), [