From 927fe352d1c5174ca9da4201ebccacbebb2f7229 Mon Sep 17 00:00:00 2001
From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Date: Tue, 7 Jul 2026 08:59:52 -0700
Subject: [PATCH 1/4] [None][fix] Pre-allocate CUDA graph padding dummy during
 warmup

The CUDA graph padding dummy request was allocated lazily at the first
padded step. If the KV cache was already saturated by that point, the
allocation failed on every subsequent step and padded batches silently
fell back to eager mode for the rest of the process lifetime.

Allocate the draft_len-0 padding dummy at the end of warmup instead,
while the KV cache still has free blocks. Refactor the dummy creation
into _get_or_create_padding_dummy() (keeping the encoder-decoder
handling) and log a warning_once when padding falls back to eager so
the failure mode is visible.

Extend test_pytorch_model_engine.py to assert the dummy exists after
warmup and that no KV blocks leak beyond it.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
---
 .../_torch/pyexecutor/cuda_graph_runner.py    | 153 +++++++++++-------
 .../_torch/pyexecutor/model_engine.py         |   6 +
 .../executor/test_pytorch_model_engine.py     |  11 +-
 3 files changed, 114 insertions(+), 56 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index e0bd91ac52a8..eda612af863d 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -481,8 +481,6 @@ def replay(self, key: KeyType,
     def _get_padded_batch(self, batch: ScheduledRequests,
                           resource_manager: ResourceManager,
                           runtime_draft_len: int) -> int:
-        kv_cache_manager = resource_manager.get_resource_manager(
-            self.config.kv_cache_manager_key)
         can_run_cuda_graph = batch.can_run_cuda_graph
         batch_size = batch.batch_size
         new_batch_size = batch_size
@@ -520,66 +518,111 @@ def _get_padded_batch(self, batch: ScheduledRequests,
         if padding_size + batch.batch_size > self.config.batch_size:
             return 0
 
-        runtime_tokens_per_gen_step = (
-            self.spec_config.get_runtime_tokens_per_gen_step(runtime_draft_len)
-            if self.spec_config is not None else 1 + runtime_draft_len)
-        runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1
-
         # No padding if it would create too many concurrent requests.
         # This is not strictly required, but we should probably
         # respect the requirement just in case that changes in the future.
         # Use per-draft-len dummy requests for dynamic draft length support.
-        if runtime_draft_len not in self.padding_dummy_requests:
-            dummy_encoder_output_len = None
-            if self.is_encoder_decoder:
-                cross_kv_cache_manager = resource_manager.get_resource_manager(
-                    ResourceManagerType.CROSS_KV_CACHE_MANAGER)
-                if cross_kv_cache_manager is None:
-                    return 0
-                dummy_encoder_output_len = self._get_padding_dummy_encoder_output_len(
-                    cross_kv_cache_manager)
-
-            # Get draft KV cache manager only for one-model speculative decoding.
-            # In two-model mode, each model has its own KV cache manager, so
-            # draft_kv_cache_manager should be None.
-            draft_kv_cache_manager = get_draft_kv_cache_manager(
-                self.spec_config, resource_manager)
-
-            # Use unique dummy request ID per draft length
-            dummy_request_id = CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len
-            dummy_request = kv_cache_manager.add_dummy_requests(
-                [dummy_request_id],
-                token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM]
-                if self.is_encoder_decoder else None,
-                is_gen=True,
-                max_num_draft_tokens=runtime_draft_token_buffer_width,
-                use_mrope=self.config.use_mrope,
-                max_beam_width=self.config.max_beam_width,
-                encoder_output_lens=[dummy_encoder_output_len]
-                if dummy_encoder_output_len is not None else None,
-                draft_kv_cache_manager=draft_kv_cache_manager)
-
-            if dummy_request is None:
-                return 0
-            else:
-                dummy_request = dummy_request[0]
-            dummy_request.is_cuda_graph_dummy = True
-            if self.is_encoder_decoder:
-                if not self._add_cross_dummy_request(
-                        dummy_request, resource_manager,
-                        dummy_encoder_output_len, draft_kv_cache_manager):
-                    return 0
-
-            spec_res_mgr = resource_manager.get_resource_manager(
-                ResourceManagerType.SPEC_RESOURCE_MANAGER)
-            if spec_res_mgr:
-                spec_res_mgr.add_dummy_requests([dummy_request_id])
-            self.padding_dummy_requests[runtime_draft_len] = dummy_request
-
-        padding_dummy_request = self.padding_dummy_requests[runtime_draft_len]
+        padding_dummy_request = self._get_or_create_padding_dummy(
+            resource_manager, runtime_draft_len)
+        if padding_dummy_request is None:
+            logger.warning_once(
+                "Failed to allocate the CUDA graph padding dummy request "
+                f"(draft_len={runtime_draft_len}) from a saturated KV cache; "
+                "falling back to eager mode for padded batches.",
+                key=f"cuda_graph_padding_dummy_fallback_{runtime_draft_len}")
+            return 0
+
         batch.generation_requests.extend([padding_dummy_request] * padding_size)
         return padding_size
 
+    def _get_or_create_padding_dummy(
+            self, resource_manager: ResourceManager,
+            runtime_draft_len: int) -> Optional[LlmRequest]:
+        """Returns the padding dummy request for the given draft length,
+        allocating it from the KV cache manager on first use.
+
+        Returns None when the KV cache manager cannot allocate the dummy
+        (e.g. saturated cache); the batch cannot be padded until a retry
+        succeeds.
+        """
+        if runtime_draft_len in self.padding_dummy_requests:
+            return self.padding_dummy_requests[runtime_draft_len]
+
+        kv_cache_manager = resource_manager.get_resource_manager(
+            self.config.kv_cache_manager_key)
+
+        runtime_tokens_per_gen_step = (
+            self.spec_config.get_runtime_tokens_per_gen_step(runtime_draft_len)
+            if self.spec_config is not None else 1 + runtime_draft_len)
+        runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1
+
+        dummy_encoder_output_len = None
+        if self.is_encoder_decoder:
+            cross_kv_cache_manager = resource_manager.get_resource_manager(
+                ResourceManagerType.CROSS_KV_CACHE_MANAGER)
+            if cross_kv_cache_manager is None:
+                return None
+            dummy_encoder_output_len = self._get_padding_dummy_encoder_output_len(
+                cross_kv_cache_manager)
+
+        # Get draft KV cache manager only for one-model speculative decoding.
+        # In two-model mode, each model has its own KV cache manager, so
+        # draft_kv_cache_manager should be None.
+        draft_kv_cache_manager = get_draft_kv_cache_manager(
+            self.spec_config, resource_manager)
+
+        # Use unique dummy request ID per draft length
+        dummy_request_id = CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len
+        dummy_request = kv_cache_manager.add_dummy_requests(
+            [dummy_request_id],
+            token_nums=[ENC_DEC_CUDA_GRAPH_DUMMY_TOKEN_NUM]
+            if self.is_encoder_decoder else None,
+            is_gen=True,
+            max_num_draft_tokens=runtime_draft_token_buffer_width,
+            use_mrope=self.config.use_mrope,
+            max_beam_width=self.config.max_beam_width,
+            encoder_output_lens=[dummy_encoder_output_len]
+            if dummy_encoder_output_len is not None else None,
+            draft_kv_cache_manager=draft_kv_cache_manager)
+
+        if dummy_request is None:
+            return None
+        dummy_request = dummy_request[0]
+        dummy_request.is_cuda_graph_dummy = True
+        if self.is_encoder_decoder:
+            if not self._add_cross_dummy_request(
+                    dummy_request, resource_manager, dummy_encoder_output_len,
+                    draft_kv_cache_manager):
+                return None
+
+        spec_res_mgr = resource_manager.get_resource_manager(
+            ResourceManagerType.SPEC_RESOURCE_MANAGER)
+        if spec_res_mgr:
+            spec_res_mgr.add_dummy_requests([dummy_request_id])
+        self.padding_dummy_requests[runtime_draft_len] = dummy_request
+        return dummy_request
+
+    def preallocate_padding_dummy(self,
+                                  resource_manager: ResourceManager) -> None:
+        """Eagerly allocates the draft_len-0 padding dummy while the KV cache
+        still has free blocks (called at the end of ModelEngine.warmup).
+
+        _get_padded_batch otherwise allocates the dummy lazily at the first
+        padded step; when the KV cache is already saturated by then, the
+        allocation fails on every step and padded batches permanently fall
+        back to eager mode.
+        """
+        if not (self.enabled and self.padding_enabled):
+            return
+        if self._get_or_create_padding_dummy(resource_manager, 0) is None:
+            logger.warning(
+                "Could not pre-allocate the CUDA graph padding dummy request "
+                "at warmup; allocation will be retried at the first padded "
+                "step.")
+        else:
+            logger.info(
+                "Pre-allocated the CUDA graph padding dummy request at warmup.")
+
     def _add_cross_dummy_request(
             self, dummy_request: LlmRequest, resource_manager: ResourceManager,
             encoder_output_len: int,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index a1048249d08d..f95abfda5b72 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1092,6 +1092,12 @@ def warmup(self, resource_manager: ResourceManager) -> None:
             self._general_warmup(resource_manager, warmup_requests_configs)
             log_mem_snapshot("warmup/after_memory_pool_prepop")
 
+        # Allocate the CUDA graph padding dummy now, while the KV cache is
+        # empty. Waiting for the first padded step can race KV saturation:
+        # once the cache is full, the lazy allocation in _get_padded_batch
+        # fails every step and padded batches silently run eager.
+        self.cuda_graph_runner.preallocate_padding_dummy(resource_manager)
+
     def _general_warmup(self, resource_manager: ResourceManager,
                         warmup_requests_configs: List[Tuple[int, int]]):
         """
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
index e5497d999e23..ae13eec2128e 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -337,7 +337,16 @@ def test_warmup(self):
 
         num_free_before = kv_cache_manager.get_num_free_blocks()
         model_engine.warmup(resource_manager)
-        # Make sure we don't leak any blocks.
+
+        # Warmup pre-allocates the CUDA graph padding dummy (draft_len 0),
+        # which intentionally keeps holding its KV blocks so padding cannot
+        # fall back to eager once the cache saturates.
+        padding_dummies = model_engine.cuda_graph_runner.padding_dummy_requests
+        self.assertIn(0, padding_dummies)
+
+        # Make sure we don't leak any blocks beyond that dummy: freeing it
+        # must restore the exact pre-warmup free-block count.
+        kv_cache_manager.free_resources(padding_dummies[0])
         self.assertEqual(num_free_before,
                          kv_cache_manager.get_num_free_blocks())
 

From 8e19a60600efadb81fad2f8fec661de3dcf06cc6 Mon Sep 17 00:00:00 2001
From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Date: Thu, 9 Jul 2026 02:30:27 -0700
Subject: [PATCH 2/4] [TRTLLM-14138][fix] Only preallocate padding dummies
 runtime padding can use

The warmup-time preallocation broke three spec-decode CI tests:

- It always allocated the draft_len-0 dummy, which spec-decode runtime
  padding never requests (it pads with the runtime draft length), so the
  dummy permanently held KV blocks and hybrid-cache slots on top of the
  lazily created one. TestQwen3_5_4B::test_dflash failed with "No free
  blocks left" for the full-attention window pool.
- It ran even when padding cannot occur at all (max_batch_size=1 with a
  batch-size-1 graph in test_eagle3_cdl_sampling), retaining resources
  the lazy path never consumes.
- It ran during the KV cache capacity estimation phase, whose cache has
  zero headroom when the user-provided max_tokens clamps the estimate;
  the held block left the estimation request unschedulable and wedged
  executor initialization (test_eagle3_cdl_sampling,
  test_dflash_qwen3_8b).

Fix:

- Preallocate one dummy per draft length of the captured CUDA graphs -
  exactly the set runtime padding requests, which also covers dynamic
  draft-length schedules - and only when padding can actually occur for
  that draft length (_can_pad_any_batch mirrors _get_padded_batch's
  rounding and capacity guards).
- Skip preallocation for estimation-phase KV cache managers.
  KvCacheCreator.build_managers now stamps is_estimating_kv_cache on the
  managers it builds.
- Extend test_pytorch_model_engine.py with tests for the three gates.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |   9 ++
 .../_torch/pyexecutor/cuda_graph_runner.py    |  67 ++++++++---
 .../_torch/pyexecutor/model_engine.py         |   4 +-
 .../executor/test_pytorch_model_engine.py     | 104 +++++++++++++++---
 4 files changed, 154 insertions(+), 30 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 97b2ca3a82f1..95058af8a3ef 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1605,6 +1605,15 @@ def build_managers(self,
                 cross_kv_cache_config, estimating_kv_cache,
                 original_max_seq_len)
 
+        # Let consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies)
+        # distinguish the throwaway estimation-phase managers from the final
+        # ones: the estimation cache is sized with no headroom for retained
+        # dummy requests.
+        for manager in (kv_cache_manager, draft_kv_cache_manager,
+                        cross_kv_cache_manager):
+            if manager is not None:
+                manager.is_estimating_kv_cache = estimating_kv_cache
+
         resources[ResourceManagerType.KV_CACHE_MANAGER] = kv_cache_manager
         resources[
             ResourceManagerType.DRAFT_KV_CACHE_MANAGER] = draft_kv_cache_manager
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index eda612af863d..a4ff278ff642 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -602,26 +602,67 @@ def _get_or_create_padding_dummy(
         self.padding_dummy_requests[runtime_draft_len] = dummy_request
         return dummy_request
 
-    def preallocate_padding_dummy(self,
-                                  resource_manager: ResourceManager) -> None:
-        """Eagerly allocates the draft_len-0 padding dummy while the KV cache
-        still has free blocks (called at the end of ModelEngine.warmup).
+    def _can_pad_any_batch(self, runtime_draft_len: int) -> bool:
+        """Returns True when _get_padded_batch can pad at least one feasible
+        batch size for the given draft length (mirrors its rounding and
+        capacity guards). For example, with graph batch sizes [1, 2, 4] and
+        max batch size 1, every batch already matches a graph size and no
+        padding dummy is ever needed.
+        """
+        use_dynamic_draft_len = (
+            self.spec_config and self.spec_config.draft_len_schedule
+            and self.spec_config.spec_dec_mode.support_dynamic_draft_len())
+        max_unpadded_batch_size = min(self.config.batch_size,
+                                      self.max_supported_batch_size)
+        for batch_size in range(1, max_unpadded_batch_size + 1):
+            padded = (self._round_up_batch_size_with_draft_len(
+                batch_size, runtime_draft_len) if use_dynamic_draft_len else
+                      self._round_up_batch_size(batch_size))
+            if batch_size < padded <= self.config.batch_size:
+                return True
+        return False
+
+    def preallocate_padding_dummies(self,
+                                    resource_manager: ResourceManager) -> None:
+        """Eagerly allocates the padding dummies while the KV cache still has
+        free blocks (called at the end of ModelEngine.warmup, after graph
+        capture).
 
-        _get_padded_batch otherwise allocates the dummy lazily at the first
+        _get_padded_batch otherwise allocates the dummies lazily at the first
         padded step; when the KV cache is already saturated by then, the
         allocation fails on every step and padded batches permanently fall
         back to eager mode.
+
+        Only the draft lengths of captured graphs are preallocated — those are
+        exactly the ones runtime padding requests — and only when padding can
+        actually occur for that draft length. Anything else would permanently
+        hold KV blocks (and spec/hybrid cache slots) that the lazy path never
+        consumes, regressing tightly-sized deployments.
         """
         if not (self.enabled and self.padding_enabled):
             return
-        if self._get_or_create_padding_dummy(resource_manager, 0) is None:
-            logger.warning(
-                "Could not pre-allocate the CUDA graph padding dummy request "
-                "at warmup; allocation will be retried at the first padded "
-                "step.")
-        else:
-            logger.info(
-                "Pre-allocated the CUDA graph padding dummy request at warmup.")
+        kv_cache_manager = resource_manager.get_resource_manager(
+            self.config.kv_cache_manager_key)
+        if kv_cache_manager is None or getattr(kv_cache_manager,
+                                               'is_estimating_kv_cache', False):
+            # The estimation-phase KV cache is sized with no headroom for
+            # retained dummies; holding blocks there can leave the estimation
+            # requests unschedulable. That executor is discarded anyway, so
+            # preallocation only matters for the final one.
+            return
+        for draft_len in sorted({key[1] for key in self.graphs}):
+            if not self._can_pad_any_batch(draft_len):
+                continue
+            if self._get_or_create_padding_dummy(resource_manager,
+                                                 draft_len) is None:
+                logger.warning(
+                    "Could not pre-allocate the CUDA graph padding dummy "
+                    f"request (draft_len={draft_len}) at warmup; allocation "
+                    "will be retried at the first padded step.")
+            else:
+                logger.info(
+                    "Pre-allocated the CUDA graph padding dummy request "
+                    f"(draft_len={draft_len}) at warmup.")
 
     def _add_cross_dummy_request(
             self, dummy_request: LlmRequest, resource_manager: ResourceManager,
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index f95abfda5b72..0419dd921bae 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1092,11 +1092,11 @@ def warmup(self, resource_manager: ResourceManager) -> None:
             self._general_warmup(resource_manager, warmup_requests_configs)
             log_mem_snapshot("warmup/after_memory_pool_prepop")
 
-        # Allocate the CUDA graph padding dummy now, while the KV cache is
+        # Allocate the CUDA graph padding dummies now, while the KV cache is
         # empty. Waiting for the first padded step can race KV saturation:
         # once the cache is full, the lazy allocation in _get_padded_batch
         # fails every step and padded batches silently run eager.
-        self.cuda_graph_runner.preallocate_padding_dummy(resource_manager)
+        self.cuda_graph_runner.preallocate_padding_dummies(resource_manager)
 
     def _general_warmup(self, resource_manager: ResourceManager,
                         warmup_requests_configs: List[Tuple[int, int]]):
diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
index ae13eec2128e..274edb71d6c2 100644
--- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py
+++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py
@@ -111,18 +111,21 @@ def create_model_engine_and_kvcache(llm_args: TorchLlmArgs = None,
     num_layers = 1
     batch_size = 13
 
-    llm_args = llm_args if llm_args else TorchLlmArgs(
-        model="dummy",
-        max_batch_size=batch_size,
-        max_num_tokens=max_tokens,
-        cuda_graph_config=CudaGraphConfig(
-            enable_padding=True, batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128]))
-    test_batches = (5, 13)
-    for batch_size in test_batches:
-        assert batch_size not in llm_args.cuda_graph_config.batch_sizes
+    if llm_args is None:
+        llm_args = TorchLlmArgs(model="dummy",
+                                max_batch_size=batch_size,
+                                max_num_tokens=max_tokens,
+                                cuda_graph_config=CudaGraphConfig(
+                                    enable_padding=True,
+                                    batch_sizes=[1, 2, 4, 8, 16, 32, 64, 128]))
+        # The padding tests below rely on these properties of the default
+        # config: batches of 5 and 13 must round up to 8 and 16 respectively.
+        test_batches = (5, 13)
+        for test_batch_size in test_batches:
+            assert test_batch_size not in llm_args.cuda_graph_config.batch_sizes
 
-    assert (8 in llm_args.cuda_graph_config.batch_sizes
-            and 16 in llm_args.cuda_graph_config.batch_sizes)
+        assert (8 in llm_args.cuda_graph_config.batch_sizes
+                and 16 in llm_args.cuda_graph_config.batch_sizes)
 
     model_engine = DummyModelEngine(llm_args, torch.half)
 
@@ -338,11 +341,12 @@ def test_warmup(self):
         num_free_before = kv_cache_manager.get_num_free_blocks()
         model_engine.warmup(resource_manager)
 
-        # Warmup pre-allocates the CUDA graph padding dummy (draft_len 0),
-        # which intentionally keeps holding its KV blocks so padding cannot
-        # fall back to eager once the cache saturates.
+        # Warmup pre-allocates the CUDA graph padding dummy for each captured
+        # draft length (only draft_len 0 here, no speculation), which
+        # intentionally keeps holding its KV blocks so padding cannot fall
+        # back to eager once the cache saturates.
         padding_dummies = model_engine.cuda_graph_runner.padding_dummy_requests
-        self.assertIn(0, padding_dummies)
+        self.assertEqual([0], sorted(padding_dummies.keys()))
 
         # Make sure we don't leak any blocks beyond that dummy: freeing it
         # must restore the exact pre-warmup free-block count.
@@ -352,6 +356,76 @@ def test_warmup(self):
 
         kv_cache_manager.shutdown()
 
+    def test_warmup_skips_padding_dummy_when_padding_impossible(self):
+        # With max_batch_size=1 and a graph for batch size 1, every batch
+        # already matches a graph size, so no padded batch can ever occur and
+        # warmup must not retain a padding dummy (it would permanently hold
+        # KV blocks and spec/hybrid resources that are never used).
+        llm_args = TorchLlmArgs(model="dummy",
+                                max_batch_size=1,
+                                max_num_tokens=258,
+                                cuda_graph_config=CudaGraphConfig(
+                                    enable_padding=True, batch_sizes=[1, 2, 4]))
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache(
+            llm_args)
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+
+        num_free_before = kv_cache_manager.get_num_free_blocks()
+        model_engine.warmup(resource_manager)
+
+        self.assertEqual({},
+                         model_engine.cuda_graph_runner.padding_dummy_requests)
+        self.assertEqual(num_free_before,
+                         kv_cache_manager.get_num_free_blocks())
+
+        kv_cache_manager.shutdown()
+
+    def test_warmup_skips_padding_dummy_during_estimation(self):
+        # The estimation-phase KV cache is sized with no headroom for retained
+        # dummies; holding blocks there can leave the estimation requests
+        # unschedulable. Warmup must skip the preallocation for managers
+        # created for KV cache capacity estimation.
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache()
+        kv_cache_manager.is_estimating_kv_cache = True
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+
+        num_free_before = kv_cache_manager.get_num_free_blocks()
+        model_engine.warmup(resource_manager)
+
+        self.assertEqual({},
+                         model_engine.cuda_graph_runner.padding_dummy_requests)
+        self.assertEqual(num_free_before,
+                         kv_cache_manager.get_num_free_blocks())
+
+        kv_cache_manager.shutdown()
+
+    def test_preallocate_padding_dummies_uses_captured_draft_lens(self):
+        # Speculative engines pad batches with the runtime draft length, so
+        # the preallocation must create dummies for the draft lengths of the
+        # captured graphs — not draft_len 0 unconditionally (an extra
+        # draft_len-0 dummy would permanently hold KV blocks that runtime
+        # padding never uses).
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache()
+        resource_manager = ResourceManager(
+            {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager})
+
+        runner = model_engine.cuda_graph_runner
+        # Simulate a spec-decode capture: graphs keyed (batch_size, draft_len,
+        # is_first_draft, short_seq_len_mode, is_all_greedy_sample) exist only
+        # for draft_len 2.
+        runner.graphs[(8, 2, False, False, True)] = Mock()
+        try:
+            runner.preallocate_padding_dummies(resource_manager)
+
+            self.assertEqual([2], sorted(runner.padding_dummy_requests.keys()))
+        finally:
+            runner.graphs.clear()
+            for dummy in runner.padding_dummy_requests.values():
+                kv_cache_manager.free_resources(dummy)
+            kv_cache_manager.shutdown()
+
     def test_layerwise_nvtx_marker(self):
         llm_args = TorchLlmArgs(
             model="dummy",

From 2e222c4a8c058f21e839539b54b7a0f4351df423 Mon Sep 17 00:00:00 2001
From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Date: Thu, 9 Jul 2026 03:04:47 -0700
Subject: [PATCH 3/4] [TRTLLM-14138][refactor] Simplify estimation gating and
 padding round-up dispatch

Store the existing is_estimating_kv_cache constructor flag on
KVCacheManager/KVCacheManagerV2 instead of stamping it onto manager
instances in KvCacheCreator.build_managers; the constructors already
receive the flag at every creation site, and MambaHybridCacheManager
and the AutoDeploy shim already treat it as a constructor parameter.

Drop the duplicated dynamic-draft-len enablement check from
_get_padded_batch and _can_pad_any_batch: the dynamic mapping is only
populated when the feature is active, so
_round_up_batch_size_with_draft_len already reduces to plain
batch-size rounding without it.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |  9 ------
 .../_torch/pyexecutor/cuda_graph_runner.py    | 32 +++++++++----------
 .../_torch/pyexecutor/kv_cache_manager_v2.py  |  6 ++++
 .../_torch/pyexecutor/resource_manager.py     |  5 +++
 4 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 95058af8a3ef..97b2ca3a82f1 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1605,15 +1605,6 @@ def build_managers(self,
                 cross_kv_cache_config, estimating_kv_cache,
                 original_max_seq_len)
 
-        # Let consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies)
-        # distinguish the throwaway estimation-phase managers from the final
-        # ones: the estimation cache is sized with no headroom for retained
-        # dummy requests.
-        for manager in (kv_cache_manager, draft_kv_cache_manager,
-                        cross_kv_cache_manager):
-            if manager is not None:
-                manager.is_estimating_kv_cache = estimating_kv_cache
-
         resources[ResourceManagerType.KV_CACHE_MANAGER] = kv_cache_manager
         resources[
             ResourceManagerType.DRAFT_KV_CACHE_MANAGER] = draft_kv_cache_manager
diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
index a4ff278ff642..17d4ef3bca7e 100644
--- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
+++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py
@@ -499,15 +499,12 @@ def _get_padded_batch(self, batch: ScheduledRequests,
                 or new_batch_size > self.max_supported_batch_size):
             return 0
 
-        # When dynamic draft length is enabled (one-model path), we treat the determined runtime draft length
-        # as the source of truth and pad the batch size up to the nearest existing graph
-        # for that draft length.
-        if (self.spec_config and self.spec_config.draft_len_schedule
-                and self.spec_config.spec_dec_mode.support_dynamic_draft_len()):
-            padded_batch_size = self._round_up_batch_size_with_draft_len(
-                new_batch_size, runtime_draft_len)
-        else:
-            padded_batch_size = self._round_up_batch_size(new_batch_size)
+        # Pad the batch size up to the nearest existing graph for the runtime
+        # draft length. With dynamic draft length (one-model path) the runtime
+        # draft length is the source of truth for which graphs are eligible;
+        # otherwise this reduces to plain batch-size rounding.
+        padded_batch_size = self._round_up_batch_size_with_draft_len(
+            new_batch_size, runtime_draft_len)
 
         if batch_size == padded_batch_size:
             return 0
@@ -609,15 +606,11 @@ def _can_pad_any_batch(self, runtime_draft_len: int) -> bool:
         max batch size 1, every batch already matches a graph size and no
         padding dummy is ever needed.
         """
-        use_dynamic_draft_len = (
-            self.spec_config and self.spec_config.draft_len_schedule
-            and self.spec_config.spec_dec_mode.support_dynamic_draft_len())
         max_unpadded_batch_size = min(self.config.batch_size,
                                       self.max_supported_batch_size)
         for batch_size in range(1, max_unpadded_batch_size + 1):
-            padded = (self._round_up_batch_size_with_draft_len(
-                batch_size, runtime_draft_len) if use_dynamic_draft_len else
-                      self._round_up_batch_size(batch_size))
+            padded = self._round_up_batch_size_with_draft_len(
+                batch_size, runtime_draft_len)
             if batch_size < padded <= self.config.batch_size:
                 return True
         return False
@@ -713,9 +706,14 @@ def _round_up_batch_size(self, batch_size: int) -> int:
 
     def _round_up_batch_size_with_draft_len(self, batch_size: int,
                                             draft_len: int) -> int:
-        """Finds the smallest graph batch size >= batch_size that also matches the given draft_len."""
+        """Finds the smallest graph batch size >= batch_size that also matches the given draft_len.
+
+        The dynamic draft length mapping exists exactly when the dynamic draft
+        length feature is active (see _compute_dynamic_draft_len_mapping);
+        without it, this ignores draft_len and reduces to plain batch-size
+        rounding.
+        """
         if not self.dynamic_draft_len_mapping:
-            # Fallback to regular round up if no mapping
             return self._round_up_batch_size(batch_size)
 
         start_idx = bisect.bisect_left(self.supported_batch_sizes, batch_size)
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index 2f93f0fb028f..4120dc77f827 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -626,6 +626,7 @@ def __init__(
         execution_stream: Optional[torch.cuda.Stream] = None,
         is_disagg: bool = False,
         enable_stats: bool = False,
+        is_estimating_kv_cache: bool = False,
         **kwargs,
     ) -> None:
         self.mapping = mapping
@@ -648,6 +649,11 @@ def __init__(
             layer_mask=layer_mask,
         )
         self.is_draft = is_draft
+        # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies)
+        # can distinguish the throwaway estimation-phase managers from the
+        # final ones: the estimation cache is sized with no headroom for
+        # retained dummy requests.
+        self.is_estimating_kv_cache = is_estimating_kv_cache
         self.enable_swa_scratch_reuse = (
             kv_cache_config.enable_swa_scratch_reuse and not self.is_draft
         )
diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
index 7ef8d802878b..bd45dd915e37 100644
--- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py
+++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py
@@ -294,6 +294,11 @@ def __init__(
             layer_mask=layer_mask,
         )
         self.is_draft = is_draft
+        # Retained so consumers (e.g. CUDAGraphRunner.preallocate_padding_dummies)
+        # can distinguish the throwaway estimation-phase managers from the
+        # final ones: the estimation cache is sized with no headroom for
+        # retained dummy requests.
+        self.is_estimating_kv_cache = is_estimating_kv_cache
         self.num_local_layers = len(self.pp_layers)
         self.layer_offsets = {
             idx: offset

From 4d646f78ed2b3df03b86ca50f401e3bd35558a38 Mon Sep 17 00:00:00 2001
From: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
Date: Sun, 12 Jul 2026 22:44:59 -0700
Subject: [PATCH 4/4] [TRTLLM-14138][fix] Share spec padding-dummy slots
 between target and draft engines

In two-model speculation the target and draft engines share one spec
resource manager, and each engine's CUDAGraphRunner registers its
padding dummy under the same draft-length-derived request ID
(CUDA_GRAPH_DUMMY_REQUEST_ID - draft_len). SlotManager.add_slot only
tolerates re-adding the exact sentinel ID, so the second registration
died with an AssertionError. Warmup preallocation made this
deterministic: the target engine warms up first and registers the dummy
for draft_len == max_draft_len, and the draft engine's first-draft
graphs use the same draft length, so its preallocation re-registered
the same ID and killed the executor during init
(test_eagle3_cuda_graph_padding CI failures on pipeline 47174).

The lazy runtime path has the same latent collision, but the test never
reached it: with 2048-token generations against the 4096-token KV cache
the guaranteed-no-evict scheduler caps the batch at two requests, which
exactly matches a CUDA graph batch size, so runtime padding never
engages.

Make dummy registration idempotent in Eagle3ResourceManager and
MTPHiddenStatesManager (mirroring SuffixAutomatonManager): reserve the
slot once and share it between the engines. Dummy outputs are
discarded, so sharing the slot is safe.

Shorten the generations in the non-overlap variant of
test_eagle3_cuda_graph_padding so runtime padding actually engages and
covers the shared-slot path end to end. The overlap variant keeps long
generations for now: padded draft batches trip a pre-existing shape
mismatch in Eagle3SpecMetadata.prepare on the incremental-update path,
which is unrelated to this change and tracked separately.

Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
---
 tensorrt_llm/_torch/speculative/eagle3.py     |  8 +++-
 tensorrt_llm/_torch/speculative/mtp.py        |  8 +++-
 .../_torch/speculative/test_eagle3.py         | 46 ++++++++++++++++++-
 3 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index 3978040724b1..00251f25b48c 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -131,7 +131,13 @@ def free_resources(self, request: LlmRequest):
 
     def add_dummy_requests(self, request_ids: List[int]):
         for rid in request_ids:
-            self.slot_manager.add_slot(rid)
+            # In two-model speculation the target and draft engines share this
+            # resource manager, and each registers its own padding dummy under
+            # the same draft-length-derived request ID. Reserve the slot once
+            # and share it between them; dummy outputs are discarded, so the
+            # slot contents never matter (mirrors SuffixAutomatonManager).
+            if self.slot_manager.get_slot(rid) is None:
+                self.slot_manager.add_slot(rid)
         if self.sa_manager is not None:
             self.sa_manager.add_dummy_requests(request_ids)
 
diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py
index 6afed0fc02e9..a7820b7f6bc8 100644
--- a/tensorrt_llm/_torch/speculative/mtp.py
+++ b/tensorrt_llm/_torch/speculative/mtp.py
@@ -106,7 +106,13 @@ def free_resources(self, request: LlmRequest):
 
     def add_dummy_requests(self, request_ids: List[int]):
         for rid in request_ids:
-            self.slot_manager.add_slot(rid)
+            # The target and draft engines share this resource manager and
+            # each registers its padding dummy under the same
+            # draft-length-derived request ID. Reserve the slot once and share
+            # it between them; dummy outputs are discarded, so the slot
+            # contents never matter (mirrors SuffixAutomatonManager).
+            if self.slot_manager.get_slot(rid) is None:
+                self.slot_manager.add_slot(rid)
         if self.sa_manager is not None:
             self.sa_manager.add_dummy_requests(request_ids)
 
diff --git a/tests/unittest/_torch/speculative/test_eagle3.py b/tests/unittest/_torch/speculative/test_eagle3.py
index 4564ab550526..2850531c1982 100644
--- a/tests/unittest/_torch/speculative/test_eagle3.py
+++ b/tests/unittest/_torch/speculative/test_eagle3.py
@@ -134,6 +134,39 @@ def test_eagle3_one_model_capture_uses_real_token_count() -> None:
     assert spec_metadata.hidden_states.shape == (4, 2)
 
 
+def test_eagle3_resource_manager_shares_padding_dummy_slot() -> None:
+    """The target and draft engines of two-model EAGLE3 share one
+    Eagle3ResourceManager, and each registers its own CUDA graph padding dummy
+    under the same draft-length-derived request ID
+    (CUDA_GRAPH_DUMMY_REQUEST_ID - draft_len). The second registration must
+    reuse the already-reserved slot instead of tripping the strict re-add
+    assert in SlotManager.add_slot."""
+    from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \
+        CUDA_GRAPH_DUMMY_REQUEST_ID
+    from tensorrt_llm._torch.speculative.eagle3 import Eagle3ResourceManager
+
+    config = Eagle3DecodingConfig(max_draft_len=4,
+                                  speculative_model="/dummy/eagle3")
+    manager = Eagle3ResourceManager(config,
+                                    torch.half,
+                                    hidden_size=8,
+                                    max_num_requests=4,
+                                    max_seq_len=32,
+                                    max_num_tokens=64)
+
+    dummy_request_id = CUDA_GRAPH_DUMMY_REQUEST_ID - config.max_draft_len
+    # The target engine registers the padding dummy first (e.g. during warmup
+    # preallocation), then the draft engine registers the same ID.
+    manager.add_dummy_requests([dummy_request_id])
+    dummy_slot = manager.slot_manager.get_slot(dummy_request_id)
+    manager.add_dummy_requests([dummy_request_id])
+    assert manager.slot_manager.get_slot(dummy_request_id) == dummy_slot
+
+    # Real request IDs still get their own slots.
+    real_slot = manager.slot_manager.add_slot(7)
+    assert real_slot != dummy_slot
+
+
 @pytest.fixture(scope="function")
 def enforce_single_worker(monkeypatch):
     monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1")
@@ -879,7 +912,18 @@ def test_eagle3_cuda_graph_padding(disable_overlap_scheduler: bool):
         "The future of AI is",
     ]
 
-    sampling_params = SamplingParams(max_tokens=2048, temperature=0)
+    # Short generations keep all three requests within the KV cache budget so
+    # they run concurrently and batches of 3 actually pad up to the graph
+    # batch size 4. With long generations (2048 tokens against the 4096-token
+    # cache) the guaranteed-no-evict scheduler caps the batch at two requests,
+    # which exactly matches a graph batch size, so runtime padding (and the
+    # shared padding-dummy slot this test exists to cover) never engages.
+    # The overlap-scheduler variant must keep the long generations for now:
+    # padded draft batches trip a pre-existing shape mismatch in
+    # Eagle3SpecMetadata.prepare on the incremental-update path
+    # (hidden_states_read_indices sized without the padded requests).
+    max_tokens = 64 if disable_overlap_scheduler else 2048
+    sampling_params = SamplingParams(max_tokens=max_tokens, temperature=0)
     llm_spec.generate(prompts, sampling_params)
     llm_spec.shutdown()