[TRTLLM-14903][fix] Free partially-allocated warmup dummy KV blocks and count spec extra tokens in warmup block estimates

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):

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
 

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,

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 (AutoDeploy MoE unit tests broken on main since 2026-08-05; nvbug pending)
+unittest/auto_deploy/singlegpu/transformations/library/test_moe_fusion.py SKIP (AutoDeploy MoE unit tests broken on main since 2026-08-05; nvbug pending)
 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)

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.