[None][fix] Fix cache buffer size for window (#8320)

chuangz0 · web-flow · commit 40d129a41596 · 2025-10-16T09:01:11.000+08:00
Signed-off-by: Chuang Zhu &lt;111838961+chuangz0@users.noreply.github.com&gt;
diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp
@@ -210,7 +210,11 @@ CacheTransBufferManager::CacheTransBufferManager(
         {
             auto poolIdx = mCacheManager->getBlockManager().getLayerPoolIdx(layerId);
             auto windowSize = static_cast<size_t>(mCacheManager->getBlockManager().getPoolWindowSize(poolIdx));
-            auto validTokenNum = (windowSize < maxNumTokens.value() ? windowSize : maxNumTokens.value());
+            auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock;
+            auto validTokenNum = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value());
+            // if windowSize % (tokensPerBlock) !=0
+            validTokenNum += tokensPerBlock; // add one more block
+
             bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer;
         }
     }
@@ -237,7 +241,7 @@ CacheTransBufferManager::CacheTransBufferManager(
     allocateBuffer();
 }
 
-size_t CacheTransBufferManager::preAllocBufferSize(
+size_t CacheTransBufferManager::preAllocBufferSize(size_t tokensPerBlock,
     std::map<SizeType32, SizeType32> const& cacheSizeBytesPerTokenPerWindow,
     std::optional<executor::CacheTransceiverConfig> const& cacheTransceiverConfig)
 {
@@ -256,9 +260,9 @@ size_t CacheTransBufferManager::preAllocBufferSize(
         TransferBufferSize = 0;
         for (auto const& [windowSize, cacheSizeBytesPerToken] : cacheSizeBytesPerTokenPerWindow)
         {
-            auto validTokenNum
-                = (static_cast<size_t>(windowSize) < maxNumTokens.value() ? static_cast<size_t>(windowSize)
-                                                                          : maxNumTokens.value());
+            auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock;
+            auto validTokenNum = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value());
+            validTokenNum += tokensPerBlock; // add one more block
             TransferBufferSize += validTokenNum * cacheSizeBytesPerToken;
         }
     }
diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h
@@ -60,7 +60,8 @@ class CacheTransBufferManager
     CacheTransBufferManager(
         KVCacheManager::BaseKVCacheManager* cacheManager, std::optional<size_t> maxNumTokens = std::nullopt);
 
-    static size_t preAllocBufferSize(std::map<SizeType32, SizeType32> const& cacheSizeBytesPerTokenPerWindow,
+    static size_t preAllocBufferSize(size_t tokensPerBlock,
+        std::map<SizeType32, SizeType32> const& cacheSizeBytesPerTokenPerWindow,
         std::optional<executor::CacheTransceiverConfig> const& cacheTransceiverConfig = std::nullopt);
 
     std::optional<int> assignBufferIndexForSend();
diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp
@@ -524,7 +524,18 @@ class CacheSender::Impl
 
             if (isReady)
             {
-                asyncSendAndRemoveResponse(it->first, std::move(it->second));
+                if (dynamic_cast<executor::kv_cache::AgentConnectionManager*>(mManager) != nullptr)
+                {
+                    // our nixl impl seems only support recv and send in the same thread
+                    //  if we use zmq as control path, we may avoid this issue
+                    sendAndRemoveResponse(it->first, std::move(it->second));
+                }
+                else
+                {
+                    // if we send data in another thread, multiple rank may send data for different requests at the same
+                    // time with gen DP case.
+                    asyncSendAndRemoveResponse(it->first, std::move(it->second));
+                }
                 removeResponse(it);
             }
             else
diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp
@@ -306,7 +306,7 @@ TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr<nvinfer
         auto const cacheSizeBytesPerTokenPerWindow = calculateCacheSizePerTokenForDisagg(
             mModelConfig, mWorldConfig, getMaxAttentionWindowVec(), mModelConfig.useCrossAttention(), 2);
         auto cacheTransPreAllocaSize = kv_cache_manager::CacheTransBufferManager::preAllocBufferSize(
-            cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig);
+            mModelConfig.getTokensPerBlock(), cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig);
 
         auto const [freePrimaryMemBytes, freeSecondaryMemBytes]
             = BaseKVCacheManager::calculateFreeMemBytes(mRuntime->getBufferManager(), kvCacheConfig);
diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp
@@ -182,5 +182,6 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m)
         .def(nb::init<tb::kv_cache_manager::BaseKVCacheManager*, std::optional<size_t>>(), nb::arg("cache_manager"),
             nb::arg("max_num_tokens") = std::nullopt)
         .def_static("pre_alloc_buffer_size", &tb::kv_cache_manager::CacheTransBufferManager::preAllocBufferSize,
-            nb::arg("cache_size_bytes_per_token_per_window"), nb::arg("cache_transceiver_config") = nb::none());
+            nb::arg("tokens_per_block"), nb::arg("cache_size_bytes_per_token_per_window"),
+            nb::arg("cache_transceiver_config") = nb::none());
 }
