[None][feat] Use max_gpu_total_bytes to control v2's capacity

[jiaganc]

Summary by CodeRabbit

• Refactor
  • Improved KV cache memory capacity handling with clearer logic for different cache manager configurations.
  • Refined GPU memory quota calculation to better prioritize explicit GPU memory limits.
  • Enhanced logging messages for memory management operations, providing clearer information about KV cache sizing and memory constraints.

Description

This can avoid the misleading max_tokens, since the estimated bytes per token are not accurate be accurate when there are swa layers. In this case max_tokens will confuse the user.

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR refactors memory quota and cache estimation logic by introducing a ceiling division utility function for consistent computations, revising quota calculation to prioritize explicit GPU memory limits, and adding conditional handling for different KVCacheManager variants with improved logging clarity.

Changes

Cohort / File(s)	Summary
Ceiling Division Utility tensorrt_llm/_torch/pyexecutor/_util.py	Added ceil_div() helper function and replaced multiple inline ceiling division operations throughout token cache block estimation and CUDA graph warmup calculations. Enhanced logging descriptions for memory constraints and KV cache capacity tracking.
Quota Calculation Refactor tensorrt_llm/_torch/pyexecutor/resource_manager.py	Restructured quota resolution logic to prioritize explicit max_gpu_total_bytes when provided, conditionally compute and clamp quota from max_tokens, and removed prior warning/clamping branches in favor of centralized logging flow.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 42.86% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	The PR description is incomplete and vague. While a title template is present, the Description section only contains a brief, unclear sentence about avoiding 'misleading max_tokens' due to 'swa layers', and Test Coverage section is empty.	Provide a clear explanation of the problem being solved, why the changes are necessary, which specific scenarios or bugs are addressed, and list at least one test case that validates the changes. Update the title with proper format [ticket][type] Summary.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main feature: using max_gpu_total_bytes to control v2's capacity, which matches the primary change in the changeset.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

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

[coderabbitai]

Actionable comments posted: 1

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/pyexecutor/_util.py (1)

492-503: ⚠️ Potential issue | 🟠 Major

Recompute legacy max_tokens after final memory clamp.

Line 492 derives max_tokens before Line 498 may further reduce kv_cache_max_memory via max_gpu_total_bytes.
For legacy KVCacheManager (see tensorrt_llm/_torch/pyexecutor/resource_manager.py, Line 883+), capacity is driven by max_tokens, so this ordering can over-allocate relative to the final memory cap.

Proposed fix

-            # For KvCacheManager, its logic still relies on max_tokens to control capacity
-            self._kv_cache_config.max_tokens = int(
-                kv_cache_max_memory // self._get_kv_size_per_token())
+            # For KVCacheManager, assign max_tokens after all memory clamps.

         # ---------------------------handle max_gpu_total_bytes---------------------------------
         # if user provided max_gpu_total_bytes, set max memory from max_gpu_total_bytes
         if self._kv_cache_config.max_gpu_total_bytes > 0:
             kv_cache_max_memory = min(kv_cache_max_memory,
                                       self._kv_cache_config.max_gpu_total_bytes)
             logger.info(
                 f"max_gpu_total_bytes={self._kv_cache_config.max_gpu_total_bytes / (GB):.2f} GiB is provided. New max memory is {kv_cache_max_memory / (GB):.2f} GiB"
             )
+
+        if not issubclass(self._kv_cache_manager_cls, KVCacheManagerV2):
+            self._kv_cache_config.max_tokens = int(
+                kv_cache_max_memory // self._get_kv_size_per_token())

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 492 - 503, Recompute
legacy max_tokens after any final memory clamps: after applying the
max_gpu_total_bytes clamp to kv_cache_max_memory (the block that checks
self._kv_cache_config.max_gpu_total_bytes), recalculate and set
self._kv_cache_config.max_tokens using the final kv_cache_max_memory via
self._get_kv_size_per_token(); this ensures KVCacheManager capacity (which is
driven by self._kv_cache_config.max_tokens) reflects the post-clamp memory limit
rather than the earlier pre-clamp value.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 168-170: The profiling log f-string in _util.py currently emits a
doubled comma between the total GPU memory and the temporary KV cache message;
update the f-string that builds the log (the expression referencing fraction,
kv_size_per_token, total_gpu_memory, allocated_bytes and GB) to remove the
extraneous comma so the output reads "...device total memory X GiB, temporary kv
cache memory during profiling Y GiB" (i.e., eliminate the extra ", " before
"temporary kv cache memory during profiling").

---

Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 492-503: Recompute legacy max_tokens after any final memory
clamps: after applying the max_gpu_total_bytes clamp to kv_cache_max_memory (the
block that checks self._kv_cache_config.max_gpu_total_bytes), recalculate and
set self._kv_cache_config.max_tokens using the final kv_cache_max_memory via
self._get_kv_size_per_token(); this ensures KVCacheManager capacity (which is
driven by self._kv_cache_config.max_tokens) reflects the post-clamp memory limit
rather than the earlier pre-clamp value.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 89f1e7aa-f408-4201-a3bb-c35b8b5d9089

📥 Commits

Reviewing files that changed from the base of the PR and between d3536f1 and ef48f21.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/_util.py
• tensorrt_llm/_torch/pyexecutor/resource_manager.py

[jiaganc]

/bot run

[tensorrt-cicd]

PR_Github #37691 [ run ] triggered by Bot. Commit: adc2975 Link to invocation

[jiaganc]

/bot run

[tensorrt-cicd]

PR_Github #37704 [ run ] triggered by Bot. Commit: 76b77ed Link to invocation

[tensorrt-cicd]

PR_Github #37704 [ run ] completed with state SUCCESS. Commit: 76b77ed
/LLM/main/L0_MergeRequest_PR pipeline #29183 completed with status: 'FAILURE'

⚠️ 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

[jiaganc]

/bot run --disable-fast-fail

[tensorrt-cicd]

PR_Github #37855 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: --disable-fast-fail

Link to invocation

[jiaganc]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37951 [ run ] triggered by Bot. Commit: 76b77ed Link to invocation

[tensorrt-cicd]

PR_Github #37951 [ run ] completed with state SUCCESS. Commit: 76b77ed
/LLM/main/L0_MergeRequest_PR pipeline #29392 completed with status: 'SUCCESS'

Link to invocation

[thorjohnsen]

LGTM

[lfr-0531]

LGTM~

[jiaganc]

The precommit failure is not related to this PR and should be fixed by #11920, so merge it.