[None][test] KV cache manager v2: add V2 + VSWA multi-GPU test coverage by eopXD · Pull Request #16114 · NVIDIA/TensorRT-LLM

eopXD · 2026-07-08T09:06:21Z

Description

Note — scope changed. This PR originally carried a fix for the KV Cache
Manager V2 quota-inflation OOM on multi-GPU VSWA configs. A parallel effort,
#15991 (nvbugs/6418103, "Clamp the post-allreduce quota by the
pre-allreduce quota"), landed that fix on main first. This PR is now rebased
on top of it and scoped down to amend the test coverage that #15991 landed
without — the V2 + VSWA multi-GPU path. No production code changes remain.

Background: the bug that #15991 fixed

With world_size > 1, KVCacheManagerV2.__init__ normalizes the per-rank GPU
byte quota across ranks by converting quota → tokens, taking an all-reduce
MIN, then converting back. The back-conversion could raise the quota above
the requested budget: when the per-rank budget is below the fixed per-batch
sliding-window generation floor
(max_batch_size * generation_swa_size_per_request),
_get_max_tokens_from_quota returns 0, and _get_quota_from_max_tokens(0)
returns exactly that floor — a value independent of free_gpu_memory_fraction.
The inflated quota then drove eager KV pool allocation past available memory →
CuOOMError, reproducing only on V2 + VSWA + world_size > 1. Regression from
#15633.

#15991 clamps the normalized quota to the incoming quota
(quota = min(quota, self._get_quota_from_max_tokens(max_tokens))) so the
cross-rank step can only reduce, never over-commit.

Test Coverage (this PR)

The landed fix went in without a test guarding the V2 + VSWA path. This PR adds
it.

accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_vswa_reuse_4gpus is
now parametrized over v2_kv_cache, adding the previously-uncovered V2 path:

test_eagle3_vswa_reuse_4gpus[v2_kv_cache-one_model] — exercises the fixed
path (GPT-OSS-120B, TP=4, max_attention_window=[128, 32768],
free_gpu_memory_fraction=0.4); this variant OOM'd before [https://nvbugs/6418103][fix] Clamp the post-allreduce quota by the pre-allreduce quota (`quota = min(quota… #15991.
test_eagle3_vswa_reuse_4gpus[v1_kv_cache-one_model],
test_eagle3_vswa_reuse_4gpus[v1_kv_cache-two_model] — existing V1 coverage,
renamed by the new axis.

The v2_kv_cache + two-model combination is skipped (V2 is not compatible with
two-model overlap scheduling).

Registered in:

tests/integration/test_lists/qa/llm_function_core.txt
tests/integration/test_lists/qa/llm_function_rtx6k.txt
tests/integration/test_lists/test-db/l0_dgx_b200.yml — runs the V2 + VSWA
path in the DGX_B200-4_GPUs-PyTorch-Post-Merge stage on Blackwell to guard
the fix post-merge.

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths.
No API changes (no api-compatible / api-breaking label needed).
No new dependencies.
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.

eopXD · 2026-07-08T09:06:31Z

/bot run --disable-fail-fast

coderabbitai · 2026-07-08T09:12:07Z

📝 Walkthrough

Walkthrough

Adds a safeguard to KVCacheManagerV2 that clamps per-rank quota to prevent increases across ranks during multi-rank initialization, logging a warning when the normalized quota exceeds the existing quota. Also extends the eagle3_vswa_reuse_4gpus test and QA lists to cover v1/v2 KV cache manager variants.

Changes

KV Cache Manager V2 Quota Fix and Test Coverage

Layer / File(s)	Summary
Quota clamping in KVCacheManagerV2 init `tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`	Computes a normalized quota from max_tokens, warns when it exceeds the current quota, and clamps quota to the minimum of the two to prevent per-rank KV budget increases.
Test parameterization and QA list updates `tests/integration/defs/accuracy/test_llm_api_pytorch.py`, `tests/integration/test_lists/qa/llm_function_core.txt`, `tests/integration/test_lists/qa/llm_function_rtx6k.txt`	Parameterizes test_eagle3_vswa_reuse_4gpus with v2_kv_cache, propagates use_kv_cache_manager_v2 into KvCacheConfig, skips the incompatible two_model+v2 combination, and updates QA lists with new v1/v2 and one_model/two_model test identifiers.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: schetlur-nv, lowsfer, mikeiovine

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
Title check	✅ Passed	The title is concise, specific, and accurately describes the added V2 + VSWA multi-GPU test coverage.
Description check	✅ Passed	The PR description follows the template and includes clear description, test coverage, and checklist items.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

tensorrt-cicd · 2026-07-08T09:13:22Z

PR_Github #58210 [ run ] triggered by Bot. Commit: 24315af Link to invocation

eopXD · 2026-07-08T09:25:50Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T09:31:22Z

PR_Github #58213 [ run ] triggered by Bot. Commit: 29a40fe Link to invocation

tensorrt-cicd · 2026-07-08T09:35:00Z

PR_Github #58210 [ run ] completed with state ABORTED. Commit: 24315af

Link to invocation

eopXD · 2026-07-08T12:52:47Z

/bot run --post-merge --disable-fail-fast --stage-list "DGX_B200-4_GPUs-PyTorch-Post-Merge-*"