diff --git a/cpp/tensorrt_llm/pybind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/pybind/batch_manager/cacheTransceiver.cpp
@@ -178,5 +178,6 @@ void tb::CacheTransceiverBindings::initBindings(py::module_& m)
         .def(py::init<tb::kv_cache_manager::BaseKVCacheManager*, std::optional<size_t>>(), py::arg("cache_manager"),
             py::arg("max_num_tokens") = std::nullopt)
         .def_static("pre_alloc_buffer_size", &tb::kv_cache_manager::CacheTransBufferManager::preAllocBufferSize,
-            py::arg("cache_size_bytes_per_token_per_window"), py::arg("cache_transceiver_config") = py::none());
+            py::arg("tokens_per_block"), py::arg("cache_size_bytes_per_token_per_window"),
+            py::arg("cache_transceiver_config") = py::none());
 }
diff --git a/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp b/cpp/tests/unit_tests/batch_manager/cacheTransBufferTest.cpp
@@ -114,8 +114,8 @@ TEST_F(CacheTransBufferTest, TestPreAllocBufferSize)
             {maxBlocksPerSeq * tokensPerBlock, cacheSizeBytesPerToken}};
         tensorrt_llm::executor::CacheTransceiverConfig cacheTransceiverConfig{
             tensorrt_llm::executor::CacheTransceiverConfig::BackendType::UCX, maxNumTokens};
-        size_t bufferSizeBytes
-            = CacheTransBufferManager::preAllocBufferSize(cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig);
+        size_t bufferSizeBytes = CacheTransBufferManager::preAllocBufferSize(
+            tokensPerBlock, cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig);
         auto bufferId = mTransBufferManager->assignBufferIndexForSend();
         EXPECT_TRUE(bufferId.has_value());
         EXPECT_EQ(bufferId.value(), 0);
@@ -156,8 +156,8 @@ TEST_F(CacheTransBufferTest, TestPreAllocBufferSize2)
             tensorrt_llm::executor::CacheTransceiverConfig::BackendType::UCX, maxNumTokens};
         std::map<SizeType32, SizeType32> cacheSizeBytesPerTokenPerWindow{
             {maxBlocksPerSeq * tokensPerBlock, cacheSizeBytesPerToken}};
-        size_t bufferSizeBytes
-            = CacheTransBufferManager::preAllocBufferSize(cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig);
+        size_t bufferSizeBytes = CacheTransBufferManager::preAllocBufferSize(
+            tokensPerBlock, cacheSizeBytesPerTokenPerWindow, cacheTransceiverConfig);
         auto bufferId = mTransBufferManager->assignBufferIndexForSend();
         EXPECT_TRUE(bufferId.has_value());
         EXPECT_EQ(bufferId.value(), 0);
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py
@@ -145,7 +145,8 @@ def __init__(self, kv_cache_manager: KVCacheManager, max_num_tokens: int):
                                                max_num_tokens)
 
     @staticmethod
-    def pre_alloc_buffer_size(kv_cache_size_per_token: int,
+    def pre_alloc_buffer_size(tokens_per_block: int,
+                              kv_cache_size_per_token: int,
                               cache_transceiver_config: CacheTransceiverConfig):
         return CacheTransBufferManagerCpp.pre_alloc_buffer_size(
-            kv_cache_size_per_token, cache_transceiver_config)
+            tokens_per_block, kv_cache_size_per_token, cache_transceiver_config)
diff --git a/tests/integration/defs/cpp/test_multi_gpu.py b/tests/integration/defs/cpp/test_multi_gpu.py
@@ -108,22 +108,6 @@ def run_cache_transceiver_tests(build_dir: _pl.Path,
                      env=mgpu_env,
                      timeout=timeout)
 
