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..3ee6fb5ca226 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,27 @@ def _add_layer( # number of layers in the KVCacheManagerPy self._num_manager_layers = len(layers) + constraints = list(config.constraints) + 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..35baf5c50fef 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: diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 0e8cf9a4d356..0ad222feeb49 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -3328,6 +3328,14 @@ def free_warmup_requests() -> None: max_num_draft_tokens=_kv_draft) available_tokens = min(available_tokens, draft_available_tokens) + if isinstance(kv_cache_manager, KVCacheManagerV2): + # V2 reserves one generation token beyond the draft/extra tokens. + available_tokens -= 1 + minimum_tokens = ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1 + if available_tokens < minimum_tokens: + free_warmup_requests() + return None + token_num = max( ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM if is_enc_dec else 1, min( 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/executor/kv_cache/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2.py index 55c917fd6f9e..80e3d00ba9e2 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: 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.