[None][fix] Cherry-pick kv_cache_manager_v2 fixes to main by lowsfer · Pull Request #14725 · NVIDIA/TensorRT-LLM

lowsfer · 2026-05-29T03:56:24Z

Description

Cherry-picks two merged kv_cache_manager_v2 fixes from feat/deepseek_v4 to main:

#14531 / e4a5978ca1c81fe304559b251b826350f21b4d7c: count only newly-added scratch blocks during resize.
#14478 / 210a290405ef9f262c49b0c5589f6657e156b012: release stale-range holders after _commit_block.

Test Coverage

Relevant coverage is in tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py, including the delayed commit path added by #14478. CI will cover this PR.

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit 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.

…s in resize (NVIDIA#14531) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> (cherry picked from commit e4a5978) Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

NVIDIA#14478) Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 210a290) Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

coderabbitai · 2026-05-29T04:02:53Z

📝 Walkthrough

Walkthrough

KV cache improvements for precise SWA scratch allocation using range intersection and stale-block lifecycle cleanup, complemented by test infrastructure enabling delayed-commit decode validation across both immediate and deferred commit timing scenarios.

Changes

KV Cache Core and Test Enhancements

Layer / File(s)	Summary
SWA scratch allocation refinement `tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`	Import `intersect` and compute SWA scratch allocation by intersecting newly added block ranges with per-lifecycle scratch ranges to derive normal block counts more precisely.
Attention-lifecycle page cleanup `tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`	Clear attention-lifecycle pages for already-committed blocks when current ordinal is within the stale range.
Test infrastructure for delayed-commit scenarios `tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`	Extend `run_request()` and `run_naive()` with `delay_commit` parameter to skip intermediate commits during decode. Expand `test_naive()` parameterization to test both immediate and delayed commit paths.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.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 describes the PR as cherry-picking kv_cache_manager_v2 fixes to main, which accurately reflects the changeset.
Description check	✅ Passed	The PR description covers the key points: what is being cherry-picked, which commits/PRs, test coverage references, and a local test attempt explanation. All required template sections are addressed.
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.}

coderabbitai

🧹 Nitpick comments (1)

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)
583-600: ⚡ Quick win

Coverage is still insufficient for the scratch-resize regression.

This matrix now exercises delayed-commit decode, but it does not pin the other _kv_cache.py fix from this cherry-pick: growing a scratch-enabled cache when the new scratch range overlaps blocks that were already allocated. Please add a focused regression in tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py that performs two resize() calls under a tight quota and asserts the second grow succeeds without charging pages for the pre-existing scratch range. Coverage is insufficient for tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py. As per coding guidelines: tests/**: Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM. Keep feedback actionable: suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py` around
lines 583 - 600, Add a focused regression test in
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py that
creates a scratch-enabled cache, sets a tight quota, calls resize() to grow the
cache into a scratch range, then calls resize() again where the new scratch
range overlaps blocks already allocated and assert the second grow succeeds
without charging pages for the pre-existing scratch range (use existing helpers
like prepare() and run_naive() patterns or directly exercise the
KvCacheV2.resize() behavior); target the runtime module
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py to ensure the
overlapping-scratch grow path is exercised; add one deterministic unit test
(e.g., test_scratch_resize_regression) that performs two resize() calls and
asserts quota/accounting unchanged for the already-allocated scratch pages.
Files to change:
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (add test)
and ensure coverage for
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py is sufficient
(current coverage: insufficient, add the test) — no other changes required.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 583-600: Add a focused regression test in
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py that
creates a scratch-enabled cache, sets a tight quota, calls resize() to grow the
cache into a scratch range, then calls resize() again where the new scratch
range overlaps blocks already allocated and assert the second grow succeeds
without charging pages for the pre-existing scratch range (use existing helpers
like prepare() and run_naive() patterns or directly exercise the
KvCacheV2.resize() behavior); target the runtime module
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py to ensure the
overlapping-scratch grow path is exercised; add one deterministic unit test
(e.g., test_scratch_resize_regression) that performs two resize() calls and
asserts quota/accounting unchanged for the already-allocated scratch pages.
Files to change:
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (add test)
and ensure coverage for
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py is sufficient
(current coverage: insufficient, add the test) — no other changes required.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 214a88bd-d246-4d3b-b9db-f118431af457

📥 Commits

Reviewing files that changed from the base of the PR and between 4826f6a and 2664aa9.

📒 Files selected for processing (2)

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py

lowsfer · 2026-05-29T04:15:39Z

/bot run --disable-fail-fast --stage-list "*PerfSanity*"

tensorrt-cicd · 2026-05-29T04:22:23Z

PR_Github #50965 [ run ] triggered by Bot. Commit: 2664aa9 Link to invocation

tensorrt-cicd · 2026-05-29T10:16:25Z

PR_Github #50965 [ run ] completed with state SUCCESS. Commit: 2664aa9
/LLM/main/L0_MergeRequest_PR pipeline #40421 (Partly Tested) 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

lowsfer · 2026-05-29T14:51:27Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-29T14:57:11Z

PR_Github #51055 [ run ] triggered by Bot. Commit: 2664aa9 Link to invocation

tensorrt-cicd · 2026-05-29T19:34:01Z

PR_Github #51055 [ run ] completed with state SUCCESS. Commit: 2664aa9
/LLM/main/L0_MergeRequest_PR pipeline #40502 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

lowsfer · 2026-05-30T07:19:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-30T07:25:22Z

PR_Github #51177 [ run ] triggered by Bot. Commit: 2664aa9 Link to invocation

tensorrt-cicd · 2026-05-30T11:05:22Z

PR_Github #51177 [ run ] completed with state SUCCESS. Commit: 2664aa9
/LLM/main/L0_MergeRequest_PR pipeline #40608 completed with status: 'SUCCESS'

CI Report

Link to invocation

lancelly

LGTM, thx!

lancelly and others added 2 commits May 29, 2026 03:15

[None][fix] kv_cache_manager_v2: count only newly-added scratch block…

e3847c9

…s in resize (NVIDIA#14531) Signed-off-by: Lance Liao <108499334+lancelly@users.noreply.github.com> (cherry picked from commit e4a5978) Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

[None][fix] Release stale-range holder after _commit_block in KVCache… (

2664aa9

NVIDIA#14478) Signed-off-by: Iman Tabrizian <10105175+tabrizian@users.noreply.github.com> (cherry picked from commit 210a290) Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>

github-actions Bot assigned lowsfer May 29, 2026

coderabbitai Bot reviewed May 29, 2026

View reviewed changes

lowsfer requested review from Tabrizian and lancelly June 2, 2026 08:29

lancelly approved these changes Jun 2, 2026

View reviewed changes

lowsfer merged commit efb71c7 into NVIDIA:main Jun 2, 2026
14 checks passed

Uh oh!

Conversation

lowsfer commented May 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented May 29, 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!

lowsfer commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

lowsfer commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

tensorrt-cicd commented May 29, 2026

Uh oh!

lowsfer commented May 30, 2026

Uh oh!

tensorrt-cicd commented May 30, 2026

Uh oh!

tensorrt-cicd commented May 30, 2026

Uh oh!

lancelly left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

lowsfer commented May 29, 2026 •

edited

Loading