From c2b979d3be6f017f32df6fa41ca472414d0d47cf Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:24:09 +0000 Subject: [PATCH 1/2] [TRTLLM-12714][fix] Suspend CUDA-graph padding dummies before pool rebalance adjust() CUDAGraphRunner retains one padding dummy request per captured draft length, whose KVCacheManagerV2 cache stays ACTIVE across iterations but never appears in PyExecutor.active_requests. The rebalance hook therefore never suspends them, and the first live adjust() fails its all-caches-suspended precondition, terminating the executor event loop along with all in-flight requests. Suspend the dummies alongside the active requests and resume them after adjust(). Suspension is what the precondition actually asks for: it tears down the cache's base-page-index buffers and releases its page locks, which is what makes the pages safe to migrate. Those buffers are written only when a page lock is taken and are never refreshed when a page later migrates, so a cache left ACTIVE across adjust() can end up addressing slots that now belong to other sequences. The dummies are suspended rather than freed so that the warmup pre-allocation added in #16072 survives a rebalance. Freeing them would return to lazy re-creation against a KV cache that is under load -- rebalance only fires after 2000 sampled caches and a 120s cooldown -- and possibly just shrunk, which is exactly the case where allocation fails and padded batches silently fall back to eager mode for the rest of the process lifetime. A dummy that cannot be resumed is released and dropped from the runner instead, because nothing reschedules a padding dummy and _get_or_create_padding_dummy returns a cached dummy without checking that its cache is live. Reproduced on main with gemma-3-1b-it (VSWA, 2 pool groups) and CudaGraphConfig(enable_padding=True) under TLLM_KV_CACHE_MANAGER_V2_BACKEND=python, where the precondition is a plain assert rather than the C++ backend's debug-gated TLLM_CHECK_DEBUG: adjust() raised AssertionError and killed the executor loop. With this change the same run completes and the GPU pool ratio moves from 0.500/0.500 to 0.667/0.333. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 63 +++++++++ .../_torch/executor/test_kv_pool_rebalance.py | 122 ++++++++++++++++++ 2 files changed, 185 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2fa5a5d0607b..d765c9c0023f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4434,6 +4434,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() @@ -4442,6 +4443,68 @@ 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. + + Note: the release only covers the main KV cache manager, matching + ``_add_cross_dummy_request``'s own failure path. A dummy that also + holds draft or cross-attention KV cache entries keeps those until + the runner is cleared. + """ + 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.") + mgr.free_resources(dummy) + runner.padding_dummy_requests.pop(draft_len, None) @contextmanager def control_action(self, diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index ed0cc5636e7b..fc6a78f09748 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/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,36 @@ 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() + if has_cuda_graph_runner: + exe.model_engine.cuda_graph_runner = MagicMock() + exe.model_engine.cuda_graph_runner.padding_dummy_requests = padding_dummies + 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 +240,99 @@ 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. + """ + dummy = _make_request(999) + exe = _make_executor(padding_dummies={0: dummy}) + exe.kv_cache_manager.resume_request.return_value = False + + self._fire(exe, monkeypatch) + + exe.kv_cache_manager.free_resources.assert_called_once_with(dummy) + assert exe.model_engine.cuda_graph_runner.padding_dummy_requests == {} + + 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 # --------------------------------------------------------------------------- # From 73a0e2887ee88da704a8dbab1e09f81d07939cfd Mon Sep 17 00:00:00 2001 From: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Date: Mon, 10 Aug 2026 16:45:52 +0000 Subject: [PATCH 2/2] [TRTLLM-12714][fix] Centralize CUDA-graph padding dummy release across managers _get_or_create_padding_dummy spreads one dummy request ID across up to four managers: the main KV cache manager, the one-model draft KV cache manager, the speculative resource manager slot and, for encoder-decoder, the 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. Add CUDAGraphRunner.release_padding_dummy(), the release path symmetric with _get_or_create_padding_dummy: it frees the dummy from every manager that allocated part of it and drops it from the runner so a later padded step re-creates it. The manager list lives in _padding_dummy_managers() next to the creation path so the two stay in step, and is deduplicated by identity since double-freeing one manager is not safe in general. The rebalance hook's resume-failure branch now goes through it instead of calling free_resources on the main KV cache manager alone. That branch is the only place the dummies are released at all: the normal path suspends and resumes them. Addresses review feedback from chienchunhung and yizhang-nv on #16157. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> --- .../_torch/pyexecutor/cuda_graph_runner.py | 50 +++++++++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 ++-- .../_torch/executor/test_kv_pool_rebalance.py | 14 +++++- .../executor/test_pytorch_model_engine.py | 43 ++++++++++++++++ 4 files changed, 109 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 07618d87a1a9..fdb5ccefa189 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -697,6 +697,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 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index d765c9c0023f..f0d81d800c96 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4488,10 +4488,9 @@ def _resume_padding_dummies_after_rebalance( one that cannot be resumed is released and dropped from the runner instead, falling back to lazy re-creation on a later padded step. - Note: the release only covers the main KV cache manager, matching - ``_add_cross_dummy_request``'s own failure path. A dummy that also - holds draft or cross-attention KV cache entries keeps those until - the runner is cleared. + 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: @@ -4503,8 +4502,7 @@ def _resume_padding_dummies_after_rebalance( "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.") - mgr.free_resources(dummy) - runner.padding_dummy_requests.pop(draft_len, None) + runner.release_padding_dummy(self.resource_manager, draft_len) @contextmanager def control_action(self, diff --git a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py index fc6a78f09748..d2f12754754e 100644 --- a/tests/unittest/_torch/executor/test_kv_pool_rebalance.py +++ b/tests/unittest/_torch/executor/test_kv_pool_rebalance.py @@ -80,9 +80,14 @@ def _make_executor( # 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 @@ -304,15 +309,20 @@ 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) - exe.kv_cache_manager.free_resources.assert_called_once_with(dummy) - assert exe.model_engine.cuda_graph_runner.padding_dummy_requests == {} + 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) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 9a1242441523..78bfe5818f29 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1448,6 +1448,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",