[KV cache manager] Simplify block allocation in KVCacheManager::addSequence

eopXD · eopXD · commit 2a546347e246 · 2025-08-11T06:29:16.000-07:00
Signed-off-by: eopXD &lt;yuehtingc@nvidia.com&gt;
diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -1891,11 +1891,7 @@ void KVCacheManager::addSequence(
 
     for (auto const [windowSize, metadata] : mBlockManager.getWindowSizesMetadata())
     {
-        auto const maxTokenNum = metadata.maxTokenNum;
-        auto const temporaryAttentionWindow = metadata.temporaryAttentionWindow;
-
-        // Consider the temporaryAttentionWindow when allocating blocks.
-        auto const effectiveInputLength = std::min(inputLength, maxTokenNum + temporaryAttentionWindow);
+        auto const effectiveInputLength = std::min(inputLength, windowSize);
         auto const numContextBlocks = tc::ceilDiv(effectiveInputLength, getTokensPerBlock());
         if (!sequence.isCyclic() && mEnableBlockReuse)
         {

Original file line number	Diff line number	Diff line change
`@@ -1891,11 +1891,7 @@ void KVCacheManager::addSequence(`
`1891`	`1891`
`1892`	`1892`	`for (auto const [windowSize, metadata] : mBlockManager.getWindowSizesMetadata())`
`1893`	`1893`	`{`
`1894`		`- auto const maxTokenNum = metadata.maxTokenNum;`
`1895`		`- auto const temporaryAttentionWindow = metadata.temporaryAttentionWindow;`
`1896`		`-`
`1897`		`- // Consider the temporaryAttentionWindow when allocating blocks.`
`1898`		`- auto const effectiveInputLength = std::min(inputLength, maxTokenNum + temporaryAttentionWindow);`
	`1894`	`+ auto const effectiveInputLength = std::min(inputLength, windowSize);`
`1899`	`1895`	`auto const numContextBlocks = tc::ceilDiv(effectiveInputLength, getTokensPerBlock());`
`1900`	`1896`	`if (!sequence.isCyclic() && mEnableBlockReuse)`
`1901`	`1897`	`{`