[TRTLLM-12714][feat] KVCacheManagerV2: wire PyExecutor rebalance hook (single GPU, aggregated for now) by thorjohnsen · Pull Request #14578 · NVIDIA/TensorRT-LLM

thorjohnsen · 2026-05-26T14:57:59Z

…-GPU agg)

The V2 KVCacheManager.adjust() auto-tuner is implemented but had zero production callers. This patch wires a minimal between-iteration hook in PyExecutor._executor_loop that suspends every active request, calls adjust(), and resumes -- gated to single-GPU aggregated serving with no beam search, drafter, or disagg transceiver.

The fast path costs one property read per iteration; adjust() itself is guarded by V2's existing 2000-sample + 120s cooldown gate. Pool VAs are preserved by the CUDA VMM-backed pool design, so captured CUDA graphs remain valid across resize. Resume failures stay suspended and the scheduler reactivates them through the existing prepare_context / try_allocate_generation path on the next iteration.

Adds a public resume_request() on KVCacheManagerV2 as a thin wrapper around the existing _resume_and_restore helper.

Out of scope (deferred): overlap scheduler, PP collective vote, disagg transceiver quiescence, draft V2 mirror, PEFT cleanup on suspend.

Summary by CodeRabbit

Release Notes

New Features
- Added prototype KV cache pool rebalance capability to automatically adjust memory allocation ratios during token generation for improved efficiency.
Performance
- Optimized KV cache storage manager pool shrinking with improved fast-path handling for reduced overhead.
Tests
- Added comprehensive test coverage for pool rebalance behavior, accuracy validation, and slot allocator operations.

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

…-GPU agg) The V2 KVCacheManager.adjust() auto-tuner is implemented but had zero production callers. This patch wires a minimal between-iteration hook in PyExecutor._executor_loop that suspends every active request, calls adjust(), and resumes -- gated to single-GPU aggregated serving with no beam search, drafter, or disagg transceiver. The fast path costs one property read per iteration; adjust() itself is guarded by V2's existing 2000-sample + 120s cooldown gate. Pool VAs are preserved by the CUDA VMM-backed pool design, so captured CUDA graphs remain valid across resize. Resume failures stay suspended and the scheduler reactivates them through the existing prepare_context / try_allocate_generation path on the next iteration. Adds a public resume_request() on KVCacheManagerV2 as a thin wrapper around the existing _resume_and_restore helper. Out of scope (deferred): overlap scheduler, PP collective vote, disagg transceiver quiescence, draft V2 mirror, PEFT cleanup on suspend. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

