From 761756f488e06217b9d414f8f606a5d6178fe376 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 6 Jul 2026 15:54:26 -0700 Subject: [PATCH 1/2] [https://nvbugs/6418912][test] Fix KV cache reuse accounting tests Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../batch_manager/kvCacheManagerTest.cpp | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index b1c91ae09d4c..6eb7e282389b 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -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) @@ -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, @@ -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) @@ -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); From 2ae8fa9ab90b2570cabc07d6aa853f36e13a1fe1 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 6 Jul 2026 16:56:47 -0700 Subject: [PATCH 2/2] [https://nvbugs/6418912][test] Run batch manager tests pre-merge Move the A30 batch_manager C++ shard out of post-merge and register a dedicated pre-merge stage so KV-cache changes are gated before landing. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- jenkins/L0_Test.groovy | 1 + tests/integration/test_lists/test-db/l0_a30.yml | 16 +++++++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index b7b78fc5971c..b28dbb989d29 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -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], diff --git a/tests/integration/test_lists/test-db/l0_a30.yml b/tests/integration/test_lists/test-db/l0_a30.yml index c1a6ec959f87..4b178bb25c61 100644 --- a/tests/integration/test_lists/test-db/l0_a30.yml +++ b/tests/integration/test_lists/test-db/l0_a30.yml @@ -38,11 +38,25 @@ l0_a30: - '*a30*' linux_distribution_name: ubuntu* terms: - stage: post_merge + stage: pre_merge 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]