fix: remove duplicate layer multiplication in KV cache size calculation #6481

jaedeok-nvidia · 2025-07-30T09:53:54Z

Remove unnecessary multiplication by number of layers in calculateWindowSizeToShare function as cache size is already accumulated across layers.

and clamp window sizes by the max seq len before calculating KV cache pool size. This will prevent KV cache pools from being skewed by the largest window size.

Summary by CodeRabbit

Bug Fixes
- Updated the calculation of contributions for each window size to improve accuracy.
- Adjusted related test cases to reflect the revised calculation method and expected results.
- Ensured window sizes are capped to the maximum sequence length to prevent cache sizing issues.

coderabbitai · 2025-07-30T09:54:01Z

📝 Walkthrough

Walkthrough

The changes update the calculation of contributions for each window size in the BlockManager::calculateWindowSizeToShare method, removing the multiplication by the number of layers. Corresponding unit tests are adjusted to reflect this revised logic, updating expected values and test assertions accordingly. Additionally, the KVCacheManager constructor clamps window sizes to the maximum sequence length before calculating cache blocks.

Changes

Cohort / File(s)	Change Summary
BlockManager Contribution Calculation `cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp`	Revised contribution calculation in `BlockManager::calculateWindowSizeToShare` to multiply only by window size, not by the number of layers. Loop variable for layers is now unused.
Unit Test Adjustments `cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp`	Updated unit tests to align with new contribution logic: expected contributions, shares, and block counts now reflect the removal of layer count from calculations. Iteration over layers now uses an unused placeholder.
KVCacheManager Window Size Clamping `tensorrt_llm/_torch/pyexecutor/resource_manager.py`	Added clamping of each window size to be at most `max_seq_len` in the `KVCacheManager` constructor before calculating KV cache blocks, preventing cache pool sizing distortion.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

Suggested reviewers

hlu1
yilin-void

Note

⚡️ Unit Test Generation is now available in beta!

Learn more here, or try it out under "Finishing Touches" below.

📜 Recent review details

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 1cb92cc and 9c9bf9a.

📒 Files selected for processing (3)

cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp (1 hunks)
cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp (4 hunks)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (1 hunks)

🚧 Files skipped from review as they are similar to previous changes (2)

tensorrt_llm/_torch/pyexecutor/resource_manager.py
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

⏰ 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 (4)

cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp (4)

404-412: LGTM! Single window size test case is correct.

The test correctly validates that when there's only one window size, it receives the entire share (1.0f), which aligns with the expected behavior.

414-458: LGTM! Variable window size test correctly implements the layer-independent calculation.

The test properly reflects the changes described in the PR:

Uses placeholder _ to ignore the layers vector (line 434)

Calculates contributions as windowSize * cacheSizePerToken without layer multiplication (line 436)

Expected values are mathematically correct:

1024/13312 ≈ 0.0769 ✓

4096/13312 ≈ 0.3077 ✓

8192/13312 ≈ 0.6154 ✓

Rounded block counts are accurate for numPrimaryBlocks=16384

460-500: LGTM! Test with different cache sizes per window is mathematically correct.

The test correctly handles the scenario with varying cache sizes per token per window:

Properly ignores layers vector using placeholder _ (line 482)

Calculates contributions without layer dependency (line 485)

Manual verification calculations are accurate:

2048/(2048+49152+16384) ≈ 0.0303 ✓

49152/(2048+49152+16384) ≈ 0.7273 ✓

16384/(2048+49152+16384) ≈ 0.2424 ✓

502-525: LGTM! Edge case test with equal contributions is well-designed.

This test validates an important edge case where different window sizes with varying cache sizes per token result in equal contributions (8192 each), leading to equal shares (1/3 each). The test logic correctly verifies that the formula handles this scenario properly.

✨ Finishing Touches

📝 Generate Docstrings

🧪 Generate unit tests

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

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

Review comments: Directly reply to a review comment made by CodeRabbit. Example:
- I pushed a fix in commit <commit_id>, please review it.
- Explain this complex logic.
- Open a follow-up GitHub issue for this discussion.
Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
- @coderabbitai explain this code block.
- @coderabbitai modularize this function.
PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
- @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
- @coderabbitai read src/utils.ts and explain its main purpose.
- @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
- @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

@coderabbitai pause to pause the reviews on a PR.
@coderabbitai resume to resume the paused reviews.
@coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
@coderabbitai full review to do a full review from scratch and review all the files again.
@coderabbitai summary to regenerate the summary of the PR.
@coderabbitai generate docstrings to generate docstrings for this PR.
@coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
@coderabbitai generate unit tests to generate unit tests for this PR.
@coderabbitai resolve resolve all the CodeRabbit review comments.
@coderabbitai configuration to show the current CodeRabbit configuration for the repository.
@coderabbitai help to get help.

Other keywords and placeholders

Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
Add @coderabbitai or @coderabbitai title anywhere in the PR title to generate the title automatically.

Documentation and Community

Visit our Documentation for detailed information on how to use CodeRabbit.
Join our Discord Community to get help, request features, and share feedback.
Follow us on X/Twitter for updates and announcements.

jaedeok-nvidia · 2025-07-30T10:53:33Z

/bot run --disable-fail-fast

tensorrt-cicd · 2025-07-30T10:59:04Z

PR_Github #13528 [ run ] triggered by Bot

tensorrt-cicd · 2025-07-30T22:10:46Z

