[TRTLLM-12714][fix] Suspend CUDA-graph padding dummies before pool rebalance adjust()

tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py

@@ -867,6 +867,56 @@ def _get_or_create_padding_dummy(
         self.padding_dummy_requests[runtime_draft_len] = dummy_request
         return dummy_request
 
+    def _padding_dummy_managers(
+            self,
+            resource_manager: ResourceManager) -> List[BaseResourceManager]:
+        """The managers ``_get_or_create_padding_dummy`` registers a dummy with.
+
+        Kept next to the creation path so the two stay in step.  Duplicates are
+        dropped by identity: freeing the same manager twice for one request is
+        not safe in general.
+        """
+        candidates = [
+            resource_manager.get_resource_manager(
+                self.config.kv_cache_manager_key),
+            get_draft_kv_cache_manager(self.spec_config, resource_manager),
+            resource_manager.get_resource_manager(
+                ResourceManagerType.SPEC_RESOURCE_MANAGER),
+        ]
+        if self.is_encoder_decoder:
+            candidates.append(
+                resource_manager.get_resource_manager(
+                    ResourceManagerType.CROSS_KV_CACHE_MANAGER))
+
+        managers: List[BaseResourceManager] = []
+        for manager in candidates:
+            if manager is not None and not any(manager is seen
+                                               for seen in managers):
+                managers.append(manager)
+        return managers
+
+    def release_padding_dummy(self, resource_manager: ResourceManager,
+                              runtime_draft_len: int) -> bool:
+        """Releases the padding dummy for ``runtime_draft_len`` from every
+        manager that allocated part of it, and drops it from the runner so a
+        later padded step re-creates it.
+
+        One dummy request ID is spread across up to four managers -- the main
+        KV cache manager, the one-model draft KV cache manager, the
+        speculative resource manager slot and the encoder-decoder cross-KV
+        cache manager.  Releasing only the main one leaves the others holding
+        the ID, and re-creation reuses the same
+        ``CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len``.
+
+        Returns True if a dummy was held for that draft length.
+        """
+        dummy_request = self.padding_dummy_requests.pop(runtime_draft_len, None)
+        if dummy_request is None:
+            return False
+        for manager in self._padding_dummy_managers(resource_manager):
+            manager.free_resources(dummy_request)
+        return True
+
     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

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -4465,6 +4465,7 @@ def _maybe_rebalance_kv_pools(self) -> None:
             if mgr.is_request_active(req.py_request_id):
                 mgr.suspend_request(req)
                 paused.append(req)
+        paused_dummies = self._suspend_padding_dummies_for_rebalance(mgr)
 
         try:
             mgr.impl.adjust()
@@ -4473,6 +4474,66 @@ def _maybe_rebalance_kv_pools(self) -> None:
 
         for req in paused:
             mgr.resume_request(req)
+        self._resume_padding_dummies_after_rebalance(mgr, paused_dummies)
+
+    def _suspend_padding_dummies_for_rebalance(
+            self, mgr: KVCacheManagerV2) -> List[Tuple[int, LlmRequest]]:
+        """Suspend the CUDA-graph padding dummies before ``adjust()``.
+
+        ``CUDAGraphRunner`` retains one padding dummy per captured draft
+        length, whose KV cache stays ACTIVE across iterations but never
+        appears in ``active_requests`` -- so the suspend loop above never
+        reaches them and ``adjust()`` fails its all-caches-suspended
+        precondition, terminating the executor event loop.
+
+        Suspending is what the precondition actually asks for: it tears
+        down the cache's base-page-index buffers and releases its page
+        locks, which is exactly what makes the pages safe to migrate.  The
+        dummies are deliberately *not* freed -- they are pre-allocated at
+        warmup because re-allocating one from a loaded KV cache can fail
+        and silently drop padded batches to eager mode for the rest of the
+        process lifetime.
+
+        Returns the ``(draft_len, dummy)`` pairs that were suspended.
+        """
+        runner = getattr(self.model_engine, "cuda_graph_runner", None)
+        if runner is None:
+            return []
+        suspended: List[Tuple[int, LlmRequest]] = []
+        for draft_len, dummy in runner.padding_dummy_requests.items():
+            if mgr.is_request_active(dummy.py_request_id):
+                mgr.suspend_request(dummy)
+                suspended.append((draft_len, dummy))
+        return suspended
+
+    def _resume_padding_dummies_after_rebalance(
+            self, mgr: KVCacheManagerV2,
+            suspended: List[Tuple[int, LlmRequest]]) -> None:
+        """Resume the padding dummies suspended for ``adjust()``.
+
+        A real request that fails to resume is left suspended for the
+        scheduler to reactivate, but nothing reschedules a padding dummy,
+        and ``_get_or_create_padding_dummy`` returns a cached dummy without
+        checking that its cache is live.  A dummy left suspended would
+        therefore be padded into a batch with a torn-down block table, so
+        one that cannot be resumed is released and dropped from the runner
+        instead, falling back to lazy re-creation on a later padded step.
+
+        The release goes through ``CUDAGraphRunner.release_padding_dummy`` so
+        that every manager the dummy was registered with is covered, not just
+        the main KV cache manager.
+        """
+        runner = getattr(self.model_engine, "cuda_graph_runner", None)
+        if runner is None:
+            return
+        for draft_len, dummy in suspended:
+            if mgr.resume_request(dummy):
+                continue
+            logger.warning(
+                "Could not resume the CUDA graph padding dummy request "
+                f"(draft_len={draft_len}) after a KV pool rebalance; "
+                "releasing it and falling back to lazy re-creation.")
+            runner.release_padding_dummy(self.resource_manager, draft_len)
 
     @contextmanager
     def control_action(self,

tests/unittest/_torch/executor/test_kv_pool_rebalance.py

@@ -51,9 +51,14 @@ def _make_executor(
     need_adjustment: bool = True,
     active_requests=None,
     previous_batch=None,
+    padding_dummies=None,
+    has_cuda_graph_runner: bool = True,
 ) -> MagicMock:
     """Construct a MagicMock shaped like PyExecutor with exactly the
     attributes the rebalance code path reads.
+
+    ``padding_dummies`` is the runner's ``{draft_len: dummy}`` map; its
+    dummies count as active on GPU, like real pre-allocated ones.
     """
     exe = MagicMock(spec=PyExecutor)
 
@@ -71,12 +76,41 @@ def _make_executor(
     exe.kv_cache_manager.impl = MagicMock()
     exe.kv_cache_manager.impl.need_adjustment = need_adjustment
 
+    # CUDA-graph runner holding the pre-allocated padding dummies.  A bare
+    # MagicMock would be truthy and non-iterable in the suspend helper.
+    padding_dummies = dict(padding_dummies or {})
+    exe.model_engine = MagicMock()
+    exe.resource_manager = MagicMock()
+    if has_cuda_graph_runner:
+        exe.model_engine.cuda_graph_runner = MagicMock()
+        exe.model_engine.cuda_graph_runner.padding_dummy_requests = padding_dummies
+        # Mirror the real helper: pop the dummy from the runner's map.
+        exe.model_engine.cuda_graph_runner.release_padding_dummy.side_effect = (
+            lambda _rm, draft_len: padding_dummies.pop(draft_len, None) is not None
+        )
+    else:
+        exe.model_engine.cuda_graph_runner = None
+
     # is_request_active returns True for every id we tracked, False for
     # everything else.  Tests set active_requests to a list of mocks with
     # py_request_id attributes.
     exe.active_requests = active_requests or []
     active_ids = {r.py_request_id for r in exe.active_requests}
+    active_ids |= {d.py_request_id for d in padding_dummies.values()}
     exe.kv_cache_manager.is_request_active.side_effect = lambda rid: rid in active_ids
+    exe.kv_cache_manager.resume_request.return_value = True
+
+    # Bind the padding-dummy helpers to their real implementations so that
+    # _maybe_rebalance_kv_pools tests exercise the actual path rather than
+    # MagicMock stand-ins supplied by spec=PyExecutor.
+    exe._suspend_padding_dummies_for_rebalance = (
+        lambda mgr: PyExecutor._suspend_padding_dummies_for_rebalance(exe, mgr)
+    )
+    exe._resume_padding_dummies_after_rebalance = (
+        lambda mgr, suspended: PyExecutor._resume_padding_dummies_after_rebalance(
+            exe, mgr, suspended
+        )
+    )
 
     # Previous batch (overlap loop).
     exe.previous_batch = previous_batch
@@ -211,6 +245,104 @@ def test_unexpected_adjust_failure_propagates(self, monkeypatch):
             PyExecutor._maybe_rebalance_kv_pools(exe)
 
 
+# --------------------------------------------------------------------------- #
+# CUDA-graph padding dummies
+# --------------------------------------------------------------------------- #
+
+
+class TestPaddingDummies:
+    """``adjust()`` requires every living KV cache to be suspended, but the
+    CUDA-graph padding dummies stay ACTIVE across iterations and never appear
+    in ``active_requests``.  The hook must suspend them too -- and must not
+    free them, since they are pre-allocated at warmup precisely so that a
+    loaded KV cache cannot deny them later (PR #16072).
+    """
+
+    @staticmethod
+    def _fire(exe, monkeypatch):
+        exe._consume_previous_batch_for_rebalance = MagicMock()
+        monkeypatch.setattr("torch.cuda.current_stream", MagicMock())
+        PyExecutor._maybe_rebalance_kv_pools(exe)
+
+    def test_padding_dummy_is_suspended_before_adjust(self, monkeypatch):
+        dummy = _make_request(999)
+        exe = _make_executor(active_requests=[_make_request(1)], padding_dummies={0: dummy})
+
+        call_order = []
+        exe.kv_cache_manager.suspend_request.side_effect = lambda req: call_order.append(
+            ("suspend", req.py_request_id)
+        )
+        exe.kv_cache_manager.impl.adjust.side_effect = lambda: call_order.append(("adjust", None))
+
+        self._fire(exe, monkeypatch)
+
+        assert call_order.index(("suspend", 999)) < call_order.index(("adjust", None))
+
+    def test_padding_dummy_is_resumed_and_never_freed(self, monkeypatch):
+        """The dummy must survive the rebalance: freeing it would return to
+        the lazy re-allocation that #16072 removed, which can fail against a
+        loaded cache and drop padded batches to eager mode permanently.
+        """
+        dummy = _make_request(999)
+        runner_map = {0: dummy}
+        exe = _make_executor(padding_dummies=runner_map)
+
+        self._fire(exe, monkeypatch)
+
+        exe.kv_cache_manager.resume_request.assert_called_once_with(dummy)
+        exe.kv_cache_manager.free_resources.assert_not_called()
+        assert exe.model_engine.cuda_graph_runner.padding_dummy_requests == {0: dummy}
+
+    def test_every_captured_draft_length_is_covered(self, monkeypatch):
+        """#16072 retains one dummy per captured draft length, not just one."""
+        dummies = {0: _make_request(999), 3: _make_request(996)}
+        exe = _make_executor(padding_dummies=dummies)
+
+        self._fire(exe, monkeypatch)
+
+        suspended = {
+            c.args[0].py_request_id for c in exe.kv_cache_manager.suspend_request.call_args_list
+        }
+        assert suspended == {999, 996}
+
+    def test_unresumable_padding_dummy_is_released_and_dropped(self, monkeypatch):
+        """Nothing reschedules a padding dummy, and the runner hands back a
+        cached dummy without checking that its cache is live -- so one that
+        cannot be resumed must not be left suspended in the runner's map.
+
+        The release goes through the runner rather than the KV cache manager
+        directly: a dummy is registered with up to four managers, and the
+        runner owns that list.
+        """
+        dummy = _make_request(999)
+        exe = _make_executor(padding_dummies={0: dummy})
+        exe.kv_cache_manager.resume_request.return_value = False
+        runner = exe.model_engine.cuda_graph_runner
+
+        self._fire(exe, monkeypatch)
+
+        runner.release_padding_dummy.assert_called_once_with(exe.resource_manager, 0)
+        exe.kv_cache_manager.free_resources.assert_not_called()
+
+    def test_already_suspended_padding_dummy_is_left_alone(self, monkeypatch):
+        dummy = _make_request(999)
+        exe = _make_executor(padding_dummies={0: dummy})
+        exe.kv_cache_manager.is_request_active.side_effect = lambda rid: False
+
+        self._fire(exe, monkeypatch)
+
+        exe.kv_cache_manager.suspend_request.assert_not_called()
+        exe.kv_cache_manager.resume_request.assert_not_called()
+
+    def test_missing_cuda_graph_runner_is_tolerated(self, monkeypatch):
+        """Engines without a CUDA-graph runner must still rebalance."""
+        exe = _make_executor(active_requests=[_make_request(1)], has_cuda_graph_runner=False)
+
+        self._fire(exe, monkeypatch)
+
+        exe.kv_cache_manager.impl.adjust.assert_called_once()
+
+
 # --------------------------------------------------------------------------- #
 # _consume_previous_batch_for_rebalance
 # --------------------------------------------------------------------------- #

tests/unittest/_torch/executor/test_pytorch_model_engine.py

@@ -1593,6 +1593,49 @@ def test_preallocate_padding_dummies_uses_captured_draft_lens(self):
                 kv_cache_manager.free_resources(dummy)
             kv_cache_manager.shutdown()
 
+    def test_release_padding_dummy_covers_every_manager(self):
+        # A padding dummy's request ID is spread across several managers, so
+        # releasing only the main KV cache manager leaves the others holding
+        # it — and re-creation reuses the same ID.
+        model_engine, kv_cache_manager = create_model_engine_and_kvcache()
+        spec_manager = Mock()
+        cross_manager = Mock()
+        resource_manager = ResourceManager({
+            ResourceManagerType.KV_CACHE_MANAGER:
+            kv_cache_manager,
+            ResourceManagerType.SPEC_RESOURCE_MANAGER:
+            spec_manager,
+            ResourceManagerType.CROSS_KV_CACHE_MANAGER:
+            cross_manager,
+        })
+
+        runner = model_engine.cuda_graph_runner
+        try:
+            self.assertIsNotNone(
+                runner._get_or_create_padding_dummy(resource_manager, 0))
+            dummy = runner.padding_dummy_requests[0]
+
+            self.assertTrue(runner.release_padding_dummy(resource_manager, 0))
+
+            # Dropped from the runner so the lazy path re-creates it...
+            self.assertEqual({}, runner.padding_dummy_requests)
+            # ...and the spec resource manager slot is released too, not just
+            # the main KV cache manager.
+            spec_manager.free_resources.assert_called_once_with(dummy)
+            # The cross-KV manager is only involved for encoder-decoder, which
+            # this engine is not.
+            self.assertFalse(runner.is_encoder_decoder)
+            cross_manager.free_resources.assert_not_called()
+
+            # Releasing again is a no-op rather than a double free.
+            self.assertFalse(runner.release_padding_dummy(resource_manager, 0))
+            spec_manager.free_resources.assert_called_once()
+        finally:
+            for dummy in runner.padding_dummy_requests.values():
+                kv_cache_manager.free_resources(dummy)
+            runner.padding_dummy_requests.clear()
+            kv_cache_manager.shutdown()
+
     def test_layerwise_nvtx_marker(self):
         llm_args = TorchLlmArgs(
             model="dummy",