…uler Adds the same between-iteration hook to PyExecutor._executor_loop_overlap that the prior commit added to _executor_loop, plus an always-drain step inside _maybe_rebalance_kv_pools so it works correctly in both loops. The overlap loop runs iter N+1's GPU work in parallel with CPU consumption of iter N's results, so at the top of any iteration previous_batch may hold an unconsumed in-flight iteration. Suspending its _KVCache instances mid-flight would race against running kernels. The new _consume_previous_batch_for_rebalance helper drains GPU work via torch.cuda.current_stream().synchronize() and processes previous_batch inline (mirroring the same _update_requests / _send_kv_async / _flush_pending_transfer_responses / _process_previous_batch sequence the loop performs). In the non-overlap loop previous_batch is always None, so the helper is a no-op there -- no special case. Cost: one pipeline-bubble per rebalance, fired at most once every 120s (the V2 auto-tuner's existing cooldown). Gates are unchanged from the prior commit; the drain mechanism is what makes overlap-mode safe, not narrower scope. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

…hook Adds ``KvCacheConfig.enable_kv_pool_rebalance`` (default True, status="prototype") as a kill switch for the PyExecutor rebalance hook introduced in the prior two commits. Plumbed through _util.py:create_py_executor_instance into PyExecutor's constructor and checked at the head of _can_pause_for_rebalance. When False, the hook is skipped entirely -- pool ratios remain at their warmup-derived values. Useful as an escape hatch for workloads where the suspend/adjust/resume cost or the 120s-cadence pipeline bubbles are undesirable, and for debugging regressions to isolate whether the hook is responsible. The flag only takes effect when use_kv_cache_manager_v2=True; the V1 manager has no rebalance code path. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

The rebalance hook is a beta feature -- ship it opt-in rather than opt-out. Users must now set ``kv_cache_config.enable_kv_pool_rebalance = True`` to invoke the V2 auto-tuner from PyExecutor; the prior default (True) silently activated it for everyone on V2. Docstring and ``_can_pause_for_rebalance`` comment updated to reflect opt-in semantics. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

…nce hook Adds tests for the rebalance hook landed in the prior commits on this branch. * tests/unittest/_torch/executor/test_kv_pool_rebalance.py -- 14 functional tests covering every short-circuit branch of _can_pause_for_rebalance, the suspend/adjust/resume cycle of _maybe_rebalance_kv_pools (including the kill-switch and the "adjust failed but resume still runs" path), and the overlap-mode drain helper _consume_previous_batch_for_rebalance. No real engine: calls methods as unbound class lookups on a MagicMock(spec=PyExecutor), following the existing test_py_executor pattern. All pass locally. * tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py -- Greedy-decode token-exact comparison on Gemma-3-1B with explicit VSWA (so we get >=2 pool groups and adjust() has real shrink work to do). Inline _inject_pool_ratio_mismatch helper bypasses the 2000-sample / 120s cooldown gates and skews target ratios past the 1.25x mismatch threshold. Parametrized on overlap mode. Requires TLLM_WORKER_USE_SINGLE_PROCESS=1 so the test process can reach the in-proc PyExecutor for injection. The accuracy test currently fails: when the hook fires for real, KVCacheManagerV2._storage_manager.shrink_pool_group asserts at L668 on an overflow-slot count mismatch. This is a pre-existing bug in the V2 shrink path that has gone undetected because adjust() had zero production callers before this branch. The hook plumbing itself is correct (the unit tests prove the call chain). Per the team's decision, the two accuracy parametrizations are added to waives.txt as SKIP with a tag pointing at the V2 bug; the tests are kept in-repo as a repro and unwaived when the bug is fixed. Also wires the new tests into l0_a10 (unit) and l0_h100 (accuracy) test lists. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

…acy waives to NVBug Swaps the placeholder tag (kvcm-v2-shrink-pool-group-assert) on the two test_rebalance_matches_baseline parametrizations for the real https://nvbugs/6225866 URL, matching the rest of waives.txt's convention. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

… shrink The shrink path in `_storage_manager.shrink_pool_group` asserted that `len(_overflow_slots) == _num_active_slots - _target_capacity`, and `finish_shrink` gated finalization on the same equation. Both implicitly assumed _num_active_slots > _target_capacity. When the auto-tuner asks to shrink a pool whose new size is still above the high-water mark of slot IDs ever issued, the RHS goes negative and the assertion fires. `_num_active_slots` is a strict monotone high-water mark of issued slot IDs, so slot IDs in the to-be-removed range have never been handed out. There is genuinely no migration to do. Two changes: - `_storage/_core.py:finish_shrink` -- relax the gate to `max(0, _num_active_slots - _target_capacity)`, which collapses to `0 == 0` in the underused case. The body already handled this case correctly via `min(_num_active_slots, _capacity)`. - `_storage_manager.py:shrink_pool_group` -- add an early-return fast path when `_num_active_slots <= new_num_slots`: call `prepare_for_shrink` + `finish_shrink` + `resize_pools` and return, skipping the migration scan entirely. Added `TestSlotAllocatorShrink` with a direct repro of the NVBug condition and a sanity check that the non-trivial migration path still finalizes correctly. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

…curacy tests The shrink-path bug tracked in NVBug 6225866 is fixed in the previous commit. Both parametrizations of test_rebalance_matches_baseline pass locally on H100 with Gemma-3-1B + VSWA + injected ratio mismatch (token-exact match between rebalance=off and rebalance=on, no_overlap and overlap scheduler paths). Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

…sertion Pre-commit's ruff-format hook prefers the trailing-message-as-call style for assert with a long expression. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

coderabbitai · 2026-05-27T22:23:07Z

📝 Walkthrough

Walkthrough

This PR implements automatic KV cache pool rebalance tuning in PyExecutor. A new enable_kv_pool_rebalance configuration flag gates the feature; when enabled and conditions are met, the executor pauses active requests, invokes the KV cache manager's ratio-adjustment logic, then resumes them—all integrated into the main execution loop to avoid disrupting ongoing inference.

Changes

KV Pool Rebalance Implementation

Layer / File(s)	Summary
Configuration and Executor Wiring `tensorrt_llm/llmapi/llm_args.py`, `tensorrt_llm/_torch/pyexecutor/py_executor.py`, `tensorrt_llm/_torch/pyexecutor/_util.py`	`KvCacheConfig.enable_kv_pool_rebalance` (prototype flag, default `False`) flows through `PyExecutor.__init__` parameter and is forwarded by the executor factory function.
PyExecutor Rebalance Logic `tensorrt_llm/_torch/pyexecutor/py_executor.py`	Adds `_can_pause_for_rebalance()` gate, `_consume_previous_batch_for_rebalance()` to drain KV state in overlap mode, and `_maybe_rebalance_kv_pools()` to synchronize CUDA, suspend requests, call manager `adjust()`, and resume requests. Integrated into both `_executor_loop` and `_executor_loop_overlap` before batch scheduling.
KV Manager Request Resumption `tensorrt_llm/_torch/pyexecutor/resource_manager.py`	Adds `KVCacheManagerV2.resume_request()` method to restore previously-suspended KV caches by delegating to `_resume_and_restore()`.
Storage Manager Shrink Optimizations `tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py`, `tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py`	Refines `SlotAllocator.finish_shrink()` consistency checks, and introduces fast-path in `StorageManager.shrink_pool_group()` to skip eviction/defrag work when no active slots exist in the target removal range.
PyExecutor Rebalance Unit Tests `tests/unittest/_torch/executor/test_kv_pool_rebalance.py`	Comprehensive unit tests with mock helpers validating gate logic across configuration branches, full rebalance cycle with request suspend/resume, exception handling, and batch consumption flows.
Slot Allocator Shrink Unit Tests `tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`	Regression tests covering `SlotAllocator` shrink for underused and partially-released pools, verifying proper completion without overflow leaks.
Integration Accuracy Tests `tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py`, `tests/integration/test_lists/test-db/l0_a10.yml`, `tests/integration/test_lists/test-db/l0_h100.yml`	Integration accuracy suite that verifies token-for-token output equivalence with and without rebalance by injecting mid-generation pool ratio mismatches; parameterized over scheduler overlap modes; test list updates for A10 and H100 pre-merge pipelines.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Suggested reviewers

leslie-fang25
Shixiaowei02
yizhang-nv
joyang-nv
lfr-0531

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 35.14% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The PR description is largely incomplete. The required sections for Description and Test Coverage are missing; only an informal summary precedes the empty template sections.	Fill in the Description section explaining the issue/solution, and the Test Coverage section listing relevant tests. The informal text above the template is insufficient as a formal PR description.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: wiring a PyExecutor rebalance hook for KVCacheManagerV2, with appropriate scoping details (single GPU, aggregated).
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 2

🧹 Nitpick comments (3)

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)

