[None][fix] Fix a capacity issue in KVCacheManagerV2 for SWA compatibility by heyuhhh · Pull Request #12968 · NVIDIA/TensorRT-LLM

heyuhhh · 2026-04-12T14:56:11Z

Summary by CodeRabbit

Bug Fixes
- Improved cache memory allocation and resource management during text generation operations for better overall performance.
- Fixed handling of allocated memory capacity to optimize the execution of generation tasks.
- Enhanced resource utilization efficiency for concurrent generation and context processing operations.
- Optimized cache management strategies to improve system efficiency and reliability.

Description

For generation phase we only need ceil(window_size / tokens_per_block) blocks in SWA. Otherwise CacheManager can't handle the situation that applied blocks are more than reserved blocks, which may cause illegal memory access in warmup stages

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.

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

heyuhhh · 2026-04-12T14:57:59Z

/bot help

github-actions · 2026-04-12T14:58:07Z

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) --high-priority]

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.

--high-priority (OPTIONAL) : Run the pipeline with high priority. This option is restricted to authorized users only and will route the job to a high-priority queue.

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.

heyuhhh · 2026-04-12T14:58:50Z

/bot run --disable-fail-fast

coderabbitai · 2026-04-12T14:59:37Z

📝 Walkthrough

Walkthrough

Modified KV cache resizing logic in KVCacheManagerV2.add_dummy_requests to conditionally pass history_length parameter based on request type. Generation requests now pass full dummy_capacity as history length, while context requests pass None.

Changes

