From 37d2f0b502276ebb19c8639defbe3bca285a1d23 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed, 17 Jun 2026 11:34:42 +0800 Subject: [PATCH 1/2] [https://nvbugs/6335137][fix] Fix VSWA/Hybrid model crash with MAX_UTILIZATION Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 60b6cc14ebee..2a46214b837e 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -384,8 +384,12 @@ std::tuple MaxUtilizationScheduler::operator()( std::unordered_set seenTaskIds; // Keep track of blocks contributed by requests in context phase - auto [newlyContributedContextBlocks, newlyContributedCrossContextBlocks] - = prefillWithChunkedContextsAlreadyExecuting(activeRequests, kvCacheManager); + std::unordered_set newlyContributedContextBlocks, newlyContributedCrossContextBlocks; + if (skippingIsRelevant) + { + std::tie(newlyContributedContextBlocks, newlyContributedCrossContextBlocks) + = prefillWithChunkedContextsAlreadyExecuting(activeRequests, kvCacheManager); + } // Find last active in case we need to evict auto startedReqLambda = [this](std::shared_ptr const& req) From c84e554025c4cbd3fea67cb23ad5ed90ac133f17 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 23 Jun 2026 14:53:39 +0800 Subject: [PATCH 2/2] Add VSWA max-utilization scheduler regression test (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../batch_manager/capacityScheduler.cpp | 1 + .../batch_manager/capacitySchedulerTest.cpp | 60 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp index 90cf7094f708..c2ef0e957bfc 100644 --- a/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp +++ b/cpp/tensorrt_llm/batch_manager/capacityScheduler.cpp @@ -445,6 +445,7 @@ std::tuple MaxUtilizationScheduler::operator()( // Keep track of blocks contributed by requests in context phase std::unordered_set newlyContributedContextBlocks, newlyContributedCrossContextBlocks; + // analyzePrefixReuse asserts on variable-window managers, so this guard prevents a crash for VSWA/hybrid models. if (skippingIsRelevant) { std::tie(newlyContributedContextBlocks, newlyContributedCrossContextBlocks) diff --git a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp index 267361cd73da..07af4c9dbb20 100644 --- a/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/capacitySchedulerTest.cpp @@ -143,6 +143,32 @@ class CapacitySchedulerTest : public ::testing::Test // NOLINT(cppcoreguidelines streamPtr, maxNumTokensPerSeq, maxNumTokensPerSeq, enableReuse, cacheType); } + static std::shared_ptr getVariableWindowKvCacheManager( + SizeType32 maxNumRequests, bool enableReuse) + { + auto const numLayers = 2; + auto const nbKvHeads = 2; + auto constexpr sizePerHead = 64; + auto constexpr tokensPerBlock = 4; + auto constexpr minAttentionWindow = 8; + auto constexpr maxAttentionWindow = 16; + auto constexpr maxNumTokensPerSeq = 128; + auto constexpr blocksInPrimaryPoolPerWindow = 16; + auto constexpr blocksInSecondaryPoolPerWindow = 16; + auto const kvDtype = nvinfer1::DataType::kHALF; + CudaStreamPtr streamPtr = std::make_shared(); + + using BlocksPerWindow = std::map>; + auto const blocksPerWindow + = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPoolPerWindow, blocksInSecondaryPoolPerWindow}}, + {minAttentionWindow, {blocksInPrimaryPoolPerWindow, blocksInSecondaryPoolPerWindow}}}; + auto const maxAttentionWindowVec = std::vector{maxAttentionWindow, minAttentionWindow}; + + return std::make_shared(numLayers, nbKvHeads, sizePerHead, tokensPerBlock, + blocksPerWindow, maxNumRequests, 1, maxAttentionWindowVec, kvDtype, 0, streamPtr, maxNumTokensPerSeq, + tokensPerBlock, enableReuse); + } + static std::shared_ptr getPeftCacheManager() { return std::make_shared(); @@ -1770,6 +1796,40 @@ TEST_F(CapacitySchedulerTest, DelayDuplicateRequestChunked) } } +TEST_F(CapacitySchedulerTest, MaxUtilizationVariableWindowSkipsChunkedContextReuseAnalysis) +{ + SizeType32 maxNumRequests = 2; + bool enableReuse = true; + + auto kvCacheManager = getVariableWindowKvCacheManager(maxNumRequests, enableReuse); + ASSERT_TRUE(kvCacheManager->getBlockManager().isVariableWindow()); + ASSERT_TRUE(kvCacheManager->isEnableBlockReuse()); + + auto peftCacheManager = getPeftCacheManager(); + auto capacityScheduler + = CapacityScheduler(maxNumRequests, CapacitySchedulerPolicy::kMAX_UTILIZATION, kvCacheManager != nullptr); + + auto constexpr promptLen = 12; + auto constexpr maxNewTokens = 4; + auto request = createRequest(promptLen, maxNewTokens, 0); + request->setContextCurrentPosition(kvCacheManager->getTokensPerBlock()); + request->setContextChunkSize(kvCacheManager->getTokensPerBlock()); + ASSERT_FALSE(request->isFirstContextChunk()); + + RequestList activeRequests; + activeRequests.push_back(request); + + RequestVector scheduledRequests; + RequestVector scheduledDisaggGenInitRequests; + RequestVector pausedRequests; + EXPECT_NO_THROW(std::tie(scheduledRequests, scheduledDisaggGenInitRequests, pausedRequests) + = capacityScheduler(activeRequests, kvCacheManager, peftCacheManager)); + EXPECT_EQ(scheduledRequests.size(), 1); + EXPECT_EQ(scheduledRequests.front()->mRequestId, request->mRequestId); + EXPECT_TRUE(scheduledDisaggGenInitRequests.empty()); + EXPECT_TRUE(pausedRequests.empty()); +} + TEST_F(CapacitySchedulerTest, DelayFiveRequestsComplicated) { SizeType32 kvCacheMaxNumTokens = 1000;