2423-2458: QA list update is unnecessary for this PR scope.

These changes are unittest-only and the added coverage is aligned with the storage shrink regression scope, so no tests/integration/test_lists/qa/* update is needed here.
As per coding guidelines: "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py` around
lines 2423 - 2458, The PR needs an explicit statement that QA list updates are
unnecessary because the change is limited to unit tests; update the PR
description (or commit message) to say something like "QA list update
unnecessary: changes are unittest-only
(tests/unittest/kv_cache_manager_v2_tests::TestSlotAllocatorShrink ->
test_shrink_underused_pool, test_shrink_touched_pool) and only add coverage for
storage shrink regression," so reviewers can easily see this follows the coding
guideline about unit-scope changes.

tests/unittest/_torch/executor/test_kv_pool_rebalance.py (1)

180-192: ⚡ Quick win

Remove unused caplog parameter.

The test function declares caplog as a parameter but never uses it. Remove it to avoid confusion and keep the test signature clean.

♻️ Proposed fix

-    def test_adjust_failure_does_not_skip_resume(self, monkeypatch, caplog):
+    def test_adjust_failure_does_not_skip_resume(self, monkeypatch):

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/executor/test_kv_pool_rebalance.py` around lines 180 -
192, The test test_adjust_failure_does_not_skip_resume currently includes an
unused caplog parameter; remove caplog from the function signature so the test
becomes def test_adjust_failure_does_not_skip_resume(self, monkeypatch): and
update any callers or decorators if present; keep the rest of the body intact
(references to _make_request, _make_executor, exe, monkeypatch, and assertions
on exe.kv_cache_manager) and ensure imports/mocks are unchanged so
PyExecutor._maybe_rebalance_kv_pools behavior is still validated.

tensorrt_llm/llmapi/llm_args.py (1)

2681-2693: 💤 Low value

Align description terminology with status field.

The description text says "Beta: enable at your own risk" but the status field is set to "prototype". According to the Field() docstring (line 90), "prototype" means "Not yet stable and subject to breaking changes; intended for experimentation only", while "beta" means "Recommended for use per the latest documentation".

For consistency, consider changing "Beta:" to "Prototype:" in the description, or simply remove the "Beta:" prefix since the status="prototype" already conveys the stability level.

📝 Suggested fix to align terminology

     enable_kv_pool_rebalance: bool = Field(
         default=False,
         status="prototype",
         description=
         "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for "
         "rebalancing pool-group ratios between iterations. When True the "
         "PyExecutor calls ``adjust()`` opportunistically; the auto-tuner "
         "itself remains gated by V2's internal 2000-sample / 120s cooldown. "
         "When False (default) the rebalance hook is skipped entirely and "
-        "pool ratios remain at their warmup-derived values. Beta: enable at "
-        "your own risk. Only used when using KV cache manager v2 "
+        "pool ratios remain at their warmup-derived values. Only used when "
+        "using KV cache manager v2 (experimental)."
-        "(experimental)."
     )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/llmapi/llm_args.py` around lines 2681 - 2693, The description
for the Field declaration of enable_kv_pool_rebalance is inconsistent with its
status="prototype": update the text to match the status by replacing the "Beta:"
prefix with "Prototype:" (or remove the stability prefix entirely) so the
description aligns with the Field(status="prototype") metadata; locate the
enable_kv_pool_rebalance Field in llm_args.py and edit its description string
accordingly to reflect "Prototype:" (or omit the prefix) and keep the rest of
the wording intact.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2647-2668: The _can_pause_for_rebalance gate currently only
excludes pipeline parallelism; update it to enforce true single-rank execution
by adding a world-size check: in the _can_pause_for_rebalance method (and any
callers that assume single-rank semantics such as
_consume_previous_batch_for_rebalance), return False when self.dist.world_size
!= 1 (or > 1) so multi-rank TP/DP/CP runs cannot enter the rebalance path;
locate the function _can_pause_for_rebalance and add the explicit
self.dist.world_size check alongside the existing pp_size check.
- Around line 2721-2724: The current broad except around mgr.impl.adjust()
swallows all exceptions; narrow it to only catch the expected rebalance failure
by replacing the generic handler with an except OutOfPagesError as e: and call
logger.warning(...) there, allowing any other exceptions (ValueError,
AssertionError, etc.) to propagate and fail fast; reference the
mgr.impl.adjust() invocation and the logger.warning call when making this change
and ensure OutOfPagesError is imported/available in the module.

---

Nitpick comments:
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 2681-2693: The description for the Field declaration of
enable_kv_pool_rebalance is inconsistent with its status="prototype": update the
text to match the status by replacing the "Beta:" prefix with "Prototype:" (or
remove the stability prefix entirely) so the description aligns with the
Field(status="prototype") metadata; locate the enable_kv_pool_rebalance Field in
llm_args.py and edit its description string accordingly to reflect "Prototype:"
(or omit the prefix) and keep the rest of the wording intact.

In `@tests/unittest/_torch/executor/test_kv_pool_rebalance.py`:
- Around line 180-192: The test test_adjust_failure_does_not_skip_resume
currently includes an unused caplog parameter; remove caplog from the function
signature so the test becomes def test_adjust_failure_does_not_skip_resume(self,
monkeypatch): and update any callers or decorators if present; keep the rest of
the body intact (references to _make_request, _make_executor, exe, monkeypatch,
and assertions on exe.kv_cache_manager) and ensure imports/mocks are unchanged
so PyExecutor._maybe_rebalance_kv_pools behavior is still validated.

In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 2423-2458: The PR needs an explicit statement that QA list updates
are unnecessary because the change is limited to unit tests; update the PR
description (or commit message) to say something like "QA list update
unnecessary: changes are unittest-only
(tests/unittest/kv_cache_manager_v2_tests::TestSlotAllocatorShrink ->
test_shrink_underused_pool, test_shrink_touched_pool) and only add coverage for
storage shrink regression," so reviewers can easily see this follows the coding
guideline about unit-scope changes.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6604d89d-6a85-4540-9c7d-0acc873e6078

📥 Commits

Reviewing files that changed from the base of the PR and between 1f8312d and fc16558.

📒 Files selected for processing (11)

tensorrt_llm/_torch/pyexecutor/_util.py
tensorrt_llm/_torch/pyexecutor/py_executor.py
tensorrt_llm/_torch/pyexecutor/resource_manager.py
tensorrt_llm/llmapi/llm_args.py
tensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.py
tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py
tests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.py
tests/integration/test_lists/test-db/l0_a10.yml
tests/integration/test_lists/test-db/l0_h100.yml
tests/unittest/_torch/executor/test_kv_pool_rebalance.py
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

thorjohnsen · 2026-05-27T22:29:13Z

/bot run --disable-fail-fast

…sError The previous broad `except Exception` around `mgr.impl.adjust()` would swallow programmer errors (AssertionError, ValueError, etc.) and turn them into a warning log line, hiding real bugs. Only `OutOfPagesError` represents an expected runtime condition where the auto-tuner cannot fit the requested ratios; everything else should propagate and fail fast. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

…or-rebalance-hook Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/pyexecutor/py_executor.py # tests/integration/test_lists/test-db/l0_a10.yml # tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

Pre-commit picked up new docstring lint rules (D205, D301) and ruff-format style updates that landed on main since this PR opened. No behaviour change -- pure formatting / docstring reflow on three files this PR touched. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

thorjohnsen · 2026-05-28T15:57:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-28T16:03:41Z

PR_Github #50823 [ run ] triggered by Bot. Commit: 7887867 Link to invocation

nvpohanh · 2026-06-01T08:16:52Z

@lowsfer @yizhang-nv could you review this? Thanks

thorjohnsen · 2026-06-01T12:39:20Z

/bot run --disable-fail-fast

thorjohnsen · 2026-06-01T12:41:51Z

/bot run --disable-fail-fast

thorjohnsen · 2026-06-01T16:08:17Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-01T16:14:54Z

PR_Github #51376 [ run ] triggered by Bot. Commit: 1cb1bac Link to invocation

tensorrt-cicd · 2026-06-01T21:23:59Z

PR_Github #51376 [ run ] completed with state FAILURE. Commit: 1cb1bac
/LLM/main/L0_MergeRequest_PR pipeline #40788 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

thorjohnsen · 2026-06-01T23:45:43Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-01T23:51:25Z

PR_Github #51444 [ run ] triggered by Bot. Commit: 276dcb4 Link to invocation

tensorrt-cicd · 2026-06-02T07:44:00Z

PR_Github #51444 [ run ] completed with state SUCCESS. Commit: 276dcb4
/LLM/main/L0_MergeRequest_PR pipeline #40854 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

thorjohnsen · 2026-06-02T14:58:38Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-02T15:04:22Z

PR_Github #51620 [ run ] triggered by Bot. Commit: 276dcb4 Link to invocation

tensorrt-cicd · 2026-06-02T19:27:13Z

PR_Github #51620 [ run ] completed with state SUCCESS. Commit: 276dcb4
/LLM/main/L0_MergeRequest_PR pipeline #41005 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

Superjomn

LGTM on the llmapi changes.

nvpohanh · 2026-06-05T02:06:02Z

@lowsfer @yizhang-nv to review this. thanks

schetlur-nv

Approving for runtime devs

… (single GPU, aggregated for now) (NVIDIA#14578) Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com> Signed-off-by: NVFB <186336021+NVFB@users.noreply.github.com>

… (single GPU, aggregated for now) (NVIDIA#14578) Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

thorjohnsen self-assigned this May 26, 2026

thorjohnsen changed the title ~~[None][feat] KVCacheManagerV2: wire PyExecutor rebalance hook (single…~~ [TRTLLM-12714][feat] KVCacheManagerV2: wire PyExecutor rebalance hook (single GPU, aggregated, no overlap for now) May 26, 2026

thorjohnsen added 3 commits May 26, 2026 15:57

thorjohnsen added the KV-Cache Management kv-cache management for efficient LLM inference label May 26, 2026

thorjohnsen added 2 commits May 26, 2026 19:58

thorjohnsen changed the title ~~[TRTLLM-12714][feat] KVCacheManagerV2: wire PyExecutor rebalance hook (single GPU, aggregated, no overlap for now)~~ [TRTLLM-12714][feat] KVCacheManagerV2: wire PyExecutor rebalance hook (single GPU, aggregated for now) May 26, 2026

thorjohnsen added 3 commits May 27, 2026 21:51

[None][chore] KVCacheManagerV2: apply ruff-format to finish_shrink as…

fc16558

…sertion Pre-commit's ruff-format hook prefers the trailing-message-as-call style for assert with a long expression. Signed-off-by: Thor Johnsen <41591019+thorjohnsen@users.noreply.github.com>

thorjohnsen marked this pull request as ready for review May 27, 2026 22:12

thorjohnsen requested review from a team as code owners May 27, 2026 22:13

thorjohnsen requested review from JunyiXu-nv and byshiue May 27, 2026 22:13

coderabbitai Bot reviewed May 27, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

thorjohnsen added 3 commits May 27, 2026 22:32

jieli-matrix approved these changes May 28, 2026

View reviewed changes

Merge branch 'main' into feat/kvcm-v2-pyexecutor-rebalance-hook

7887867

nvpohanh requested review from lowsfer and yizhang-nv June 1, 2026 08:16

nvpohanh closed this Jun 1, 2026

nvpohanh reopened this Jun 1, 2026

Merge branch 'main' into feat/kvcm-v2-pyexecutor-rebalance-hook

23ec3c4

thorjohnsen added 2 commits June 1, 2026 09:51

Merge branch 'main' into feat/kvcm-v2-pyexecutor-rebalance-hook

8134b89

Merge branch 'main' into feat/kvcm-v2-pyexecutor-rebalance-hook

1cb1bac

Merge branch 'main' into feat/kvcm-v2-pyexecutor-rebalance-hook

276dcb4

Superjomn approved these changes Jun 3, 2026

View reviewed changes

lowsfer approved these changes Jun 5, 2026

View reviewed changes

schetlur-nv approved these changes Jun 5, 2026

View reviewed changes

thorjohnsen merged commit 86f9602 into NVIDIA:main Jun 5, 2026
7 checks passed

thorjohnsen deleted the feat/kvcm-v2-pyexecutor-rebalance-hook branch June 5, 2026 18:36

thorjohnsen mentioned this pull request Jun 22, 2026

[TRTLLM-12714][feat] KV pool rebalance: gate for multi-GPU and coordinate attention-DP #15518

Open

thorjohnsen mentioned this pull request Jul 9, 2026

[TRTLLM-12714][fix] Free CUDA-graph padding dummy KV caches before pool rebalance adjust() #16157

Open

Uh oh!

Conversation

thorjohnsen commented May 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Release Notes

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented May 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

❌ Failed checks (2 warnings)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

thorjohnsen commented May 27, 2026

Uh oh!

thorjohnsen commented May 28, 2026

Uh oh!

tensorrt-cicd commented May 28, 2026

Uh oh!

nvpohanh commented Jun 1, 2026

Uh oh!

thorjohnsen commented Jun 1, 2026

Uh oh!

thorjohnsen commented Jun 1, 2026

Uh oh!

thorjohnsen commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

thorjohnsen commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 1, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

thorjohnsen commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

Superjomn left a comment

Choose a reason for hiding this comment

Uh oh!

nvpohanh commented Jun 5, 2026

Uh oh!

schetlur-nv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

thorjohnsen commented May 26, 2026 •

edited

Loading

coderabbitai Bot commented May 27, 2026 •

edited

Loading