Skip to content

[https://nvbugs/6418912][test] Fix KV cache reuse accounting tests #16012

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
chienchunhung wants to merge 2 commits into from
Closed

[https://nvbugs/6418912][test] Fix KV cache reuse accounting tests #16012

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
761756f
[https://nvbugs/6418912][test] Fix KV cache reuse accounting tests
chienchunhung Jul 6, 2026
2ae8fa9
[https://nvbugs/6418912][test] Run batch manager tests pre-merge
chienchunhung Jul 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6858,8 +6858,8 @@ TEST(KVCacheManagerReuseAccountingTest, CountReusableBlocksPartialMatch)
= kvCacheManager->getNeededBlocksOneStep(req1, /*twoStepsLookAhead=*/false, onlyWindowSize);
EXPECT_EQ(neededOneStep, promptLength / tokensPerBlock); // All 4 context blocks

// Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks.
EXPECT_EQ(req1.getEstimatedReusableTokens(), 0);
// Free-but-cached blocks still consume free-pool capacity, but their tokens do not require recomputation.
EXPECT_EQ(req1.getEstimatedReusableTokens(), summaryShared.reusableBlocksAll * tokensPerBlock);
}

TEST(KVCacheManagerReuseAccountingTest, GetRemainingBlocksToCompletionWithPartialReuse)
Expand Down Expand Up @@ -6976,7 +6976,7 @@ TEST(KVCacheManagerReuseAccountingTest, GetNeededBlocksOneStepWithFullReuse)
tensorrt_llm::testing::KvCacheManagerTestUtil::simulatePrefillCompletion(req0);
kvCacheManager->removeSequence(req0.mRequestId, req0);

// Second request with identical tokens - all context blocks should be reusable
// Second request with identical tokens - the recoverable cached prefix should be reusable
auto req1 = LlmRequest{
1,
maxNewTokens,
Expand All @@ -6989,11 +6989,12 @@ TEST(KVCacheManagerReuseAccountingTest, GetNeededBlocksOneStepWithFullReuse)
// getNeededBlocksOneStep must NOT subtract free reusable blocks.
auto const neededOneStep
= kvCacheManager->getNeededBlocksOneStep(req1, /*twoStepsLookAhead=*/false, onlyWindowSize);
auto const numSharedBlocks = promptLength / tokensPerBlock; // 3 blocks
EXPECT_EQ(neededOneStep, numSharedBlocks); // All 3 context blocks
auto const numContextBlocks = promptLength / tokensPerBlock; // 3 blocks
EXPECT_EQ(neededOneStep, numContextBlocks); // All 3 context blocks

// Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks.
EXPECT_EQ(req1.getEstimatedReusableTokens(), 0);
// Sequence insertion omits the final prompt token, so only the preceding full blocks are recoverable.
auto const expectedReusableBlocks = (promptLength - 1) / tokensPerBlock;
EXPECT_EQ(req1.getEstimatedReusableTokens(), expectedReusableBlocks * tokensPerBlock);
}

TEST(KVCacheManagerReuseAccountingTest, ReuseDisabledReturnsFullBlockCount)
Expand Down Expand Up @@ -7131,8 +7132,8 @@ TEST(KVCacheManagerReuseAccountingTest, MultipleRequestsWithSharedPrefix)
= kvCacheManager->getNeededBlocksOneStep(req1, /*twoStepsLookAhead=*/false, onlyWindowSize);
EXPECT_EQ(neededOneStep, promptLength / tokensPerBlock); // All 4 context blocks

// Blocks are free (released via removeSequence), so onlyAllocated=true yields 0 reusable blocks.
EXPECT_EQ(req1.getEstimatedReusableTokens(), 0);
// Free-but-cached blocks still consume free-pool capacity, but their tokens do not require recomputation.
EXPECT_EQ(req1.getEstimatedReusableTokens(), summaryPrefix.reusableBlocksAll * tokensPerBlock);

// getRemainingBlocksToCompletion: 4 context + 1 gen = 5 blocks (no subtraction; blocks are free)
auto const remaining = kvCacheManager->getRemainingBlocksToCompletion(req1, onlyWindowSize);
Expand Down
1 change: 1 addition & 0 deletions jenkins/L0_Test.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4402,6 +4402,7 @@ def launchTestJobs(pipeline, testFilter)
"A30-PyTorch-1": ["a30", "l0_a30", 1, 2],
"A30-PyTorch-2": ["a30", "l0_a30", 2, 2],
"A10-CPP-1": ["a10", "l0_a10", 1, 1],
"A30-CPP-1": ["a30", "l0_a30", 1, 1],
"A30-AutoDeploy-1": ["a30", "l0_a30", 1, 1],
"A100X-PyTorch-1": ["a100x", "l0_a100", 1, 1],
"L40S-PyTorch-1": ["l40s", "l0_l40s", 1, 2],
Expand Down
16 changes: 15 additions & 1 deletion tests/integration/test_lists/test-db/l0_a30.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,11 +38,25 @@ l0_a30:
- '*a30*'
linux_distribution_name: ubuntu*
terms:
stage: post_merge
stage: pre_merge

@chzblych chzblych Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It doesn't always sound good to me if we see a test failure in post-merge, then we just add the test to pre-merge. When everyone does this, the pre-merge pipeline will be bloat eventually.

Perhaps the question is that why existing CPP tests on A10 in pre-merge cannot cover the failure? The fix looks like testing-only changes.

@SimengLiu-nv SimengLiu-nv Jul 7, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @chzblych , kvcache cpp test is fundamental and important in my opinion. There are quite a lot of active development around kvcache. The whole suite does not take more than 15 mins to execute anyway. I strongly support putting it in the pre_merge stage.

@chienchunhung chienchunhung Jul 7, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

FYI this PR is subsumed by PR#16019, which is a superset of multiple fixes for the NVBUG.

I agree we should clearly define what tests should be included in pre-merge (run more frequently) vs. stay in post-merge (run more comprehensively) to strike the right balance between coverage and overhead. I will defer the decision to kv cache dev folks. cc: @thorjohnsen

backend: cpp
tests:
# ------------- CPP tests ---------------
- cpp/test_unit_tests.py::test_unit_tests[batch_manager-80]
- condition:
ranges:
system_gpu_count:
gte: 1
lte: 1
wildcards:
gpu:
- '*a30*'
linux_distribution_name: ubuntu*
terms:
stage: post_merge
backend: cpp
tests:
# ------------- CPP tests ---------------
- cpp/test_unit_tests.py::test_unit_tests[common-80]
- cpp/test_unit_tests.py::test_unit_tests[executor-80]
- cpp/test_unit_tests.py::test_unit_tests[kernels-80]
Expand Down
Loading