-    # Nixl transfer agent tests
-    new_env = get_multi_gpu_env(kv_cache_type=KVCacheType.NIXL)
-
-    # Cache transceiver tests
-    cache_trans_test_8_proc = [
-        "mpirun",
-        "-n",
-        "8",
-        "--allow-run-as-root",
-        "cacheTransceiverTest",
-    ]
-    _cpp.run_command(cache_trans_test_8_proc,
-                     cwd=tests_dir,
-                     env=new_env,
-                     timeout=600)
-
 
 def run_user_buffer_tests(build_dir: _pl.Path, nprocs=2, timeout=300):
     tests_dir = build_dir / "tests" / "unit_tests" / "multi_gpu"
@@ -500,8 +484,8 @@ def test_fused_gemm_allreduce(build_google_tests, nprocs, build_dir):
 
 @pytest.mark.parametrize("build_google_tests", ["80", "86", "89", "90"],
                          indirect=True)
-@pytest.mark.parametrize("kvcache_type", [KVCacheType.MPI, KVCacheType.UCX],
-                         ids=["mpi_kvcache", "ucx_kvcache"])
+@pytest.mark.parametrize("kvcache_type", [KVCacheType.NIXL, KVCacheType.UCX],
+                         ids=["nixl_kvcache", "ucx_kvcache"])
 @pytest.mark.parametrize("nprocs", [2, 8], ids=["2proc", "8proc"])
 def test_cache_transceiver(build_google_tests, nprocs, kvcache_type, build_dir):
 
diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml
@@ -189,9 +189,8 @@ l0_dgx_h100:
   # ------------- CPP tests ---------------
   - cpp/test_multi_gpu.py::test_mpi_utils[90]
   - cpp/test_multi_gpu.py::test_fused_gemm_allreduce[4proc-90]
-  - cpp/test_multi_gpu.py::test_cache_transceiver[2proc-mpi_kvcache-90]
   - cpp/test_multi_gpu.py::test_cache_transceiver[2proc-ucx_kvcache-90]
-  - cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mpi_kvcache-90]
+  - cpp/test_multi_gpu.py::test_cache_transceiver[8proc-nixl_kvcache-90]
   - cpp/test_multi_gpu.py::test_cache_transceiver[8proc-ucx_kvcache-90]
   - cpp/test_multi_gpu.py::test_user_buffer[2proc-90]
   - cpp/test_multi_gpu.py::test_enc_dec[t5-90]
diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
@@ -334,8 +334,6 @@ unittest/_torch/multi_gpu_modeling/test_llama3.py::test_llama_3_3 SKIP (https://
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-ep4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/5541494)
 accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[dep4_latency_moe_trtllm-torch_compile=False] SKIP (https://nvbugs/5541494)
 accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_nvfp4_4gpus[moe_backend=TRTLLM-mtp_nextn=2-tp4-fp8kv=True-attention_dp=True-cuda_graph=True-overlap_scheduler=True-torch_compile=False] SKIP (https://nvbugs/5541494)
-cpp/test_multi_gpu.py::test_cache_transceiver[2proc-mpi_kvcache-90] SKIP (https://nvbugs/5546202)
-cpp/test_multi_gpu.py::test_cache_transceiver[8proc-mpi_kvcache-90] SKIP (https://nvbugs/5546202)
 examples/test_eagle.py::test_llm_eagle_1gpu_modelopt_ckpt[llama3.1-eagle-8b-hf_v0.5-float16-bs8] SKIP (https://nvbugs/5546507)
 examples/test_eagle.py::test_llm_eagle_1gpu[EAGLE-Vicuna-7B-v1.3-float16-bs1-eagle1] SKIP (https://nvbugs/5546507)
 examples/test_eagle.py::test_llm_eagle_1gpu[EAGLE-Vicuna-7B-v1.3-float16-bs1-eagle2] SKIP (https://nvbugs/5546507)

Original file line number	Diff line number	Diff line change
`@@ -210,7 +210,11 @@ CacheTransBufferManager::CacheTransBufferManager(`
`210`	`210`	`{`
`211`	`211`	`auto poolIdx = mCacheManager->getBlockManager().getLayerPoolIdx(layerId);`
`212`	`212`	`auto windowSize = static_cast<size_t>(mCacheManager->getBlockManager().getPoolWindowSize(poolIdx));`
`213`		`- auto validTokenNum = (windowSize < maxNumTokens.value() ? windowSize : maxNumTokens.value());`
	`213`	`+ auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock;`
	`214`	`+ auto validTokenNum = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value());`
	`215`	`+ // if windowSize % (tokensPerBlock) !=0`
	`216`	`+ validTokenNum += tokensPerBlock; // add one more block`
	`217`	`+`
`214`	`218`	`bufferSizeFromMaxNumToken += validTokenNum * kvCacheByteSizePerTokenPerLayer;`
`215`	`219`	`}`
`216`	`220`	`}`
`@@ -237,7 +241,7 @@ CacheTransBufferManager::CacheTransBufferManager(`
`237`	`241`	`allocateBuffer();`
`238`	`242`	`}`
`239`	`243`
`240`		`-size_t CacheTransBufferManager::preAllocBufferSize(`
	`244`	`+size_t CacheTransBufferManager::preAllocBufferSize(size_t tokensPerBlock,`
`241`	`245`	`std::map<SizeType32, SizeType32> const& cacheSizeBytesPerTokenPerWindow,`
`242`	`246`	`std::optional<executor::CacheTransceiverConfig> const& cacheTransceiverConfig)`
`243`	`247`	`{`
`@@ -256,9 +260,9 @@ size_t CacheTransBufferManager::preAllocBufferSize(`
`256`	`260`	`TransferBufferSize = 0;`
`257`	`261`	`for (auto const& [windowSize, cacheSizeBytesPerToken] : cacheSizeBytesPerTokenPerWindow)`
`258`	`262`	`{`
`259`		`- auto validTokenNum`
`260`		`- = (static_cast<size_t>(windowSize) < maxNumTokens.value() ? static_cast<size_t>(windowSize)`
`261`		`- : maxNumTokens.value());`
	`263`	`+ auto alignedWindowSize = (windowSize + tokensPerBlock - 1) / tokensPerBlock * tokensPerBlock;`
	`264`	`+ auto validTokenNum = (alignedWindowSize < maxNumTokens.value() ? alignedWindowSize : maxNumTokens.value());`
	`265`	`+ validTokenNum += tokensPerBlock; // add one more block`
`262`	`266`	`TransferBufferSize += validTokenNum * cacheSizeBytesPerToken;`
`263`	`267`	`}`
`264`	`268`	`}`
Original file line number	Diff line number	Diff line change
`@@ -182,5 +182,6 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m)`
`182`	`182`	`.def(nb::init<tb::kv_cache_manager::BaseKVCacheManager*, std::optional<size_t>>(), nb::arg("cache_manager"),`
`183`	`183`	`nb::arg("max_num_tokens") = std::nullopt)`
`184`	`184`	`.def_static("pre_alloc_buffer_size", &tb::kv_cache_manager::CacheTransBufferManager::preAllocBufferSize,`
`185`		`- nb::arg("cache_size_bytes_per_token_per_window"), nb::arg("cache_transceiver_config") = nb::none());`
	`185`	`+ nb::arg("tokens_per_block"), nb::arg("cache_size_bytes_per_token_per_window"),`
	`186`	`+ nb::arg("cache_transceiver_config") = nb::none());`
`186`	`187`	`}`
Original file line number	Diff line number	Diff line change
`@@ -178,5 +178,6 @@ void tb::CacheTransceiverBindings::initBindings(py::module_& m)`
`178`	`178`	`.def(py::init<tb::kv_cache_manager::BaseKVCacheManager*, std::optional<size_t>>(), py::arg("cache_manager"),`
`179`	`179`	`py::arg("max_num_tokens") = std::nullopt)`
`180`	`180`	`.def_static("pre_alloc_buffer_size", &tb::kv_cache_manager::CacheTransBufferManager::preAllocBufferSize,`
`181`		`- py::arg("cache_size_bytes_per_token_per_window"), py::arg("cache_transceiver_config") = py::none());`
	`181`	`+ py::arg("tokens_per_block"), py::arg("cache_size_bytes_per_token_per_window"),`
	`182`	`+ py::arg("cache_transceiver_config") = py::none());`
`182`	`183`	`}`