Release recv buffers on transfer exceptions

Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com>

Author: liji-nv

cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp

@@ -1093,17 +1093,25 @@ class CacheReceiver::Impl
         TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId));
         auto session = sendRequestInfo(llmRequest);
         session.setTime(TransferSession::kTimeRequestInfo);
-        bool isReady = receiveReadySignal(session);
-        if (!isReady)
+        try
+        {
+            bool isReady = receiveReadySignal(session);
+            if (!isReady)
+            {
+                // Reuse the error state for the cancelled request.
+                llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR);
+                llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now());
+                session.releasePreAssignedRecvBuffers();
+                return;
+            }
+            receiveSync(session);
+            session.clearPreAssignedRecvBuffers();
+        }
+        catch (...)
         {
-            // Reuse the error state for the cancelled request.
-            llmRequest.setState(LlmRequestState::kDISAGG_TRANS_ERROR);
-            llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now());
             session.releasePreAssignedRecvBuffers();
-            return;
+            throw;
         }
-        receiveSync(session);
-        session.clearPreAssignedRecvBuffers();
         llmRequest.setKvCacheTransferEnd(std::chrono::steady_clock::now());
 
         TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(),