[TRTLLM-12399][fix] Fix KV cache adaptive ratio sampling

[lowsfer]

Summary by CodeRabbit

• Bug Fixes
  • Improved KV cache management stability by refining cache-level adjustment triggers based on sampling metrics rather than closure events.
  • Added safeguards for edge cases: enforces minimum block requirements per life cycle and handles scenarios where capacity is below history length.

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)

• 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.

[lowsfer]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR refactors KVCacheManager's adjustment trigger logic from counting closed KV caches to counting sampled KV caches. _KVCache.close() conditionally updates sampling statistics only when capacity exceeds zero. StorageManager.ratio_from_length validates capacity constraints and enforces minimum block allocations per lifecycle.

Changes

KV Cache Sampling-Based Adjustment

Layer / File(s)	Summary
KVCacheManager Field Replacement tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py	Internal tracking fields transition from closed-KV-cache/closed-request counters to sampled-KV-cache and last-update timestamp counters.
KVCacheManager Initialization tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py	__init__ initializes new sampled-KV-cache counters (_num_sampled_kv_caches, _last_update_num_sampled_kv_caches) and adjustment timing fields.
Adjustment Gate Logic tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py	need_adjustment property and _try_update_target_ratios() method now gate on sampled-KV-cache deltas and thresholds instead of closed-cache counts.
KVCache Closure Guard tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py	_KVCache.close() guards manager sampling updates with if self.capacity > 0, skipping sampling/ratio updates for zero-capacity caches while preserving shutdown sequence.
Storage Manager Validation tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py	ratio_from_length validates capacity >= history_length and enforces minimum 1 block per non-SSM lifecycle via max(num_blocks - stale, 1).

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.00% 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 pull request description is incomplete, containing only the template structure without substantive content in required sections.	Fill in the 'Description' section explaining the issue and solution, 'Test Coverage' section listing relevant tests, and provide more context about the KV cache adaptive ratio sampling fix.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: fixing KV cache adaptive ratio sampling, with proper JIRA ticket reference and fix type designation.
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

Tip

💬 Introducing Slack Agent: The best way for teams to turn conversations into code.

Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.

• Generate code and open pull requests
• Plan features and break down work
• Investigate incidents and troubleshoot customer tickets together
• Automate recurring tasks and respond to alerts with triggers
• Summarize progress and report instantly

Built for teams:

• Shared memory across your entire org—no repeating context
• Per-thread sandboxes to safely plan and execute work
• Governance built-in—scoped access, auditability, and budget controls

One agent for your entire SDLC. Right inside Slack.

👉 Get started

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 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 730-732: The warnings.warn call in _storage_manager.py (the branch
checking if capacity < history_length) needs an explicit stacklevel to satisfy
Ruff B028; update the call to warnings.warn to pass stacklevel=2 (e.g.,
warnings.warn("Bad sampling for capacity and history_length", stacklevel=2)) so
the warning points to the caller rather than the warning site, leaving the
surrounding logic that sets capacity = history_length unchanged.

🪄 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: 5add0062-08d1-41d3-9f41-e37c3ad493ec

📥 Commits

Reviewing files that changed from the base of the PR and between 47d7ecc and c1a42cb.

📒 Files selected for processing (3)

• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py

[tensorrt-cicd]

PR_Github #47215 [ run ] triggered by Bot. Commit: c1a42cb Link to invocation

[tensorrt-cicd]

PR_Github #47215 [ run ] completed with state SUCCESS. Commit: c1a42cb
/LLM/main/L0_MergeRequest_PR pipeline #37171 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

[jiaganc]

Approved with a nit comment.

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #47376 [ run ] triggered by Bot. Commit: c1a42cb Link to invocation

[tensorrt-cicd]

PR_Github #47376 [ run ] completed with state SUCCESS. Commit: c1a42cb
/LLM/main/L0_MergeRequest_PR pipeline #37307 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

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #47509 [ run ] triggered by Bot. Commit: c1a42cb Link to invocation

[tensorrt-cicd]

PR_Github #47509 [ run ] completed with state SUCCESS. Commit: c1a42cb
/LLM/main/L0_MergeRequest_PR pipeline #37427 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

Link to invocation

[lowsfer]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #47628 [ run ] triggered by Bot. Commit: dede1ad Link to invocation

[tensorrt-cicd]

PR_Github #47628 [ run ] completed with state SUCCESS. Commit: dede1ad
/LLM/main/L0_MergeRequest_PR pipeline #37533 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

[lowsfer]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #47824 [ run ] triggered by Bot. Commit: dede1ad Link to invocation

[tensorrt-cicd]

PR_Github #47824 [ run ] completed with state SUCCESS. Commit: dede1ad
/LLM/main/L0_MergeRequest_PR pipeline #37709 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

[lowsfer]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #47917 [ run ] triggered by Bot. Commit: dede1ad Link to invocation

[tensorrt-cicd]

PR_Github #47917 [ run ] completed with state SUCCESS. Commit: dede1ad
/LLM/main/L0_MergeRequest_PR pipeline #37764 completed with status: 'SUCCESS'

CI Report

Link to invocation