Cohort / File(s)	Summary
KV Cache Resizing Logic `tensorrt_llm/_torch/pyexecutor/resource_manager.py`	Updated dummy KV cache resizing to use conditional `history_length` parameter. When `is_gen` is true, `history_hint` is set to `dummy_capacity`; otherwise `None`. Alters resize semantics for generation vs. context dummy requests.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 2 | ❌ 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 (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly references the main change (capacity issue in KVCacheManagerV2 for SWA compatibility) and matches the code modification shown in the summary.
Description check	✅ Passed	The PR description includes a substantive explanation of the issue and solution, but the Test Coverage section is empty and lacks specific test details required by the template.

_{✏️ 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.}

coderabbitai

Actionable comments posted: 1

🤖 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/resource_manager.py`:
- Around line 2297-2302: The generation warmup sets history_hint and passes it
to kv_cache.resize (history_hint = dummy_capacity if is_gen else None), but
subsequent resizes omit history_length and can desync history from capacity;
update every resize call performed during generation (both the main cache
kv_cache.resize(...) and the draft cache resize call(s) such as
draft_kv_cache.resize(...) or draft_cache.resize(...)) to pass
history_length=history_hint when is_gen is true (i.e., reuse the existing
history_hint/dummy_capacity variable) so all generation resizes remain in sync
and preserve the stale-block optimization.

🪄 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: Pro

Run ID: 3087e83d-2ceb-48dc-9cc7-afb54932032f

📥 Commits

Reviewing files that changed from the base of the PR and between 5653803 and fc1d412.

📒 Files selected for processing (1)

tensorrt_llm/_torch/pyexecutor/resource_manager.py

tensorrt-cicd · 2026-04-12T15:04:30Z

PR_Github #42880 [ run ] triggered by Bot. Commit: fc1d412 Link to invocation

tensorrt-cicd · 2026-04-12T18:55:12Z

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

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

heyuhhh · 2026-04-13T04:43:39Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-13T04:50:31Z

PR_Github #42948 [ run ] triggered by Bot. Commit: c2ad989 Link to invocation

tensorrt-cicd · 2026-04-13T10:26:12Z

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

heyuhhh · 2026-04-13T10:55:13Z

/bot run

tensorrt-cicd · 2026-04-13T11:01:24Z

PR_Github #43027 [ run ] triggered by Bot. Commit: c2ad989 Link to invocation

tensorrt-cicd · 2026-04-13T13:11:57Z

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

heyuhhh · 2026-04-13T14:50:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-13T14:55:59Z

PR_Github #43064 [ run ] triggered by Bot. Commit: bd3a85d Link to invocation

tensorrt-cicd · 2026-04-13T17:41:19Z

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

heyuhhh · 2026-04-14T01:16:01Z

/bot run

tensorrt-cicd · 2026-04-14T01:23:10Z

PR_Github #43125 [ run ] triggered by Bot. Commit: bd3a85d Link to invocation

tensorrt-cicd · 2026-04-14T02:00:21Z

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

tensorrt-cicd · 2026-04-15T01:39:11Z

PR_Github #43319 [ run ] triggered by Bot. Commit: ad327c8 Link to invocation

tensorrt-cicd · 2026-04-15T02:21:35Z

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

heyuhhh · 2026-04-15T03:38:31Z

/bot run

tensorrt-cicd · 2026-04-15T03:46:17Z

PR_Github #43369 [ run ] triggered by Bot. Commit: b403ff3 Link to invocation

nvpohanh · 2026-04-15T05:15:02Z

@eopXD could you review this? thanks

tensorrt-cicd · 2026-04-15T05:28:05Z

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

heyuhhh · 2026-04-15T05:56:23Z

/bot run

tensorrt-cicd · 2026-04-15T06:01:52Z

PR_Github #43405 [ run ] triggered by Bot. Commit: b403ff3 Link to invocation

tensorrt-cicd · 2026-04-15T09:37:11Z

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

heyuhhh · 2026-04-15T09:44:10Z

/bot run

tensorrt-cicd · 2026-04-15T09:49:47Z

PR_Github #43471 [ run ] triggered by Bot. Commit: b403ff3 Link to invocation

tensorrt-cicd · 2026-04-15T14:03:50Z

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

heyuhhh · 2026-04-15T15:58:54Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-15T16:08:05Z

PR_Github #43523 [ run ] triggered by Bot. Commit: 065f9a5 Link to invocation

tensorrt-cicd · 2026-04-15T20:30:59Z

PR_Github #43523 [ run ] completed with state SUCCESS. Commit: 065f9a5
/LLM/main/L0_MergeRequest_PR pipeline #34038 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

heyuhhh · 2026-04-16T02:46:09Z

/bot run

tensorrt-cicd · 2026-04-16T02:51:43Z

PR_Github #43627 [ run ] triggered by Bot. Commit: 065f9a5 Link to invocation

tensorrt-cicd · 2026-04-16T04:49:13Z

PR_Github #43627 [ run ] completed with state FAILURE. Commit: 065f9a5
/LLM/main/L0_MergeRequest_PR pipeline #34117 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

heyuhhh · 2026-04-16T04:51:33Z

/bot run

tensorrt-cicd · 2026-04-16T04:56:50Z

PR_Github #43648 [ run ] triggered by Bot. Commit: 065f9a5 Link to invocation

tensorrt-cicd · 2026-04-16T05:15:41Z

PR_Github #43648 [ run ] completed with state FAILURE. Commit: 065f9a5
/LLM/main/L0_MergeRequest_PR pipeline #34136 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

heyuhhh · 2026-04-16T07:01:56Z

/bot run

tensorrt-cicd · 2026-04-16T07:08:52Z

PR_Github #43684 [ run ] triggered by Bot. Commit: 065f9a5 Link to invocation

tensorrt-cicd · 2026-04-16T10:17:35Z

PR_Github #43684 [ run ] completed with state SUCCESS. Commit: 065f9a5
/LLM/main/L0_MergeRequest_PR pipeline #34169 completed with status: 'SUCCESS'

CI Report

Link to invocation

heyuhhh · 2026-04-16T11:32:04Z

Hi @nvpohanh @lowsfer , the CI has passed, could you please help to merge this PR? I've discussed with @yizhang-nv offline and he has approved this PR.

add history_hint in capacity resize

fc1d412

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

heyuhhh requested a review from a team as a code owner April 12, 2026 14:56

github-actions Bot assigned heyuhhh Apr 12, 2026

heyuhhh requested a review from lowsfer April 12, 2026 14:56

coderabbitai Bot reviewed Apr 12, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/resource_manager.py Outdated

heyuhhh added 2 commits April 12, 2026 20:56

Fix to use real meaning of token_num in gen phase

d710087

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

minor

4727bd5

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

heyuhhh requested a review from yizhang-nv April 13, 2026 04:11

yizhang-nv approved these changes Apr 13, 2026

View reviewed changes

minor fix

c2ad989

Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com>

Merge branch 'main' into user/yuhangh/fix_ima_in_swa

bd3a85d

Merge branch 'main' into user/yuhangh/fix_ima_in_swa

a900657

Merge branch 'main' into user/yuhangh/fix_ima_in_swa

b403ff3

nvpohanh approved these changes Apr 15, 2026

View reviewed changes

Merge branch 'main' into user/yuhangh/fix_ima_in_swa

065f9a5

lfr-0531 enabled auto-merge (squash) April 17, 2026 06:55

lfr-0531 merged commit 461f3b9 into NVIDIA:main Apr 17, 2026
5 checks passed

Conversation

heyuhhh commented Apr 12, 2026 • 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

Uh oh!

heyuhhh commented Apr 12, 2026

Uh oh!

github-actions Bot commented Apr 12, 2026

GitHub Bot Help

kill

skip

reuse-pipeline

Uh oh!

heyuhhh commented Apr 12, 2026

Uh oh!

coderabbitai Bot commented Apr 12, 2026

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented Apr 12, 2026

Uh oh!

tensorrt-cicd commented Apr 12, 2026

Uh oh!

heyuhhh commented Apr 13, 2026

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

heyuhhh commented Apr 13, 2026

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

heyuhhh commented Apr 13, 2026

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

tensorrt-cicd commented Apr 13, 2026

Uh oh!

heyuhhh commented Apr 14, 2026

Uh oh!

tensorrt-cicd commented Apr 14, 2026

Uh oh!

tensorrt-cicd commented Apr 14, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

Uh oh!

heyuhhh commented Apr 15, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

Uh oh!

nvpohanh commented Apr 15, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

Uh oh!

heyuhhh commented Apr 15, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

Uh oh!

heyuhhh commented Apr 15, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

Uh oh!

tensorrt-cicd commented Apr 15, 2026

heyuhhh commented Apr 12, 2026 •

edited by coderabbitai Bot

Loading