[https://nvbugs/6438658][fix] Account for resume utilization in KV constraints#16484
Merged
Merged
Conversation
Collaborator
Author
|
Hi, @lowsfer could you please review this PR? |
Collaborator
Author
|
/bot run |
Contributor
📝 WalkthroughWalkthroughConstraint-derived KV cache slot floors now reserve headroom based on ChangesKV cache resume headroom
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant TestInitRatioConfig
participant KVCacheManager
participant StorageManager
participant KVCache
TestInitRatioConfig->>KVCacheManager: create constrained KV caches
KVCacheManager->>StorageManager: pass max_util_for_resume
StorageManager-->>KVCacheManager: reserve scaled slot capacity
KVCacheManager->>KVCache: resume cache
KVCache-->>TestInitRatioConfig: resume succeeds
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)
2446-2470: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider adding a negative/boundary companion test.
The new test verifies the success path once headroom scaling is applied. As per path instructions for
tests/**, coverage should be assessed for sufficiency: since without the scaling,min_slots=32would put utilization at exactly1.0 > 0.95, this test would likely already catch a regression that removes scaling entirely — coverage is sufficient for that case. Optionally, add a companion test assertingresume()returnsFalseatmax_util_for_resume=1.0(or with an even tighter constraint) to make the headroom mechanism's necessity explicit and pin the boundary behavior.🤖 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 2446 - 2470, Add a boundary companion test for the resume headroom behavior around KVCacheManager and max_util_for_resume, configuring utilization at the threshold (for example, 1.0 or a tighter constraint) and asserting kv_cache.resume(stream) returns False. Reuse the existing setup symbols such as BatchDesc, KVCacheManager, and _make_config, and ensure resources are closed and the manager is shut down.Source: Path instructions
🤖 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/runtime/kv_cache_manager_v2/_storage_manager.py`:
- Around line 943-947: Validate max_util_for_resume before the scaling
calculation in the constraints flow around _compute_slots_for_batch, enforcing 0
< max_util_for_resume <= 1.0. Update KVCacheManagerConfig.__post_init__ if
appropriate so invalid values are rejected before division, while preserving the
existing max_slots calculation for valid values.
---
Nitpick comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 2446-2470: Add a boundary companion test for the resume headroom
behavior around KVCacheManager and max_util_for_resume, configuring utilization
at the threshold (for example, 1.0 or a tighter constraint) and asserting
kv_cache.resume(stream) returns False. Reuse the existing setup symbols such as
BatchDesc, KVCacheManager, and _make_config, and ensure resources are closed and
the manager is shut down.
🪄 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: 16c0f58a-2645-44d5-9bfe-f75034311daa
📒 Files selected for processing (3)
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
Collaborator
|
PR_Github #59674 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59674 [ run ] completed with state
|
lowsfer
approved these changes
Jul 16, 2026
jiaganc
force-pushed
the
codex/kv-constraints-resume-util
branch
from
July 17, 2026 03:05
f240f4f to
4cebd09
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59854 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59854 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59882 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59882 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59987 [ run ] triggered by Bot. Commit: |
jiaganc
enabled auto-merge (squash)
July 17, 2026 15:58
Collaborator
|
PR_Github #59987 [ run ] completed with state
|
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
jiaganc
force-pushed
the
codex/kv-constraints-resume-util
branch
from
July 20, 2026 03:19
4cebd09 to
ff4f191
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60247 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60247 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60314 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60314 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #60374 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #60374 [ run ] completed with state |
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
Port the Python-side fix from PR NVIDIA#16484 (nvbugs/6438658) to the C++ StorageManager so both backends behave identically: thread max_util_for_resume from KVCacheManagerConfig through the StorageManager constructor, and scale each constraint-derived min-slot floor by 1/max_util_for_resume in computeMinSlotsFromConstraints so a declared batch stays below the utilization gate checked by KvCache::resume. This mirrors main's behavior exactly: an explicit initial_pool_ratio still overrides constraints for min-slots (the gating is kept), and the scale factor is applied unconditionally (no clamping). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
REVIEW NOTE: this is the cpp-branch behavior delta on top of the previous commit (which mirrors main's PR NVIDIA#16484). It is isolated so it can be reviewed or dropped independently. If we decide to match main exactly, revert just this commit. Diverges from main's resume-utilization fix in two ways, on both the Python and C++ backends: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (main drops them): a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. Follows Lanyu's C++ commit 15a90e1 (which cited PR NVIDIA#16269). - Clamp max_util_for_resume to the binding (0, 1) range before scaling, so a degenerate 0 can't divide-by-zero and a value >= 1 can't shrink floors. Renames test_initial_pool_ratio_overrides_typical_step_and_constraints to test_constraint_floor_overrides_infeasible_initial_pool_ratio to assert the new behavior (main's test_constraint_reserves_resume_headroom is kept as-is). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
Port the Python-side fix from PR NVIDIA#16484 (nvbugs/6438658) to the C++ StorageManager so both backends behave identically: thread max_util_for_resume from KVCacheManagerConfig through the StorageManager constructor, and scale each constraint-derived min-slot floor by 1/max_util_for_resume in computeMinSlotsFromConstraints so a declared batch stays below the utilization gate checked by KvCache::resume. This mirrors main's behavior exactly: an explicit initial_pool_ratio still overrides constraints for min-slots (the gating is kept), and the scale factor is applied unconditionally (no clamping). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
REVIEW NOTE: this is the cpp-branch behavior delta on top of the previous commit (which mirrors main's PR NVIDIA#16484). It is isolated so it can be reviewed or dropped independently. If we decide to match main exactly, revert just this commit. Diverges from main's resume-utilization fix in two ways, on both the Python and C++ backends: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (main drops them): a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. Follows Lanyu's C++ commit 15a90e1 (which cited PR NVIDIA#16269). - Clamp max_util_for_resume to the binding (0, 1) range before scaling, so a degenerate 0 can't divide-by-zero and a value >= 1 can't shrink floors. Renames test_initial_pool_ratio_overrides_typical_step_and_constraints to test_constraint_floor_overrides_infeasible_initial_pool_ratio to assert the new behavior (main's test_constraint_reserves_resume_headroom is kept as-is). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
Port the Python-side fix from PR NVIDIA#16484 (nvbugs/6438658) to the C++ StorageManager so both backends behave identically: thread max_util_for_resume from KVCacheManagerConfig through the StorageManager constructor, and scale each constraint-derived min-slot floor by 1/max_util_for_resume in computeMinSlotsFromConstraints so a declared batch stays below the utilization gate checked by KvCache::resume. This mirrors main's behavior exactly: an explicit initial_pool_ratio still overrides constraints for min-slots (the gating is kept), and the scale factor is applied unconditionally (no clamping). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
REVIEW NOTE: this is the cpp-branch behavior delta on top of the previous commit (which mirrors main's PR NVIDIA#16484). It is isolated so it can be reviewed or dropped independently. If we decide to match main exactly, revert just this commit. Diverges from main's resume-utilization fix in two ways, on both the Python and C++ backends: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (main drops them): a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. Follows Lanyu's C++ commit 15a90e1 (which cited PR NVIDIA#16269). - Clamp max_util_for_resume to the binding (0, 1) range before scaling, so a degenerate 0 can't divide-by-zero and a value >= 1 can't shrink floors. Renames test_initial_pool_ratio_overrides_typical_step_and_constraints to test_constraint_floor_overrides_infeasible_initial_pool_ratio to assert the new behavior (main's test_constraint_reserves_resume_headroom is kept as-is). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
Port the Python-side fix from PR NVIDIA#16484 (nvbugs/6438658) to the C++ StorageManager so both backends behave identically: thread max_util_for_resume from KVCacheManagerConfig through the StorageManager constructor, and scale each constraint-derived min-slot floor by 1/max_util_for_resume in computeMinSlotsFromConstraints so a declared batch stays below the utilization gate checked by KvCache::resume. This mirrors main's behavior exactly: an explicit initial_pool_ratio still overrides constraints for min-slots (the gating is kept), and the scale factor is applied unconditionally (no clamping). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
REVIEW NOTE: this is the cpp-branch behavior delta on top of the previous commit (which mirrors main's PR NVIDIA#16484). It is isolated so it can be reviewed or dropped independently. If we decide to match main exactly, revert just this commit. Diverges from main's resume-utilization fix in two ways, on both the Python and C++ backends: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (main drops them): a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. Follows Lanyu's C++ commit 15a90e1 (which cited PR NVIDIA#16269). - Clamp max_util_for_resume to the binding (0, 1) range before scaling, so a degenerate 0 can't divide-by-zero and a value >= 1 can't shrink floors. Renames test_initial_pool_ratio_overrides_typical_step_and_constraints to test_constraint_floor_overrides_infeasible_initial_pool_ratio to assert the new behavior (main's test_constraint_reserves_resume_headroom is kept as-is). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
Port the Python-side fix from PR NVIDIA#16484 (nvbugs/6438658) to the C++ StorageManager so both backends behave identically: thread max_util_for_resume from KVCacheManagerConfig through the StorageManager constructor, and scale each constraint-derived min-slot floor by 1/max_util_for_resume in computeMinSlotsFromConstraints so a declared batch stays below the utilization gate checked by KvCache::resume. This mirrors main's behavior exactly: an explicit initial_pool_ratio still overrides constraints for min-slots (the gating is kept), and the scale factor is applied unconditionally (no clamping). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
REVIEW NOTE: this is the cpp-branch behavior delta on top of the previous commit (which mirrors main's PR NVIDIA#16484). It is isolated so it can be reviewed or dropped independently. If we decide to match main exactly, revert just this commit. Diverges from main's resume-utilization fix in two ways, on both the Python and C++ backends: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (main drops them): a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. Follows Lanyu's C++ commit 15a90e1 (which cited PR NVIDIA#16269). - Clamp max_util_for_resume to the binding (0, 1) range before scaling, so a degenerate 0 can't divide-by-zero and a value >= 1 can't shrink floors. Renames test_initial_pool_ratio_overrides_typical_step_and_constraints to test_constraint_floor_overrides_infeasible_initial_pool_ratio to assert the new behavior (main's test_constraint_reserves_resume_headroom is kept as-is). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
Port the Python-side fix from PR NVIDIA#16484 (nvbugs/6438658) to the C++ StorageManager so both backends behave identically: thread max_util_for_resume from KVCacheManagerConfig through the StorageManager constructor, and scale each constraint-derived min-slot floor by 1/max_util_for_resume in computeMinSlotsFromConstraints so a declared batch stays below the utilization gate checked by KvCache::resume. This mirrors main's behavior exactly: an explicit initial_pool_ratio still overrides constraints for min-slots (the gating is kept), and the scale factor is applied unconditionally (no clamping). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
added a commit
to lowsfer/TensorRT-LLM
that referenced
this pull request
Jul 21, 2026
REVIEW NOTE: this is the cpp-branch behavior delta on top of the previous commit (which mirrors main's PR NVIDIA#16484). It is isolated so it can be reviewed or dropped independently. If we decide to match main exactly, revert just this commit. Diverges from main's resume-utilization fix in two ways, on both the Python and C++ backends: - Constraints stay feasibility floors even under an explicit initial_pool_ratio (main drops them): a declared batch that needs more than its target pool-group share clamps that share up instead of starving during warmup. Follows Lanyu's C++ commit 15a90e1 (which cited PR NVIDIA#16269). - Clamp max_util_for_resume to the binding (0, 1) range before scaling, so a degenerate 0 can't divide-by-zero and a value >= 1 can't shrink floors. Renames test_initial_pool_ratio_overrides_typical_step_and_constraints to test_constraint_floor_overrides_infeasible_initial_pool_ratio to assert the new behavior (main's test_constraint_reserves_resume_headroom is kept as-is). Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
1 task
eopXD
added a commit
to eopXD/TensorRT-LLM
that referenced
this pull request
Jul 22, 2026
…verage Add a Python-only per-iteration suspend/resume (preemption) observability counter to KVCacheManagerV2 -- iterSuspendedRequests / iterResumedRequests, surfaced via get_stats() -- together with the tests that exercise it. This coverage is genuinely new on main: no existing test can assert on the counter because the counter did not exist, and there is no end-to-end suspend/resume integration cell for GPT-OSS on the V2 manager. (Overlapping bring-up cells -- eagle3-1gpu, chunked-prefill -- already exist on main and are deliberately NOT re-added here.) Stat (mirrors the NVIDIA#15633 offload/onboard path; no C++/nanobind change): incremented at _KVCache.suspend() and on a successful resume(), gated on _should_record_stats(), drained once per iteration-stats fetch into KVCacheV2IterationStatsReport and emitted as top-level iteration keys. Coverage added: - tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py: the counters increment on suspend+resume and surface via get_iteration_stats() (model-free, deterministic). - tests/unittest/executor/test_stats_serializer.py: the serializer emits the top-level counters (not nested inside the per-pool-group breakdown). - tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestGPTOSS:: test_w4_1gpu_suspend_resume -- V2-only ACTIVE<->SUSPENDED->resume round trip with page-index reconnect; gates on iterSuspendedRequests>0 and iterResumedRequests>0, no crash/deadlock, and a >=3-token deterministic prefix vs the uncontended reference. Verification: the suspend/resume unit + serializer tests and I-10 passed on the rc21 build (main ef6ebc2 + NVIDIA#15633): I-10 suspended=31/resumed=35, offload/onboard=0, prefixes 17/11/19/8; api_stability 64/64. The unit tests are model-free and base-insensitive. I-10's workload sizing was tuned on rc21 and has NOT yet been re-run on this later main; the resume-utilization change (NVIDIA#16484) can shift when suspension fires, so retune max_tokens if the gate-(2) counters come back 0. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
eopXD
added a commit
to eopXD/TensorRT-LLM
that referenced
this pull request
Jul 22, 2026
…verage Add a Python-only per-iteration suspend/resume (preemption) observability counter to KVCacheManagerV2 -- iterSuspendedRequests / iterResumedRequests, surfaced via get_stats() -- together with the tests that exercise it. This coverage is genuinely new on main: no existing test can assert on the counter because the counter did not exist, and there is no end-to-end suspend/resume integration cell for GPT-OSS on the V2 manager. (Overlapping bring-up cells -- eagle3-1gpu, chunked-prefill -- already exist on main and are deliberately NOT re-added here.) Stat (mirrors the NVIDIA#15633 offload/onboard path; no C++/nanobind change): incremented at _KVCache.suspend() and on a successful resume(), gated on _should_record_stats(), drained once per iteration-stats fetch into KVCacheV2IterationStatsReport and emitted as top-level iteration keys. Coverage added: - tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py: the counters increment on suspend+resume and surface via get_iteration_stats() (model-free, deterministic). - tests/unittest/executor/test_stats_serializer.py: the serializer emits the top-level counters (not nested inside the per-pool-group breakdown). - tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestGPTOSS:: test_w4_1gpu_suspend_resume -- V2-only ACTIVE<->SUSPENDED->resume round trip with page-index reconnect; gates on iterSuspendedRequests>0 and iterResumedRequests>0, no crash/deadlock, and a >=3-token deterministic prefix vs the uncontended reference. Registered in tests/integration/test_lists/test-db/l0_b200.yml (1-GPU B200, pre-merge). Verification: the suspend/resume unit + serializer tests and I-10 passed on the rc21 build (main ef6ebc2 + NVIDIA#15633): I-10 suspended=31/resumed=35, offload/onboard=0, prefixes 17/11/19/8; api_stability 64/64. The unit tests are model-free and base-insensitive. I-10's workload sizing was tuned on rc21 and has NOT yet been re-run on this later main; the resume-utilization change (NVIDIA#16484) can shift when suspension fires, so retune max_tokens if the gate-(2) counters come back 0. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Bug Fixes
Tests
Description
Constraint-derived pool floors currently reserve the exact slot count required by each constrained batch. However,
_KVCache.resume()rejects a resume when projected utilization exceedsmax_util_for_resume, so a supposedly supported constraint can still fail before the pool is full.Pass
max_util_for_resumeintoStorageManagerand scale constraint-derived minimum slot counts by the inverse utilization threshold. This preserves sufficient headroom for constrained batches to resume. Explicit user-provided pool ratios continue to take precedence over inferred constraints.Test Coverage
test_constraint_reserves_resume_headroomto exercise a complete constrained batch atmax_util_for_resume=0.95.python3 -m py_compilepassed for all changed files.PR Checklist
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.