diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 3fdd8190dd39..fa689f4039da 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -2234,10 +2234,30 @@ def _create_warmup_request( if num_ctx_requests + num_gen_requests > self.batch_size: return None # Not enough batch size to fill the request - blocks_to_use = num_full_seqs * math.ceil( - max_seq_len / kv_cache_manager.tokens_per_block) + math.ceil( - num_left_over_tokens / kv_cache_manager.tokens_per_block - ) + num_gen_requests * self.max_beam_width + # Mirror add_dummy_requests' actual allocation: on top of the raw + # token count, every sequence gets num_extra_kv_tokens + + # num_extra_decoding_steps add_token calls, and generation dummies + # additionally reserve max_draft_loop_tokens for the draft loop. + # In one-engine spec modes that is (max_draft_len - 1) extra KV + # tokens plus max_draft_len draft-loop tokens per gen dummy, i.e. + # 2 * max_draft_len - 1 on top of the single prompt token. + # Under-counting these let warmup start an allocation that fails + # midway and, before the partial-allocation cleanup existed, + # permanently leaked most of the estimation-sized KV pool + # (TRTLLM-14903). + def blocks_for_seq(num_tokens: int) -> int: + return math.ceil(num_tokens / kv_cache_manager.tokens_per_block) + + extra_ctx_tokens = (getattr(kv_cache_manager, "num_extra_kv_tokens", 0) + or 0) + num_extra_decoding_steps + extra_gen_tokens = extra_ctx_tokens + self.max_draft_loop_tokens + blocks_to_use = num_full_seqs * blocks_for_seq(max_seq_len + + extra_ctx_tokens) + if num_left_over_tokens > 0: + blocks_to_use += blocks_for_seq(num_left_over_tokens + + extra_ctx_tokens) + blocks_to_use += (num_gen_requests * self.max_beam_width * + blocks_for_seq(1 + extra_gen_tokens)) if blocks_to_use > available_blocks and isinstance( kv_cache_manager, KVCacheManager): diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 9248c1cbe87f..eb3d4193bd8c 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -986,59 +986,92 @@ def add_dummy_requests( _populate_dummy_mrope_config(req, token_num, is_gen) requests.append(req) - # Use add_sequence_batch for all dummy requests, then add extra tokens. - # This must happen before is_gen state modifications below, which may - # set prompt_len to 0 and trigger assertion in setPrepopulatedPromptLen. - if batch_request_infos: - self.impl.add_sequence_batch(batch_request_infos, - batch_llm_requests) - for req_id, token_num, _ in batch_request_infos: - for _ in range(self.num_extra_kv_tokens): - self.impl.add_token(req_id) - for _ in range(num_extra_decoding_steps): - self.impl.add_token(req_id) - - if draft_batch_request_infos and draft_kv_cache_manager is not None: - draft_kv_cache_manager.impl.add_sequence_batch( - draft_batch_request_infos, draft_batch_llm_requests) - for req_id, _, _ in draft_batch_request_infos: - for _ in range(self.num_extra_kv_tokens): - draft_kv_cache_manager.impl.add_token(req_id) - - # Set is_gen state after add_sequence_batch to avoid modifying - # prompt_len before the C++ side reads it. - if is_gen: - for i, req in enumerate(requests): - token_num = token_nums[ - i] if token_nums is not None else 1 + max_num_draft_tokens - if self.mapping.has_cp_helix(): - token_num = max(token_num, 2) - req.state = LlmRequestState.GENERATION_IN_PROGRESS - req.prompt_len = token_num - 1 - req.py_prompt_len = req.prompt_len - if self.mapping.has_cp_helix(): - if self.mapping.cp_size - 1 == self.mapping.cp_rank: - req.py_helix_is_inactive_rank = False - req.prompt_len = token_num - 1 - req.py_prompt_len = req.prompt_len - req.seqlen_this_rank_cp = req.prompt_len - req.total_input_len_cp = token_num * self.mapping.cp_size - 1 - req.py_decoding_iter = 1 - else: - req.py_helix_is_inactive_rank = True - req.prompt_len = token_num - req.py_prompt_len = req.prompt_len - req.seqlen_this_rank_cp = req.prompt_len - req.total_input_len_cp = token_num * self.mapping.cp_size - 1 - req.py_decoding_iter = 1 - req.py_draft_tokens = [1] * max_num_draft_tokens - if prepare_resource: - for _ in range(_kv_draft): - self.impl.add_token(req.request_id) - if draft_kv_cache_manager is not None: + try: + # Use add_sequence_batch for all dummy requests, then add extra tokens. + # This must happen before is_gen state modifications below, which may + # set prompt_len to 0 and trigger assertion in setPrepopulatedPromptLen. + if batch_request_infos: + self.impl.add_sequence_batch(batch_request_infos, + batch_llm_requests) + for req_id, token_num, _ in batch_request_infos: + for _ in range(self.num_extra_kv_tokens): + self.impl.add_token(req_id) + for _ in range(num_extra_decoding_steps): + self.impl.add_token(req_id) + + if draft_batch_request_infos and draft_kv_cache_manager is not None: + draft_kv_cache_manager.impl.add_sequence_batch( + draft_batch_request_infos, draft_batch_llm_requests) + for req_id, _, _ in draft_batch_request_infos: + for _ in range(self.num_extra_kv_tokens): + draft_kv_cache_manager.impl.add_token(req_id) + + # Set is_gen state after add_sequence_batch to avoid modifying + # prompt_len before the C++ side reads it. + if is_gen: + for i, req in enumerate(requests): + token_num = token_nums[ + i] if token_nums is not None else 1 + max_num_draft_tokens + if self.mapping.has_cp_helix(): + token_num = max(token_num, 2) + req.state = LlmRequestState.GENERATION_IN_PROGRESS + req.prompt_len = token_num - 1 + req.py_prompt_len = req.prompt_len + if self.mapping.has_cp_helix(): + if self.mapping.cp_size - 1 == self.mapping.cp_rank: + req.py_helix_is_inactive_rank = False + req.prompt_len = token_num - 1 + req.py_prompt_len = req.prompt_len + req.seqlen_this_rank_cp = req.prompt_len + req.total_input_len_cp = token_num * self.mapping.cp_size - 1 + req.py_decoding_iter = 1 + else: + req.py_helix_is_inactive_rank = True + req.prompt_len = token_num + req.py_prompt_len = req.prompt_len + req.seqlen_this_rank_cp = req.prompt_len + req.total_input_len_cp = token_num * self.mapping.cp_size - 1 + req.py_decoding_iter = 1 + req.py_draft_tokens = [1] * max_num_draft_tokens + if prepare_resource: for _ in range(_kv_draft): - draft_kv_cache_manager.impl.add_token( - req.request_id) + self.impl.add_token(req.request_id) + if draft_kv_cache_manager is not None: + for _ in range(_kv_draft): + draft_kv_cache_manager.impl.add_token( + req.request_id) + except Exception: + # A partial allocation failure (e.g. add_token raising "no free + # blocks left" after add_sequence_batch succeeded) must not leak + # the sequences already registered. On the minimal KV pool built + # for cache-size estimation, such a leak leaves too few blocks + # for the estimation requests themselves, so the executor loop + # spins forever without ever scheduling them and LLM startup + # hangs (TRTLLM-14903). remove_sequence is a no-op for request + # ids the failed batched add never registered, so every request + # can be removed unconditionally; attempt all target and draft + # cleanup before re-raising so one cleanup failure doesn't leak + # the remaining sequences. + cleanup_error = None + for freeing_impl, freeing_requests in ( + (self.impl, batch_llm_requests), + (draft_kv_cache_manager.impl if draft_kv_cache_manager + is not None else None, draft_batch_llm_requests), + ): + if freeing_impl is None: + continue + for req in freeing_requests: + try: + freeing_impl.remove_sequence(req.py_request_id, req, + False) + except Exception as e: + cleanup_error = cleanup_error or e + if cleanup_error is not None: + # A failed release leaves the KV pool poisoned — the same + # hang mechanism this cleanup exists to prevent — so it + # must not be masked by the allocation failure alone. + raise cleanup_error + raise return requests diff --git a/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py b/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py index 074cdf99485f..56ff2a4ac813 100644 --- a/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py +++ b/tests/integration/defs/kv_cache/test_prefix_aware_scheduling.py @@ -662,6 +662,11 @@ def _run_lmbenchmark( benchmark_debug_context = _stage_debug_context(debug_context, port, server_log, output_csv) def _popen_lmbenchmark(call_args: list[str], env: Mapping[str, str]) -> subprocess.Popen: + # multi-round-qa.py imports its sibling utils.py via the implicit + # script-directory sys.path entry. PYTHONSAFEPATH=1 (set by some CI + # environments) disables that entry and the script dies at import + # time with ModuleNotFoundError: No module named 'utils'. + env = {k: v for k, v in env.items() if k != "PYTHONSAFEPATH"} return subprocess.Popen( call_args, stdout=subprocess.PIPE, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index b2ff70b86a17..f5813c58ff0b 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -433,6 +433,8 @@ unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py::test_allred unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py::test_allreduce_strategies[ONESHOT] SKIP (https://nvbugs/6517839) unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py::test_allreduce_strategies[SYMM_MEM] SKIP (https://nvbugs/6517839) unittest/auto_deploy/multigpu/smoke/test_ad_allreduce_strategies.py::test_allreduce_strategies[TWOSHOT] SKIP (https://nvbugs/6517839) +unittest/auto_deploy/singlegpu/custom_ops/moe/test_trtllm_moe.py SKIP (https://nvbugs/6564714) +unittest/auto_deploy/singlegpu/transformations/library/test_moe_fusion.py SKIP (https://nvbugs/6564714) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_wait_in_progress_on_zero_timeout SKIP (https://nvbugs/6517836) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_write_transfer_gpu_tensor SKIP (https://nvbugs/6517836) unittest/bindings/test_transfer_agent_bindings.py::TestMooncakeFunctionalTransfer::test_mooncake_write_transfer_multiple_chunks SKIP (https://nvbugs/6517836) diff --git a/tests/unittest/_torch/executor/test_resource_manager.py b/tests/unittest/_torch/executor/test_resource_manager.py index 9a230973a85a..f4ab82db1d54 100644 --- a/tests/unittest/_torch/executor/test_resource_manager.py +++ b/tests/unittest/_torch/executor/test_resource_manager.py @@ -960,6 +960,45 @@ def test_batch_cache_indices_honor_requested_blocks_with_beams(self): finally: kv_cache_manager.shutdown() + def test_add_dummy_requests_failure_frees_partial_allocation(self): + """A partial add_dummy_requests failure must free every block it + allocated (TRTLLM-14903): leaked blocks on the minimal pool built for + cache-size estimation starve the estimation requests and hang startup. + """ + kv_cache_manager = KVCacheManager( + kv_cache_config=KvCacheConfig(max_tokens=256, + enable_block_reuse=False), + kv_cache_type=tensorrt_llm.bindings.internal.batch_manager. + CacheType.SELF, + num_layers=2, + num_kv_heads=2, + head_dim=128, + tokens_per_block=64, + max_seq_len=1024, + max_batch_size=2, + mapping=Mapping(), + ) + try: + total_free = kv_cache_manager.get_num_free_blocks() + self.assertEqual(total_free, 4) + # Both sequences fit in one block each, but the per-request draft + # add_token loop needs two more blocks per request: request 0 + # drains the pool and request 1's first add_token raises, after + # three of the four blocks were already allocated. + with self.assertRaises(Exception): + kv_cache_manager.add_dummy_requests([0, 1], + token_nums=[64, 64], + is_gen=True, + max_num_draft_tokens=128) + self.assertEqual(kv_cache_manager.get_num_free_blocks(), total_free) + # The freed pool must serve follow-up allocations. + requests = kv_cache_manager.add_dummy_requests([2], token_nums=[64]) + self.assertIsNotNone(requests) + kv_cache_manager.free_resources(requests[0]) + self.assertEqual(kv_cache_manager.get_num_free_blocks(), total_free) + finally: + kv_cache_manager.shutdown() + def test_kv_cache_manager_with_execution_stream(self): """ Test that KVCacheManager uses the provided execution_stream.