PR_Github #13528 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #10137 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

Remove unnecessary multiplication by number of layers in `calculateWindowSizeToShare` function as cache size is already accumulated across layers. Signed-off-by: Jaedeok Kim <[email protected]>

… window Signed-off-by: Jaedeok Kim <[email protected]>

jaedeok-nvidia · 2025-07-31T14:00:14Z

/bot run

tensorrt-cicd · 2025-07-31T14:05:54Z

PR_Github #13686 [ run ] triggered by Bot

tensorrt-cicd · 2025-08-01T02:34:21Z

PR_Github #13686 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #10278 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

…on (NVIDIA#6481) Signed-off-by: Jaedeok Kim <[email protected]> Signed-off-by: Lanyu Liao <[email protected]>

…on (NVIDIA#6481) Signed-off-by: Jaedeok Kim <[email protected]>

coderabbitai bot requested review from chuangz0 and zhengd-nv July 30, 2025 09:54

jaedeok-nvidia requested a review from netanel-haber July 30, 2025 09:55

jaedeok-nvidia self-assigned this Jul 30, 2025

jaedeok-nvidia requested review from a team and lfr-0531 and removed request for a team July 30, 2025 09:56

jaedeok-nvidia requested a review from a team as a code owner July 30, 2025 10:50

jaedeok-nvidia requested a review from qixiang-99 July 30, 2025 10:50

coderabbitai bot requested review from hlu1 and yilin-void July 30, 2025 10:50

netanel-haber approved these changes Jul 31, 2025

View reviewed changes

jaedeok-nvidia requested a review from Funatiq July 31, 2025 13:57

jaedeok-nvidia added 2 commits July 31, 2025 22:59

fix: remove duplicate layer multiplication in KV cache size calculation

e1eaae2

Remove unnecessary multiplication by number of layers in `calculateWindowSizeToShare` function as cache size is already accumulated across layers. Signed-off-by: Jaedeok Kim <[email protected]>

Clamp all window sizes to max_seq_len not to be skewed by the largest…

9c9bf9a

… window Signed-off-by: Jaedeok Kim <[email protected]>

jaedeok-nvidia force-pushed the fix/vswa-window-share-cal branch from 1cb92cc to 9c9bf9a Compare July 31, 2025 13:59

Funatiq approved these changes Jul 31, 2025

View reviewed changes

Funatiq enabled auto-merge (squash) July 31, 2025 14:05

Funatiq merged commit fbee279 into NVIDIA:main Aug 1, 2025
3 checks passed

lancelly pushed a commit to lancelly/TensorRT-LLM that referenced this pull request Aug 6, 2025

fix: remove duplicate layer multiplication in KV cache size calculati…

37b2873

…on (NVIDIA#6481) Signed-off-by: Jaedeok Kim <[email protected]> Signed-off-by: Lanyu Liao <[email protected]>

jain-ria pushed a commit to jain-ria/TensorRT-LLM that referenced this pull request Aug 7, 2025

fix: remove duplicate layer multiplication in KV cache size calculati…

2efa345

…on (NVIDIA#6481) Signed-off-by: Jaedeok Kim <[email protected]>

coderabbitai bot mentioned this pull request Aug 20, 2025

[TRTLLM-6341][feature] Support SWA KV cache #6768

Merged

coderabbitai bot mentioned this pull request Sep 9, 2025

[https://nvbugs/5501557][fix] Fix out-of-bounds vector access for model with multiple layer types #7636

Merged

1 task

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix: remove duplicate layer multiplication in KV cache size calculation #6481

fix: remove duplicate layer multiplication in KV cache size calculation #6481

Uh oh!

jaedeok-nvidia commented Jul 30, 2025 •

edited by coderabbitai bot

Loading

Uh oh!

coderabbitai bot commented Jul 30, 2025 •

edited

Loading

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

Chat

Support

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

Documentation and Community

Uh oh!

jaedeok-nvidia commented Jul 30, 2025

Uh oh!

tensorrt-cicd commented Jul 30, 2025

Uh oh!

tensorrt-cicd commented Jul 30, 2025

Uh oh!

jaedeok-nvidia commented Jul 31, 2025

Uh oh!

tensorrt-cicd commented Jul 31, 2025

Uh oh!

tensorrt-cicd commented Aug 1, 2025

Uh oh!

Uh oh!

Uh oh!

fix: remove duplicate layer multiplication in KV cache size calculation #6481

fix: remove duplicate layer multiplication in KV cache size calculation #6481

Uh oh!

Conversation

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

Uh oh!

Summary by CodeRabbit

Summary by CodeRabbit

Uh oh!

coderabbitai bot commented Jul 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Suggested reviewers

Chat

Support

CodeRabbit Commands (Invoked using PR comments)

Other keywords and placeholders

Documentation and Community

Uh oh!

jaedeok-nvidia commented Jul 30, 2025

Uh oh!

tensorrt-cicd commented Jul 30, 2025

Uh oh!

tensorrt-cicd commented Jul 30, 2025

Uh oh!

jaedeok-nvidia commented Jul 31, 2025

Uh oh!

tensorrt-cicd commented Jul 31, 2025

Uh oh!

tensorrt-cicd commented Aug 1, 2025

Uh oh!

Uh oh!

Uh oh!

jaedeok-nvidia commented Jul 30, 2025 •

edited by coderabbitai bot

Loading

coderabbitai bot commented Jul 30, 2025 •

edited

Loading