fix bug: insert to cache only when stream's kv cache is computed

Author: xinfei-shi

rtp_llm/cpp/engine_base/stream/GenerateStream.cc

@@ -680,6 +680,10 @@ bool GenerateStream::finished() {
     return generate_status_->status == StreamState::FINISHED;
 }
 
+bool GenerateStream::isRemoteRunningWithoutLock() {
+    return generate_status_->status == StreamState::REMOTE_RUNNING;
+}
+
 bool GenerateStream::needRemoteGenerate() const {
     std::lock_guard<std::mutex> lock(*output_mutex_);
     return need_remote_generate_;

rtp_llm/cpp/engine_base/stream/GenerateStream.h

@@ -221,6 +221,7 @@ class GenerateStream {
     bool         finishedWithoutLock();
     void         cancelIfNotRunning();
     void         setFinishedWithoutLock();
+    bool         isRemoteRunningWithoutLock();
     bool         needRemoteGenerate() const;
     bool         setRemoteGenerate();
     size_t       iterCount() const;

rtp_llm/cpp/engine_base/stream/StreamCacheResource.cc

@@ -21,15 +21,15 @@ void StreamCacheResource::freeBatchBlocks(size_t batch_id, vector<int>& blocks)
     bool should_reuse_cache =
         reuseCache() && (!stream_->hasNumBeams() || (!stream_->stoppedWithoutLock() && batch_id == 0));
     // TODO(zhangjianning.zjn) cache all beams of beam search
-    if (blocks.size() == batch_resource_.blockSize(batch_id) && should_reuse_cache) {
+    if (blocks.size() == batch_resource_.blockSize(batch_id) && should_reuse_cache
+        && (stream_->finishedWithoutLock() || stream_->isRemoteRunningWithoutLock())) {
         reConstructCacheKeys();
         auto          tokens_id  = stream_->completeTokenIdsVec(batch_id);
         const auto&   cache_keys = stream_->cacheKeys(batch_id);
         vector<float> loss;
         if (stream_->getLoss()) {
             loss = rtp_llm::buffer2vector<float>(*(stream_->getLoss()));
         }
-        // TODO(xinfei.sxf) 一些场景调用了cancel的地方，是否应该free with cache
         CacheManager::FreeInfo free_info(stream_->streamId(),
                                          tokens_id,
                                          cache_keys,

rtp_llm/cpp/engine_base/stream/test/StreamCacheResourceTest.cc

@@ -307,7 +307,7 @@ TEST_F(StreamCacheResourceTest, testReuseCache) {
     batch_tokens_2[stream_->seqLength() - 3] = 9;
     batch_tokens_2[stream_->seqLength() - 2] = 10;
 
-    // stream_->reConstructCacheKeys();
+    stream_->setFinishedWithoutLock();
     stream_->releaseResource();
 
     ASSERT_EQ(cache_manager_->freeBlockNums(), 3);
@@ -361,6 +361,7 @@ TEST_F(StreamCacheResourceTest, testReuseCache) {
     tokens_2[stream_->seqLength() - 2] = 14;
     tokens_2[stream_->seqLength() - 1] = 15;
 
+    stream_->setFinishedWithoutLock();
     stream_->releaseResource();
     ASSERT_EQ(cache_manager_->freeBlockNums(), 2);
     ASSERT_EQ(cache_manager_->cacheItemNum(), 3);
@@ -377,6 +378,7 @@ TEST_F(StreamCacheResourceTest, testReuseCacheWithFastGen) {
     stream_->setIsContextStream(false);
     ASSERT_TRUE(resource.incrKVBlock(token_capacity).ok());
     ASSERT_EQ(cache_manager_->freeBlockNums(), 3);
+    stream_->setFinishedWithoutLock();
     stream_->releaseResource();
     ASSERT_EQ(cache_manager_->freeBlockNums(), 5);
 
@@ -426,6 +428,7 @@ TEST_F(StreamCacheResourceTest, testReuseCacheWithFastGen) {
     ASSERT_EQ(cache_manager_->freeBlockNums(), 1);
 
     // partial fallback
+    stream_->setFinishedWithoutLock();
     stream_->tryReleaseKVBlock(2);
     stream_->setPaused();
     ASSERT_EQ(stream_->maxBlockSize(), 2);
@@ -440,6 +443,7 @@ TEST_F(StreamCacheResourceTest, testReuseCacheWithFastGen) {
     ASSERT_EQ(stream_->maxBlockSize(), 3);
 
     // full fallback
+    stream_->setFinishedWithoutLock();
     stream_->tryReleaseKVBlock(stream_->maxBlockSize());
     stream_->setPaused();
     ASSERT_EQ(cache_manager_->freeBlockNums(), 3);
@@ -456,6 +460,7 @@ TEST_F(StreamCacheResourceTest, testReuseCacheWithFastGen) {
     ASSERT_EQ(cache_manager_->freeBlockNums(), 1);
     ASSERT_EQ(stream_->maxBlockSize(), 4);
 
+    stream_->setFinishedWithoutLock();
     stream_->releaseResource();
     ASSERT_EQ(cache_manager_->availableBlockNums(), 8);
 }
@@ -472,6 +477,7 @@ TEST_F(StreamCacheResourceTest, testTryReleaseKVBlock) {
     ASSERT_EQ(allocator_->blockRefCounter().getRefCounter(2), 2);
     ASSERT_EQ(allocator_->blockRefCounter().getRefCounter(3), 2);
 
+    stream_->setFinishedWithoutLock();
     resource.tryReleaseKVBlock(1);
     ASSERT_EQ(cache_manager_->freeBlockNums(), 6);
     ASSERT_EQ(resource.maxBlockSize(), 2);
@@ -480,6 +486,7 @@ TEST_F(StreamCacheResourceTest, testTryReleaseKVBlock) {
     ASSERT_EQ(allocator_->blockRefCounter().getRefCounter(2), 2);
     ASSERT_EQ(allocator_->blockRefCounter().getRefCounter(3), 0);
 
+    stream_->setFinishedWithoutLock();
     resource.tryReleaseKVBlock(2);
     ASSERT_EQ(cache_manager_->freeBlockNums(), 8);
     ASSERT_EQ(resource.maxBlockSize(), 0);
@@ -499,6 +506,7 @@ TEST_F(StreamCacheResourceTest, testTryReleaseKVBlock) {
     ASSERT_EQ(allocator_->blockRefCounter().getRefCounter(5), 0);
     ASSERT_EQ(allocator_->blockRefCounter().getRefCounter(6), 0);
 
+    stream_->setFinishedWithoutLock();
     resource.tryReleaseKVBlock(2);
     ASSERT_EQ(cache_manager_->freeBlockNums(), 7);
     ASSERT_EQ(resource.maxBlockSize(), 1);
@@ -509,11 +517,13 @@ TEST_F(StreamCacheResourceTest, testTryReleaseKVBlock) {
     auto tokens_2                          = stream_->complete_token_ids_->data(1);
     tokens_2[0]                            = 2;
 
+    stream_->setFinishedWithoutLock();
     resource.tryReleaseKVBlock(1);
     ASSERT_EQ(cache_manager_->freeBlockNums(), 7);
     ASSERT_EQ(resource.maxBlockSize(), 0);
     ASSERT_EQ(cache_manager_->cacheItemNum(), 2);
 
+    stream_->setFinishedWithoutLock();
     resource.tryReleaseKVBlock(1);
     ASSERT_EQ(cache_manager_->freeBlockNums(), 7);
     ASSERT_EQ(resource.maxBlockSize(), 0);