[https://nvbugs/5570599][fix] Set KVCache free_gpu_memory_fraction fo… #8780

liji-nv · 2025-10-30T07:23:21Z

…r estimation

Summary by CodeRabbit

Bug Fixes
- Fixed GPU memory allocation settings not being properly respected during KV-cache operations. User-provided memory fractions are now preserved throughout the estimation and configuration processes, with sensible defaults applied only when preferences aren't specified.

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

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

liji-nv · 2025-10-30T07:24:15Z

/bot run

liji-nv · 2025-10-30T07:25:32Z

/bot run

coderabbitai · 2025-10-30T07:28:36Z

📝 Walkthrough

Walkthrough

Modified memory fraction handling in KvCacheCreator to track user-provided values separately from configuration values, ensuring user preferences are preserved during KV-cache estimation and configuration while applying defaults only when user input is absent.

Changes

Cohort / File(s)	Change Summary
KvCacheCreator Memory Fraction Tracking `tensorrt_llm/_torch/pyexecutor/_util.py`	Added `_user_set_free_gpu_memory_fraction` instance variable to store user-provided memory fractions; modified `_get_free_gpu_memory_fraction()` to read from user-set value with 0.9 default; `try_prepare_estimation()` now logs and sets default when user value is absent; `configure_kv_cache_capacity()` restores user-provided fraction before computation to preserve user preference.

Sequence Diagram

sequenceDiagram
    participant User
    participant KvCacheCreator
    participant kv_cache_config

    User->>KvCacheCreator: Initialize with optional free_gpu_memory_fraction
    KvCacheCreator->>KvCacheCreator: Store in _user_set_free_gpu_memory_fraction

    rect rgb(200, 220, 255)
    Note over KvCacheCreator: Estimation Phase
    User->>KvCacheCreator: try_prepare_estimation()
    alt User fraction provided
        KvCacheCreator->>KvCacheCreator: Use _user_set_free_gpu_memory_fraction
    else User fraction not provided
        KvCacheCreator->>KvCacheCreator: Set default (0.9)
        KvCacheCreator->>kv_cache_config: Set free_gpu_memory_fraction = 0.9
        KvCacheCreator->>KvCacheCreator: Log action
    end
    end

    rect rgb(200, 255, 220)
    Note over KvCacheCreator: Configuration Phase
    User->>KvCacheCreator: configure_kv_cache_capacity()
    KvCacheCreator->>kv_cache_config: Restore _user_set_free_gpu_memory_fraction
    KvCacheCreator->>KvCacheCreator: Compute free memory using restored value
    end

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Verify the initialization and storage of _user_set_free_gpu_memory_fraction is complete across all KvCacheCreator instantiation paths
Confirm the restoration logic in configure_kv_cache_capacity() correctly preserves user intent without unintended side effects
Validate the default value (0.9) application is only triggered when user input is explicitly absent
Check thread-safety considerations if KvCacheCreator is used in concurrent contexts

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description Check	⚠️ Warning	The pull request description is largely incomplete and does not meet the template requirements. While the PR title and checklist box are present, the critical "Description" and "Test Coverage" sections are entirely empty with only placeholder comments. The author has checked the PR Checklist box indicating completion, but the required explanation of the issue and solution, as well as the test coverage documentation, are missing. These sections are essential for understanding the rationale and validation strategy for the changes.
Docstring Coverage	⚠️ Warning	Docstring coverage is 40.00% which is insufficient. The required threshold is 80.00%.	You can run `@coderabbitai generate docstrings` to improve docstring coverage.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title Check	✅ Passed	The pull request title follows the correct format with an NVBugs ticket ID, a "[fix]" type indicator, and a clear summary that accurately describes the main change in the changeset. The title "Set KVCache free_gpu_memory_fraction for estimation" is specific and directly related to the core modification documented in the raw summary, which details changes to how free_gpu_memory_fraction is handled during KV-cache capacity estimation and configuration. The title concisely communicates the primary intent without unnecessary noise or vague language.

✨ Finishing touches

📝 Generate docstrings

🧪 Generate unit tests (beta)

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

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (2)

tensorrt_llm/_torch/pyexecutor/_util.py (2)
80-80: Consider adding a clarifying comment.

While the variable name is self-explanatory, adding a brief inline comment would help future maintainers understand why the user-provided value is being cached separately from the config value.

Example:
+        # Cache the user-provided memory fraction to preserve user intent during estimation
         self._user_set_free_gpu_memory_fraction = self._kv_cache_config.free_gpu_memory_fraction
259-260: Refine the comment for accuracy.

The comment states the value "may get modified by the estimation process", but the value is only modified during estimation if the user didn't provide it (lines 240-245). Consider making the comment more precise.

Apply this diff:
-        # Restore the user-set free_gpu_memory_fraction as it may get modified by the estimation process.
+        # Restore the user-set free_gpu_memory_fraction (which may have been set to a default during estimation if not originally provided).
         self._kv_cache_config.free_gpu_memory_fraction = self._user_set_free_gpu_memory_fraction

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0dac57f and 40d7a5f.

📒 Files selected for processing (1)

tensorrt_llm/_torch/pyexecutor/_util.py (4 hunks)

🧰 Additional context used

📓 Path-based instructions (3)

**/*.{h,hpp,hh,hxx,cpp,cxx,cc,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Use only spaces, no tabs; indent with 4 spaces.

Files:

tensorrt_llm/_torch/pyexecutor/_util.py

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

**/*.py: Python code must target Python 3.8+.
Indent Python code with 4 spaces; do not use tabs.
Maintain module namespace when importing; prefer 'from package.subpackage import foo' then 'foo.SomeClass()' instead of importing the class directly.
Python filenames should be snake_case (e.g., some_file.py).
Python classes use PascalCase names.
Functions and methods use snake_case names.
Local variables use snake_case; prefix 'k' for variables that start with a number (e.g., k_99th_percentile).
Global variables use upper SNAKE_CASE prefixed with 'G' (e.g., G_MY_GLOBAL).
Constants use upper SNAKE_CASE (e.g., MY_CONSTANT).
Avoid shadowing variables from an outer scope.
Initialize all externally visible members of a class in the constructor.
Prefer docstrings for interfaces that may be used outside a file; comments for in-function or file-local interfaces.
Use Google-style docstrings for classes and functions (Sphinx-parsable).
Document attributes and variables inline so they render under the class/function docstring.
Avoid reflection when a simpler, explicit approach suffices (e.g., avoid dict(**locals()) patterns).
In try/except, catch the most specific exceptions possible.
For duck-typing try/except, keep the try body minimal and use else for the main logic.

Files:

tensorrt_llm/_torch/pyexecutor/_util.py

**/*.{cpp,cxx,cc,h,hpp,hh,hxx,cu,cuh,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Prepend the NVIDIA Apache-2.0 copyright header with current year to the top of all source files (e.g., .cpp, .h, .cu, .py).

Files:

tensorrt_llm/_torch/pyexecutor/_util.py

🧠 Learnings (1)

📓 Common learnings

Learnt from: thorjohnsen
PR: NVIDIA/TensorRT-LLM#6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.248Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

GitHub Check: Pre-commit Check

🔇 Additional comments (1)

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

239-245: Good defensive programming practice.

The conditional setting of free_gpu_memory_fraction during estimation ensures that the KV cache manager has a valid value to work with while preserving user intent. The logging statement provides good visibility for debugging.

tensorrt_llm/_torch/pyexecutor/_util.py

tensorrt-cicd · 2025-10-30T07:31:08Z

PR_Github #23010 [ run ] triggered by Bot. Commit: 40d7a5f

tensorrt-cicd · 2025-10-30T08:32:12Z

PR_Github #23010 [ run ] completed with state SUCCESS. Commit: 40d7a5f
/LLM/release-1.1/L0_MergeRequest_PR pipeline #318 completed with status: 'FAILURE'

tensorrt_llm/_torch/pyexecutor/_util.py

liji-nv · 2025-10-31T08:29:12Z

/bot run

tensorrt-cicd · 2025-10-31T08:34:56Z

PR_Github #23167 [ run ] triggered by Bot. Commit: 9077242

tensorrt-cicd · 2025-10-31T09:10:17Z

PR_Github #23167 [ run ] completed with state FAILURE. Commit: 9077242
/LLM/release-1.1/L0_MergeRequest_PR pipeline #334 completed with status: 'FAILURE'

liji-nv · 2025-10-31T09:17:02Z

/bot run

tensorrt-cicd · 2025-10-31T09:23:37Z

PR_Github #23182 [ run ] triggered by Bot. Commit: 9077242

tensorrt-cicd · 2025-10-31T09:59:57Z

PR_Github #23182 [ run ] completed with state FAILURE. Commit: 9077242
/LLM/release-1.1/L0_MergeRequest_PR pipeline #340 completed with status: 'FAILURE'

HuiGao-NV

LGTM

liji-nv · 2025-11-03T05:39:13Z

/bot run

tensorrt-cicd · 2025-11-03T05:44:38Z

PR_Github #23340 [ run ] triggered by Bot. Commit: 4bd974d

tensorrt-cicd · 2025-11-03T06:39:45Z

PR_Github #23340 [ run ] completed with state SUCCESS. Commit: 4bd974d
/LLM/release-1.1/L0_MergeRequest_PR pipeline #376 completed with status: 'FAILURE'

liji-nv · 2025-11-04T08:45:24Z

/bot run

tensorrt-cicd · 2025-11-04T08:51:40Z

PR_Github #23490 [ run ] triggered by Bot. Commit: 66d2532

tensorrt-cicd · 2025-11-04T10:03:53Z

PR_Github #23490 [ run ] completed with state SUCCESS. Commit: 66d2532
/LLM/release-1.1/L0_MergeRequest_PR pipeline #407 completed with status: 'FAILURE'

liji-nv · 2025-11-04T11:42:17Z

/bot run

tensorrt-cicd · 2025-11-04T11:48:09Z

PR_Github #23507 [ run ] triggered by Bot. Commit: b678ff9

tensorrt-cicd · 2025-11-04T12:57:17Z

PR_Github #23507 [ run ] completed with state SUCCESS. Commit: b678ff9
/LLM/release-1.1/L0_MergeRequest_PR pipeline #408 completed with status: 'FAILURE'

liji-nv · 2025-11-05T02:47:05Z

/bot run

tensorrt_llm/_torch/pyexecutor/resource_manager.py

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv · 2025-11-06T03:44:07Z

/bot run

tensorrt-cicd · 2025-11-06T03:49:49Z

PR_Github #23701 [ run ] triggered by Bot. Commit: 8ba6b32

tensorrt-cicd · 2025-11-06T13:58:04Z

PR_Github #23701 [ run ] completed with state SUCCESS. Commit: 8ba6b32
/LLM/release-1.1/L0_MergeRequest_PR pipeline #441 completed with status: 'SUCCESS'

NVIDIA#8780) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

NVIDIA#8780) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com>

NVIDIA#8780) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Mike Iovine <miovine@nvidia.com>

…VIDIA#8780) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv requested a review from a team as a code owner October 30, 2025 07:23

coderabbitai bot reviewed Oct 30, 2025

View reviewed changes

tensorrt_llm/_torch/pyexecutor/_util.py Show resolved Hide resolved

HuiGao-NV reviewed Oct 30, 2025

View reviewed changes

tensorrt_llm/_torch/pyexecutor/_util.py Outdated Show resolved Hide resolved

tensorrt_llm/_torch/pyexecutor/_util.py Outdated Show resolved Hide resolved

liji-nv force-pushed the dev-liji-set-fraction branch from 40d7a5f to 9077242 Compare October 31, 2025 08:29

HuiGao-NV approved these changes Nov 3, 2025

View reviewed changes

liji-nv force-pushed the dev-liji-set-fraction branch from 9077242 to 4bd974d Compare November 3, 2025 05:39

liji-nv force-pushed the dev-liji-set-fraction branch from 4bd974d to 66d2532 Compare November 4, 2025 08:45

liji-nv force-pushed the dev-liji-set-fraction branch from 66d2532 to b678ff9 Compare November 4, 2025 11:41

liji-nv force-pushed the dev-liji-set-fraction branch from b678ff9 to 6c1e4b6 Compare November 5, 2025 02:46

HuiGao-NV reviewed Nov 6, 2025

View reviewed changes

tensorrt_llm/_torch/pyexecutor/resource_manager.py Outdated Show resolved Hide resolved

liji-nv force-pushed the dev-liji-set-fraction branch 4 times, most recently from 63989fa to 0a102f1 Compare November 6, 2025 02:57

[https://nvbugs/5570599][fix] Limit KVCache memory for estimate

8ba6b32

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv force-pushed the dev-liji-set-fraction branch from 0a102f1 to 8ba6b32 Compare November 6, 2025 03:44

liji-nv enabled auto-merge (squash) November 6, 2025 05:09

litaotju approved these changes Nov 6, 2025

View reviewed changes

kaiyux approved these changes Nov 6, 2025

View reviewed changes

liji-nv merged commit 1ef38f2 into NVIDIA:release/1.1 Nov 6, 2025
5 checks passed

liji-nv mentioned this pull request Nov 19, 2025

[https://nvbugs/5537996][fix] Let KV cache manager block initialization be aware whether it is doing a dry run or not #9093

Merged

1 task

mikeiovine pushed a commit to mikeiovine/TensorRT-LLM that referenced this pull request Nov 20, 2025

[https://nvbugs/5570599][fix] Set KVCache free_gpu_memory_fraction fo… (

56bbdde

NVIDIA#8780) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv added a commit to liji-nv/TensorRT-LLM that referenced this pull request Nov 25, 2025

[https://nvbugs/5570599][fix] Set KVCache free_gpu_memory_fraction … (N…

df9d180

…VIDIA#8780) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

liji-nv mentioned this pull request Nov 25, 2025

[None][fix] Cherry-pick conflict changes for PR 7999 PR 8515 #9446

Open

1 task

mikeiovine mentioned this pull request Nov 25, 2025

[None][chore] Weekly mass integration of release/1.1 #9343

Closed

1 task

dominicshanshan mentioned this pull request Nov 29, 2025

[None][chore] Weekly mass integration of release/1.1 -- rebase #9522

Merged

1 task

[https://nvbugs/5570599][fix] Set KVCache free_gpu_memory_fraction fo… #8780

[https://nvbugs/5570599][fix] Set KVCache free_gpu_memory_fraction fo… #8780

Uh oh!

Conversation

liji-nv commented Oct 30, 2025 • edited by coderabbitai bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

kill

skip

reuse-pipeline

Uh oh!

liji-nv commented Oct 30, 2025

Uh oh!

liji-nv commented Oct 30, 2025

Uh oh!

coderabbitai bot commented Oct 30, 2025

Walkthrough

Changes

Sequence Diagram

Estimated code review effort

Pre-merge checks and finishing touches

Uh oh!

coderabbitai bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

tensorrt-cicd commented Oct 30, 2025

Uh oh!

Uh oh!

Uh oh!

liji-nv commented Oct 31, 2025

Uh oh!

tensorrt-cicd commented Oct 31, 2025

Uh oh!

tensorrt-cicd commented Oct 31, 2025

Uh oh!

liji-nv commented Oct 31, 2025

Uh oh!

tensorrt-cicd commented Oct 31, 2025

Uh oh!

tensorrt-cicd commented Oct 31, 2025

Uh oh!

HuiGao-NV left a comment

Choose a reason for hiding this comment

Uh oh!

liji-nv commented Nov 3, 2025

Uh oh!

tensorrt-cicd commented Nov 3, 2025

Uh oh!

tensorrt-cicd commented Nov 3, 2025

Uh oh!

liji-nv commented Nov 4, 2025

Uh oh!

tensorrt-cicd commented Nov 4, 2025

Uh oh!

tensorrt-cicd commented Nov 4, 2025

Uh oh!

liji-nv commented Nov 4, 2025

Uh oh!

tensorrt-cicd commented Nov 4, 2025

Uh oh!

tensorrt-cicd commented Nov 4, 2025

Uh oh!

liji-nv commented Nov 5, 2025

Uh oh!

Uh oh!

liji-nv commented Nov 6, 2025

Uh oh!

tensorrt-cicd commented Nov 6, 2025

Uh oh!

tensorrt-cicd commented Nov 6, 2025

Uh oh!

Uh oh!

Reviewers

liji-nv commented Oct 30, 2025 •

edited by coderabbitai bot

Loading