tensorrt-cicd · 2026-07-08T12:58:42Z

PR_Github #58233 [ run ] triggered by Bot. Commit: 29a40fe Link to invocation

tensorrt-cicd · 2026-07-08T13:03:31Z

PR_Github #58213 [ run ] completed with state ABORTED. Commit: 29a40fe

Link to invocation

tensorrt-cicd · 2026-07-08T15:15:26Z

PR_Github #58233 [ run ] completed with state FAILURE. Commit: 29a40fe
/LLM/main/L0_MergeRequest_PR pipeline #46873 (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

eopXD · 2026-07-09T03:36:41Z

/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-PyTorch-Post-Merge-1"

tensorrt-cicd · 2026-07-09T03:43:31Z

PR_Github #58361 [ run ] triggered by Bot. Commit: 29a40fe Link to invocation

nvpohanh · 2026-07-09T07:29:17Z

[by Codex] @yizhang-nv Could you review this PR? Thanks!

tensorrt-cicd · 2026-07-09T08:33:03Z

PR_Github #58361 [ run ] completed with state SUCCESS. Commit: 29a40fe
/LLM/main/L0_MergeRequest_PR pipeline #46987 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

eopXD · 2026-07-09T08:46:10Z

/bot run --disable-fail-fast --extra-stage "DGX_B200-4_GPUs-PyTorch-Post-Merge-1"

tensorrt-cicd · 2026-07-09T08:52:23Z

PR_Github #58435 [ run ] triggered by Bot. Commit: 29a40fe Link to invocation

tensorrt-cicd · 2026-07-09T10:15:56Z

PR_Github #58435 [ run ] completed with state SUCCESS. Commit: 29a40fe
/LLM/main/L0_MergeRequest_PR pipeline #47049 completed with status: 'SUCCESS'

CI Report

Link to invocation

The KV Cache Manager V2 quota-inflation OOM on multi-GPU VSWA (variable sliding-window attention) configurations -- a regression from NVIDIA#15633 -- was fixed on main by a parallel effort, NVIDIA#15991 (nvbugs/6418103), which clamps the post-allreduce quota by the pre-allreduce quota so the cross-rank normalization can only reduce, never over-commit. That fix landed without a test guarding the V2 + VSWA path. This change amends the missing coverage. Parametrize TestGPTOSS::test_eagle3_vswa_reuse_4gpus over v2_kv_cache so the previously-uncovered V2 path (GPT-OSS-120B, TP=4, max_attention_window=[128, 32768], free_gpu_memory_fraction=0.4) runs and guards against regressing the fix -- this variant OOM'd before NVIDIA#15991. The v2_kv_cache + two-model combination is skipped (V2 is not compatible with two-model overlap scheduling). Register the new variants in llm_function_core.txt and llm_function_rtx6k.txt. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>

Add test_eagle3_vswa_reuse_4gpus (v1/v2 x one/two model) to l0_dgx_b200.yml, next to the sibling test_eagle3_4gpus, so the V2 + VSWA path runs in the DGX_B200-4_GPUs-PyTorch-Post-Merge stage on Blackwell and guards the quota-inflation fix post-merge. Signed-off-by: Yueh-Ting Chen <yuehtingc@nvidia.com>

eopXD · 2026-07-10T06:35:10Z

/bot skip --comment "Previous CI gain success. Now that #15991 landed before this MR does, we are purely adding test coverage now. Safe to skip CI."

tensorrt-cicd · 2026-07-10T06:41:51Z

PR_Github #58619 [ skip ] triggered by Bot. Commit: f06d474 Link to invocation

tensorrt-cicd · 2026-07-10T06:49:18Z

PR_Github #58619 [ skip ] completed with state SUCCESS. Commit: f06d474
Skipping testing for commit f06d474

Link to invocation

eopXD requested review from a team as code owners July 8, 2026 09:06

github-actions Bot assigned eopXD Jul 8, 2026

yufeiwu-nv approved these changes Jul 9, 2026

View reviewed changes

nvpohanh requested a review from yizhang-nv July 9, 2026 06:28

eopXD added 2 commits July 10, 2026 14:31

eopXD changed the title ~~[None][fix] KV cache manager v2: prevent VSWA quota inflation OOM (multi-GPU)~~ [None][test] KV cache manager v2: add V2 + VSWA multi-GPU test coverage Jul 10, 2026

eopXD force-pushed the user/yuehtingc/kvcv2-vswa-quota-oom-fix branch from 29a40fe to f06d474 Compare July 10, 2026 06:35

crazydemo approved these changes Jul 10, 2026

View reviewed changes

eopXD merged commit e2ce7a6 into NVIDIA:main Jul 10, 2026
8 checks passed

Uh oh!

Conversation

eopXD commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Background: the bug that #15991 fixed

Test Coverage (this PR)

PR Checklist

GitHub Bot Help

Uh oh!

eopXD commented Jul 8, 2026

Uh oh!

coderabbitai Bot commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

eopXD commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

eopXD commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

eopXD commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

nvpohanh commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

eopXD commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

eopXD commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

tensorrt-cicd commented Jul 10, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

eopXD commented Jul 8, 2026 •

edited

Loading

coderabbitai Bot commented Jul 8, 2026 •

edited

Loading