[https://nvbugs/6266370][fix] Fix MAX_UTILIZATION reuse token budget on main

[brb-nv]

Description

We currently use a very conservative estimate of reuse tokens for microbatch scheduling. This bottlenecks prefill significantly resulting in smaller batch sizes than expected. This MR fixes that.

Test Coverage

TEST(KVCacheManagerReuseAccountingTest, NeededBlocksOneStepCreditsFreeCachedReuseInTokenBudget)

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.

Summary by CodeRabbit

• Refactor

  • Updated block-reuse token budgeting calculations in context processing, resulting in more accurate token estimates from cached blocks and improved resource utilization

• Tests

  • Added test coverage for free-cached block reuse in batch admission scenarios

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 2b8cfe44-6250-4d80-8ddb-3fefebfdb709

📥 Commits

Reviewing files that changed from the base of the PR and between 48d2b89 and d73fb44.

📒 Files selected for processing (2)

• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
• cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp

---

📝 Walkthrough

Walkthrough

getNeededBlocksOneStep now separates block-capacity and token-budget accounting for cached-prefix reuse. When blocks become free but remain cached, they can now reduce token budget (via reusableBlocksAll) while only allocated blocks reduce block capacity budget (via reusableBlocksAllocated). A test validates this dual-path accounting improves FCFS admission.

Changes

KV Cache Prefix Reuse Token Budgeting

Layer / File(s)

Summary

Cached-prefix reuse accounting with separate block/token budget paths cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp

getNeededBlocksOneStep now computes maxRecoverableSharedBlocks and applies separate caps: reusableBlocksAllocated reduces numRequiredBlocks (block budget), while reusableBlocksAll sets estimatedReusableTokens (token budget). This allows free-but-cached blocks to improve token accounting without inflating block-capacity estimates. New test confirms cached reuse is credited and FCFS micro-batch admission admits more requests.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers

• JunyiXu-nv
• nvpohanh

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies a fix to the MAX_UTILIZATION reuse token budget using a specific NVBugs reference, directly corresponding to the main change in the PR.
Description check	✅ Passed	The PR description includes a clear explanation of the issue (conservative reuse token estimates bottlenecking prefill), the solution, test coverage details, and a completed checklist following the template structure.
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.}

[brb-nv]

/bot run --disable-fail-fast

[SimengLiu-nv]

LGTM.

[tensorrt-cicd]

PR_Github #53145 [ run ] triggered by Bot. Commit: d73fb44 Link to invocation

[tensorrt-cicd]

PR_Github #53145 [ run ] completed with state FAILURE. Commit: d73fb44
/LLM/main/L0_MergeRequest_PR pipeline #42348 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

[brb-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #53241 [ run ] triggered by Bot. Commit: d73fb44 Link to invocation

[tensorrt-cicd]

PR_Github #53241 [ run ] completed with state SUCCESS. Commit: d73fb44
/LLM/main/L0_MergeRequest_PR pipeline #42436 completed with status: 'SUCCESS'

CI Report

Link to invocation