fix cache buffer (#3942)

Signed-off-by: Chuang Zhu <[email protected]>

Author: chuangz0

cpp/tensorrt_llm/batch_manager/cacheFormatter.cpp

@@ -421,7 +421,6 @@ void CacheFormatter::formatInput(LlmRequest const& llmRequest,
                 else
                 {
                     cacheBufferId = mCacheTransBufferManager->assignBufferIndexForRecv();
-                    TLLM_CHECK(cacheBufferId.has_value());
                     auto [recvSplitCachestmp, bufferCoverTargetNumtmp, onlyUseDynamicBuffer]
                         = mCacheTransBufferManager->getOrAllocateRecvBuffers(
                             cacheBufferId, targetNum, targetBufferSize, bufferManager);

cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp

@@ -38,22 +38,21 @@ CacheTransBufferManager::CacheTransBufferManager(
     {
         TLLM_CHECK(maxNumTokens.value() % tokensPerBlock == 0);
     }
-    TLLM_LOG_INFO("maxNumTokens: %d", maxNumTokens.has_value() ? maxNumTokens.value() : 0);
     auto kvCachePerToken
         = (mCacheManager->getBlockManager().getBlockSize(0) * mCacheManager->getBlockManager().getNumLayers()
               * (mCacheManager->getCacheType() == CacheType::kSELFKONLY ? 1 : 2))
         / tokensPerBlock;
     mTransferBufferSize = maxNumTokens.has_value() ? maxNumTokens.value() * kvCachePerToken
                                                    : common::getEnvMemSizeForKVCacheTransferBuffer();
-    monlyUseDynamicBuffer = mTransferBufferSize == 0;
+    mOnlyUseDynamicBuffer = mTransferBufferSize == 0;
     mRecvBufferCount = common::getEnvRequestKVCacheConcurrent() ? common::getEnvKVCacheRecvBufferCount() : 1;
     mSendBufferCount = common::getEnvParallelCacheSend() ? common::getEnvKVCacheSendMaxConcurrenceNum() : 1;
     mPreAllocBufferSize = mTransferBufferSize * (mRecvBufferCount + mSendBufferCount);
     TLLM_LOG_INFO(
         "CacheTransBufferManager: mMaxNumTokens:%ld, mRecvBufferCount:%ld, "
-        "mSendBufferCount:%ld,mTransferBufferSize:%ld, mPreAllocBufferSize:%ld",
+        "mSendBufferCount:%ld,mTransferBufferSize:%ld, mPreAllocBufferSize:%ld,monlyUseDynamicBuffer:%d",
         maxNumTokens.has_value() ? maxNumTokens.value() : 0, mRecvBufferCount, mSendBufferCount, mTransferBufferSize,
-        mPreAllocBufferSize);
+        mPreAllocBufferSize, mOnlyUseDynamicBuffer);
     bool to_allocate = common::getEnvUseMPIKvCache() || common::getEnvUseUCXKvCache();
 
     TLLM_CHECK_WITH_INFO(to_allocate, "CacheTransBufferManager: to_allocate is false");
@@ -85,22 +84,22 @@ size_t CacheTransBufferManager::preAllocBufferSize(
 
 std::optional<int> CacheTransBufferManager::assignBufferIndexForSend()
 {
-    return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, monlyUseDynamicBuffer);
+    return assignBufferIndex(mConcurrenceSendResource, mSendBufferCount, mOnlyUseDynamicBuffer);
 }
 
 void CacheTransBufferManager::freeBufferIndexForSend(std::optional<int> bufferId)
 {
-    freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, monlyUseDynamicBuffer);
+    freeBufferIndex(mConcurrenceSendResource, bufferId, mSendBufferCount, mOnlyUseDynamicBuffer);
 }
 
 std::optional<int> CacheTransBufferManager::assignBufferIndexForRecv()
 {
-    return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, monlyUseDynamicBuffer);
+    return assignBufferIndex(mConcurrenceRecvResource, mRecvBufferCount, mOnlyUseDynamicBuffer);
 }
 
 void CacheTransBufferManager::freeBufferIndexForRecv(std::optional<int> bufferId)
 {
-    freeBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, monlyUseDynamicBuffer);
+    freeBufferIndex(mConcurrenceRecvResource, bufferId, mRecvBufferCount, mOnlyUseDynamicBuffer);
 }
 
 std::tuple<std::vector<runtime::ITensor::SharedPtr>, size_t, bool> CacheTransBufferManager::getOrAllocateSendBuffers(
@@ -119,7 +118,7 @@ std::tuple<std::vector<runtime::ITensor::SharedPtr>, size_t, bool> CacheTransBuf
 
 runtime::ITensor::SharedPtr CacheTransBufferManager::getSendBuffer(std::optional<int> bufferId)
 {
-    TLLM_CHECK(bufferId.has_value() || monlyUseDynamicBuffer);
+    TLLM_CHECK(bufferId.has_value() || mOnlyUseDynamicBuffer);
     if (bufferId.has_value())
     {
         TLLM_CHECK(static_cast<size_t>(bufferId.value()) < mSendBufferCount);
@@ -131,7 +130,7 @@ runtime::ITensor::SharedPtr CacheTransBufferManager::getSendBuffer(std::optional
 
 runtime::ITensor::SharedPtr CacheTransBufferManager::getRecvBuffer(std::optional<int> bufferId)
 {
-    TLLM_CHECK(bufferId.has_value() || monlyUseDynamicBuffer);
+    TLLM_CHECK(bufferId.has_value() || mOnlyUseDynamicBuffer);
     if (bufferId.has_value())
     {
         TLLM_CHECK(static_cast<size_t>(bufferId.value()) < mRecvBufferCount);
@@ -145,7 +144,7 @@ std::tuple<std::vector<runtime::ITensor::SharedPtr>, size_t, bool> CacheTransBuf
     std::optional<int> bufferId, int targetNum, size_t targetBufferEleSize,
     runtime::BufferManager const& bufferManagerToUse, ConcurrenceResource& concurrenceResource)
 {
-    TLLM_CHECK(bufferId.has_value() || monlyUseDynamicBuffer);
+    TLLM_CHECK(bufferId.has_value() || mOnlyUseDynamicBuffer);
     std::vector<runtime::ITensor::SharedPtr> retSplitCaches;
     size_t bufferCoverTargetNum = std::min(
         static_cast<size_t>(targetNum), mTransferBufferSize / (targetBufferEleSize * common::getDTypeSize(mDataType)));
@@ -178,18 +177,17 @@ std::tuple<std::vector<runtime::ITensor::SharedPtr>, size_t, bool> CacheTransBuf
                 runtime::ITensor::makeShape({static_cast<int64_t>(targetBufferEleSize)}), mDataType));
         }
     }
-    if (monlyUseDynamicBuffer)
+    if (mOnlyUseDynamicBuffer)
     {
         bufferCoverTargetNum = targetNum;
     }
-    return std::make_tuple(retSplitCaches, bufferCoverTargetNum, monlyUseDynamicBuffer);
+    return std::make_tuple(retSplitCaches, bufferCoverTargetNum, mOnlyUseDynamicBuffer);
 }
 
 void CacheTransBufferManager::allocateBuffer()
 {
-    if (monlyUseDynamicBuffer)
+    if (mOnlyUseDynamicBuffer)
     {
-        TLLM_LOG_INFO("monlyUseDynamicBuffer: true");
         return;
     }
     mBufferEleSize = mTransferBufferSize / common::getDTypeSize(mDataType);

cpp/tensorrt_llm/batch_manager/cacheTransBuffer.h

@@ -78,7 +78,7 @@ class CacheTransBufferManager
     size_t mRecvBufferCount;
     size_t mSendBufferCount;
     size_t mTransferBufferSize;
-    bool monlyUseDynamicBuffer;
+    bool mOnlyUseDynamicBuffer;
     size_t mBufferEleSize;
     nvinfer1::DataType mDataType;
     ConcurrenceResource mConcurrenceSendResource;

cpp/tensorrt_llm/executor/cache_transmission/ucx_utils/ucxCacheCommunicator.cpp

@@ -106,8 +106,7 @@ UcxConnectionManager::UcxConnectionManager()
     try
     {
         TLLM_CUDA_CHECK(cudaGetDevice(&mDevice));
-        mUcxCtx = ucxx::createContext(
-            {{"RNDV_PIPELINE_ERROR_HANDLING", "y"}, {"MEMTYPE_CACHE", "n"}}, ucxx::Context::defaultFeatureFlags);
+        mUcxCtx = ucxx::createContext({{"RNDV_PIPELINE_ERROR_HANDLING", "y"}}, ucxx::Context::defaultFeatureFlags);
         int device = mDevice;
         try
         {

docs/source/advanced/disaggregated-service.md

@@ -87,7 +87,7 @@ TRT-LLM uses some environment variables to control the behavior of disaggregated
 
 * `TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER`: TRT-LLM typically copies non-contiguous data into a temporary buffer before sending KV cache. If set to `1`, TRT-LLM will attempt to directly transmit each KV cache block, eliminating extra copies. The default value is `0`.
 
-* `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE`: By default, TRT-LLM uses a `stream-ordered memory allocator` to allocate temporary buffers. If this environment variable is set to #Size, TRT-LLM will use `cudaMalloc` to allocate buffer of size #Size for KV cache transmission. The default value is `0`. Users can set `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=1GB` to allocate a 1 GB buffer with `cudaMalloc` for KV cache transmission.
+* `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE`: By default, TRT-LLM uses a `stream-ordered memory allocator` to allocate temporary buffers. If this environment variable is set to #Size, TRT-LLM will use `cudaMalloc` to allocate buffer of size #Size for KV cache transmission. The default value is `512MB`. Users can set `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=1GB` to allocate a 1 GB buffer with `cudaMalloc` for KV cache transmission.
 
 * `TRTLLM_KVCACHE_TRANSFER_USE_ASYNC_BUFFER`: If set to `1`, TRT-LLM will use `cudaMallocAsync` to allocate buffers for KV cache transmission. The default value is `0`. This environment variable only takes effect when `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE` is greater than 0.
 
@@ -146,13 +146,13 @@ When the environment variable `TRTLLM_USE_MPI_KVCACHE=1` is set, TRT-LLM will tr
 A. Ensure TRT-LLM is running with `UCX`-backend `CUDA-aware MPI` , and check version of `UCX` with `ucx_info -v`.
 If the version of UCX <=1.17, set the environment variables `UCX_RNDV_FRAG_MEM_TYPE=cuda` and `UCX_MEMTYPE_CACHE=n` to enable NVLink. For BlackWell architecture GPUs, UCX version >=1.19 is required to enable NVLink.
 If the version of UCX >=1.18, there are several ways to enable NVLink:
-1. Set the environment variables `UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda`, `UCX_CUDA_COPY_DMABUF=no`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`.
+1. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B`,`UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda`, `UCX_CUDA_COPY_DMABUF=no`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`.
 2. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$Size`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`. $Size represents the size of the buffer for KV cache transfer, which is recommended to be larger than the size of the KV cache for the longest request.
 
 *Q. Does TRT-LLM support using GPU direct RDMA for inter-node KV Cache transfer?*
 
 A. Yes, TRT-LLM supports using GPU direct RDMA for inter-node KV cache transfer, but it is not enabled by default. There are several ways to enable GPU direct RDMA:
-1. Set the environment variables `UCX_RNDV_FRAG_MEM_TYPE=cuda`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`.
+1. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B`,`UCX_RNDV_FRAG_MEM_TYPE=cuda`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`.
 2. Set the environment variables `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=$Size`, `UCX_MEMTYPE_CACHE=n` and `UCX_RNDV_PIPELINE_ERROR_HANDLING=y`, $Size represents the size of the buffer for KV cache transfer, which is recommended to be larger than the size of the KV cache for the longest request.
 To achieve the optimal performance when using GPU direct RDMA, it is advisable to create CUDA context before MPI initialization when TRTLLM_USE_MPI_KVCACHE=1 is set. One possible approach is to rely on MPI environment variables to set the correct device before MPI initialization.
 
@@ -163,6 +163,7 @@ A. Depending on the user's use case, certain sets of environment variables can h
 Environment Variable Set A
 
 ```
+export TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B
 export UCX_RNDV_FRAG_MEM_TYPES=cuda
 export UCX_MEMTYPE_CACHE=n
 export UCX_RNDV_PIPELINE_ERROR_HANDLING=y
@@ -172,6 +173,7 @@ This set allows KV cache transfers to utilize NVLink within nodes and GDRDMA bet
 Environment Variable Set B
 
 ```
+export TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=0B
 export UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda
 export UCX_CUDA_COPY_DMABUF=no
 export UCX_MEMTYPE_CACHE=n

examples/disaggregated/README.md

@@ -9,14 +9,15 @@ You can use multiple `trtllm-serve` commands to launch the context and generatio
 for disaggregated serving. For example, you could launch two context servers and one generation servers as follows:
 
 ```
-echo -e "pytorch_backend_config:\n  enable_overlap_scheduler: False" > extra-llm-api-config.yml
+echo -e "pytorch_backend_config:\n  enable_overlap_scheduler: False\ncache_transceiver_config:\n  max_num_tokens: 2048" > context_extra-llm-api-config.yml
+echo -e "cache_transceiver_config:\n  max_num_tokens: 2048" > gen_extra-llm-api-config.yml
 
 export TRTLLM_USE_UCX_KVCACHE=1
 #Context servers
-CUDA_VISIBLE_DEVICES=0 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --backend pytorch --extra_llm_api_options ./extra-llm-api-config.yml &> log_ctx_0 &
-CUDA_VISIBLE_DEVICES=1 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --backend pytorch --extra_llm_api_options ./extra-llm-api-config.yml &> log_ctx_1 &
+CUDA_VISIBLE_DEVICES=0 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --backend pytorch --extra_llm_api_options ./context_extra-llm-api-config.yml &> log_ctx_0 &
+CUDA_VISIBLE_DEVICES=1 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --backend pytorch --extra_llm_api_options ./context_extra-llm-api-config.yml &> log_ctx_1 &
 #Generation servers
-CUDA_VISIBLE_DEVICES=2 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8003 --backend pytorch &> log_gen_0 &
+CUDA_VISIBLE_DEVICES=2 trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8003 --backend pytorch --extra_llm_api_options ./gen_extra-llm-api-config.yml &> log_gen_0 &
 ```
 Once the context and generation servers are launched, you can launch the disaggregated
 server, which will accept requests from clients and do the orchestration between context