diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py index a367efdd011b..34e652f8dd83 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/cache_manager.py @@ -44,8 +44,10 @@ from tensorrt_llm.runtime import ModelConfig from tensorrt_llm.runtime.kv_cache_manager_v2 import ( AttentionLayerConfig, + BatchDesc, BufferConfig, DataRole, + KVCacheDesc, LayerId, PageIndexMode, ScratchDesc, @@ -1028,7 +1030,7 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """ - Add DeepSeek-V4 layers to the cache config. + Add DeepSeek-V4 layers and warmup constraints to the cache config. """ layers: List[AttentionLayerConfig] = [] layer_attn_to_layer_id: Dict[Tuple[int, DeepseekV4AttentionType], LayerId] = {} @@ -1173,9 +1175,28 @@ def _add_layer( # number of layers in the KVCacheManagerPy self._num_manager_layers = len(layers) + constraints = list(config.constraints) + # _build_base_config copies kv_cache_config.pool_ratio to initial_pool_ratio. + if config.initial_pool_ratio is None: + # DeepSeek-V4's windowed and compressed pools must also support + # the longest decode request alongside the short decode requests. + min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens + constraints.append( + BatchDesc( + [ + KVCacheDesc( + capacity=self.max_seq_len, + history_length=self.max_seq_len - 1, + ) + ] + + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)] + * (self.max_batch_size - 1) + ) + ) return replace( config, layers=layers, + constraints=constraints, ) def _init_indexer_dtype(self, sparse_attn_config: DeepSeekV4SparseAttentionConfig) -> None: 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 0774d0affc9f..635e60b3796f 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 @@ -2724,38 +2724,6 @@ def _build_base_config( * (generation_request_capacity - 1) ) - # CUDA graph generation warmup uses one request at max_seq_len and - # enough minimal decode requests to fill the resident capacity. - if ( - self.max_cuda_graph_batch_size is not None - and self.max_cuda_graph_batch_size > 0 - and self.is_estimating_kv_cache - and all(window is None for window in self.max_attention_window_vec) - ): - # Estimation graph warmup needs the smaller of the resident - # capacity and the largest captured CUDA graph batch. - constraint_batch_size = min( - generation_request_capacity, self.max_cuda_graph_batch_size - ) - else: - constraint_batch_size = generation_request_capacity - constraint_batch_size = max(1, constraint_batch_size) - min_decode_capacity = 1 + self.max_draft_len + self.num_extra_kv_tokens - # Model one request at max_seq_len plus minimal decode requests - # to fill constraint_batch_size. - constraints.append( - BatchDesc( - [ - KVCacheDesc( - capacity=self.max_seq_len, - history_length=self.max_seq_len - 1, - ) - ] - + [KVCacheDesc(capacity=min_decode_capacity, history_length=0)] - * (constraint_batch_size - 1) - ) - ) - # General and chunked-prefill warmup uses one fresh context request # at the per-iteration token budget. if self.max_num_tokens is not None: @@ -4949,7 +4917,8 @@ def release_resources( req.total_input_len_cp = token_num * self._helix_cp_size - 1 req.py_decoding_iter = 1 if prepare_resource: - new_capacity = kv_cache.capacity + _kv_draft + 1 + # token_num already includes the current generation input. + new_capacity = kv_cache.capacity + _kv_draft success = kv_cache.resize(new_capacity, history_length=history_hint) if not success: release_resources(req, free_draft_resources=draft_kv_cache is not None) diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index f31f17446190..4e6c7f7a2a54 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -902,8 +902,7 @@ def create_kv_cache_manager( # Please refer to `tensorrt_llm/_torch/pyexecutor/_util.py` for `kv_cache_manager` config = model_config.pretrained_config - # max_seq_len + 1 because the is_gen path in add_dummy_requests resizes each - # request to capacity + 1; without the extra token the last block rounds down. + # Reserve one token of headroom before rounding to a block boundary. # kv_pool_headroom oversizes max_tokens when the manager splits it across # several pools. DeepSeek-V4 needs 3; the default 1 keeps every other model # on its previous allocation. diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 34688c88513e..0f4a7902a0e4 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6028,7 +6028,8 @@ class TestSeedOss_36B(LlmapiAccuracyTestHarness): @pytest.mark.timeout(14400) @pytest.mark.skip_less_device_memory(140000) def test_auto_dtype(self): - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.8, + use_kv_cache_manager_v2=True) chat_template_kwargs = dict(thinking_budget=-1) with LLM(self.MODEL_PATH, kv_cache_config=kv_cache_config) as llm: diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py b/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py index 15f437d4f8b0..02591600143d 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py @@ -639,7 +639,7 @@ class TestMistralSmall24B(LlmapiAccuracyTestHarness): ids=["forced_chunked_prefill"], ) def test_auto_dtype(self, max_num_tokens): - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75) + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.75, use_kv_cache_manager_v2=True) with LLM( self.MODEL_PATH, kv_cache_config=kv_cache_config, diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py index fda64c2f673d..7dc9edf76c86 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py @@ -53,7 +53,9 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( BatchDesc, + GpuCacheTierConfig, KVCacheDesc, + KVCacheManagerConfig, PageIndexMode, _introspection, ) @@ -75,6 +77,53 @@ def test_typical_seq_len_preserves_deepseek_v4_fallback( assert manager._get_typical_seq_len(KvCacheConfig(avg_seq_len=avg_seq_len)) == expected +@pytest.mark.cpu_only +@pytest.mark.parametrize("pool_ratio", [None, [1.0]]) +def test_build_cache_config_long_decode_constraint(pool_ratio: list[float] | None) -> None: + manager = object.__new__(DeepseekV4CacheManager) + manager.pp_layers = [0] + manager._compress_ratios = [1] + manager.dtype = DataType.BF16 + manager._use_nvfp4_compress = False + manager.head_dim = 512 + 64 + manager.index_head_dim = 128 + manager._indexer_k_dtype = "fp8" + manager.use_fp8_ds_mla = False + manager._swa_window_size = 128 + manager.tokens_per_block = 128 + manager.max_seq_len = 1024 + manager.max_batch_size = 3 + manager.max_draft_len = manager._max_draft_len = 4 + manager.num_extra_kv_tokens = 3 + + context_constraint = BatchDesc([KVCacheDesc(capacity=259, history_length=0)]) + base_config = KVCacheManagerConfig( + tokens_per_block=128, + cache_tiers=[GpuCacheTierConfig(quota=1 << 20)], + layers=[], + constraints=[context_constraint] if pool_ratio is None else [], + initial_pool_ratio=pool_ratio, + ) + + config = manager._build_cache_config(base_config) + + if pool_ratio is None: + assert config.constraints == [ + context_constraint, + BatchDesc( + [ + KVCacheDesc(capacity=1024, history_length=1023), + KVCacheDesc(capacity=8, history_length=0), + KVCacheDesc(capacity=8, history_length=0), + ] + ), + ] + assert base_config.constraints == [context_constraint] + else: + assert config.initial_pool_ratio == pool_ratio + assert config.constraints == [] + + def test_cache_size_estimation_uses_model_attention_layer_count(): class FakeModelConfig: sparse_attention_config = SimpleNamespace( diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 55c917fd6f9e..8aafdf235aef 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py @@ -828,7 +828,7 @@ def test_default_uses_allocator_fallback() -> None: assert config.constraints == [] -def test_avg_seq_len_builds_warmup_constraints() -> None: +def test_avg_seq_len_builds_context_warmup_constraint() -> None: config = _make_cache_config_for_test( KvCacheConfig(host_cache_size=0, avg_seq_len=1024), max_batch_size=3, @@ -841,16 +841,7 @@ def test_avg_seq_len_builds_warmup_constraints() -> None: [KVCacheDesc(capacity=2048, history_length=0)] + [KVCacheDesc(capacity=1024, history_length=1021)] * 2 ) - assert config.constraints == [ - BatchDesc( - [ - KVCacheDesc(capacity=1024, history_length=1023), - KVCacheDesc(capacity=3, history_length=0), - KVCacheDesc(capacity=3, history_length=0), - ] - ), - BatchDesc([KVCacheDesc(capacity=2048, history_length=0)]), - ] + assert config.constraints == [BatchDesc([KVCacheDesc(capacity=2048, history_length=0)])] def test_avg_seq_len_updates_typical_step() -> None: @@ -1104,7 +1095,7 @@ def test_extra_tokens_are_in_context_capacity() -> None: ) assert config.typical_step == BatchDesc([KVCacheDesc(capacity=258, history_length=0)]) - assert config.constraints[1] == BatchDesc([KVCacheDesc(capacity=258, history_length=0)]) + assert config.constraints == [BatchDesc([KVCacheDesc(capacity=258, history_length=0)])] def test_try_commit_blocks_commits_partial_block_at_context_end() -> None: @@ -1510,6 +1501,44 @@ def test_external_draft_estimated_quota_supports_allocation_and_resume( manager.shutdown() +@pytest.mark.parametrize("draft_len", [0, 4]) +def test_generation_dummy_uses_available_capacity(draft_len: int) -> None: + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + init_cuda_once() + manager = KVCacheManagerV2( + KvCacheConfig(enable_block_reuse=False, max_gpu_total_bytes=4 << 20), + CacheType.SELF, + num_layers=2, + num_kv_heads=2, + head_dim=128, + tokens_per_block=32, + max_seq_len=131072, + max_batch_size=1, + max_num_tokens=128, + mapping=Mapping(), + dtype=DataType.HALF, + spec_config=MTPDecodingConfig(max_draft_len=draft_len) if draft_len else None, + ) + try: + token_num = manager.get_num_available_tokens( + token_num_upper_bound=manager.max_seq_len, max_num_draft_tokens=draft_len + ) + capacity = token_num + manager.num_extra_kv_tokens + draft_len + assert capacity % manager.tokens_per_block == 0 + # The current generation input is already included in token_num. + requests = manager.add_dummy_requests( + [0], token_nums=[token_num], is_gen=True, max_num_draft_tokens=draft_len + ) + assert requests is not None + cache = manager.kv_cache_map[requests[0].py_request_id] + assert cache.history_length == token_num - 1 + assert cache.capacity == capacity + manager.free_resources(requests[0]) + finally: + manager.shutdown() + + @pytest.fixture def max_num_turns() -> int: return 1 diff --git a/tests/unittest/grpc/smg/test_smg.py b/tests/unittest/grpc/smg/test_smg.py index bbf6f5a31717..8cd1909517ef 100644 --- a/tests/unittest/grpc/smg/test_smg.py +++ b/tests/unittest/grpc/smg/test_smg.py @@ -842,7 +842,7 @@ def grpc_vlm_service(): model_path = get_model_path(vlm_model_name) llm = LLM( model=model_path, - kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.6), + kv_cache_config=KvCacheConfig(free_gpu_memory_fraction=0.6, use_kv_cache_manager_v2=True), load_format="dummy", ) tokenizer = llm.tokenizer diff --git a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py index 03019f68f0fd..b45d2b9d1716 100644 --- a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py +++ b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import io import os import tempfile @@ -43,6 +46,7 @@ def temp_extra_llm_api_options_file(request): "kv_cache_config": { "enable_block_reuse": False, "free_gpu_memory_fraction": 0.6, + "use_kv_cache_manager_v2": True, }, } diff --git a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py index de7aa43fa301..c5a03e59ef17 100644 --- a/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py +++ b/tests/unittest/llmapi/apps/_test_trtllm_serve_multimodal_example.py @@ -40,6 +40,7 @@ def temp_extra_llm_api_options_file(request): "kv_cache_config": { "enable_block_reuse": False, "free_gpu_memory_fraction": 0.6, + "use_kv_cache_manager_v2": True, }, "max_num_tokens": 16384, # for pytorch backend # NOTE: This is for video support.