[LLM.PrefixDiskCache] Fix prefix disk cache not loaded after first response

source/backend/cpu/CPUKVCacheManager.cpp

@@ -327,13 +327,14 @@ void CPUKVCacheManager::onAlloc(KVMeta* meta, int seq_len) {
             MNN_ERROR("[Error]: Kvcache in disk size smaller than saved lengthInDiskToload:%d\n", (int)meta->seqlen_in_disk);
         }
 
+        // Update mMaxLength first, then setFlashAttentionUpperKv to avoid division by zero
+        int kv_seq_len = meta->add + meta->seqlen_in_disk;
+        mMaxLength = kv_seq_len > oldMaxLength ? kv_seq_len + mConfig.mExpandChunk : oldMaxLength;
         if (mUseFlashAttention) {
             setFlashAttentionUpperKv(MNN_FLASH_ATTENTION_BLOCK_SIZE);
         } else {
             setFlashAttentionUpperKv(mMaxLength);
         }
-        int kv_seq_len = meta->add + meta->seqlen_in_disk;
-        mMaxLength = kv_seq_len > oldMaxLength ? kv_seq_len + mConfig.mExpandChunk : oldMaxLength;
         size_t keySize = (size_t)mKvNumHead * ROUND_UP(mMaxLength, hP) * ROUND_UP(mHeadDim, lP) * mBytes;
         size_t valueSize = (size_t)mKvNumHead * UP_DIV(mMaxLength, mFlashAttentionUpperKv) * (ROUND_UP(mHeadDim, hP) * ROUND_UP(mFlashAttentionUpperKv, lP) * mBytes);
 
@@ -646,19 +647,7 @@ void CPUKVCacheManager::onClear() {
     if (mKVCacheInDisk) {
         // mSaveShareKvPrefix also need unmap file
         unmapKVCache(mCurrentKeySizePerHead * (size_t)mKvNumHead, mCurrentValueSizePerHead * (size_t)mKvNumHead);
-        if(mSaveShareKvPrefix) {
-            // set prefix cachefile validation
-            auto k_file = mBasePrefixFileName + ".k";
-            if(MNNFileExist(k_file.c_str())) {
-                auto k_sync_file = mBasePrefixFileName + "_sync.k";
-                MNNCreateFile(k_sync_file.c_str());
-            }
-            auto v_file = mBasePrefixFileName + ".v";
-            if(MNNFileExist(v_file.c_str())) {
-                auto v_sync_file = mBasePrefixFileName + "_sync.v";
-                MNNCreateFile(v_sync_file.c_str());
-            }
-        } else {
+        if(!mSaveShareKvPrefix) {
             // delete temp kvcache file
             removeKVCacheFile();
         }

source/backend/metal/MetalKVCacheManager.mm

@@ -307,19 +307,7 @@
 
         // mSaveShareKvPrefix also need unmap file
         unmapKVCache(mCurrentTotalSize, mCurrentTotalSize);
-        if(mSaveShareKvPrefix) {
-            // set prefix cachefile validation
-            auto k_file = mBasePrefixFileName + ".k";
-            if(MNNFileExist(k_file.c_str())) {
-                auto k_sync_file = mBasePrefixFileName + "_sync.k";
-                MNNCreateFile(k_sync_file.c_str());
-            }
-            auto v_file = mBasePrefixFileName + ".v";
-            if(MNNFileExist(v_file.c_str())) {
-                auto v_sync_file = mBasePrefixFileName + "_sync.v";
-                MNNCreateFile(v_sync_file.c_str());
-            }
-        } else {
+        if(!mSaveShareKvPrefix) {
             // delete temp kvcache file
             removeKVCacheFile();
         }

transformers/llm/engine/demo/rollback_demo.cpp

@@ -12,6 +12,7 @@
 #include <fstream>
 #include <sstream>
 #include <stdlib.h>
+#include <stdio.h>
 using namespace MNN::Transformer;
 
 
@@ -58,6 +59,43 @@ std::vector<std::vector<std::string>> parse_csv(const std::vector<std::string>&
     return csv_data;
 }
 
+static bool fileExists(const std::string& path) {
+    FILE* f = fopen(path.c_str(), "r");
+    if (f) { fclose(f); return true; }
+    return false;
+}
+
+static void verifySyncFiles(const std::string& prefixDir, const std::string& fileName) {
+    int layerCount = 0;
+    bool allExist = true;
+    for (int i = 0; i < 256; i++) {
+        std::string base = prefixDir + "/" + fileName + "_" + std::to_string(i);
+        std::string k_file = base + ".k";
+        std::string v_file = base + ".v";
+        if (!fileExists(k_file) && !fileExists(v_file)) {
+            break;
+        }
+        std::string k_sync = base + "_sync.k";
+        std::string v_sync = base + "_sync.v";
+        if (!fileExists(k_sync)) {
+            MNN_PRINT("[TEST FAIL] Missing: %s\n", k_sync.c_str());
+            allExist = false;
+        }
+        if (!fileExists(v_sync)) {
+            MNN_PRINT("[TEST FAIL] Missing: %s\n", v_sync.c_str());
+            allExist = false;
+        }
+        layerCount++;
+    }
+    if (allExist && layerCount > 0) {
+        MNN_PRINT("[TEST PASS] All %d layers sync files verified.\n", layerCount);
+    } else if (layerCount == 0) {
+        MNN_PRINT("[TEST FAIL] No KV cache files found in %s/\n", prefixDir.c_str());
+    } else {
+        MNN_PRINT("[TEST FAIL] %d layers checked, some sync files missing.\n", layerCount);
+    }
+}
+
 static int benchmark(Llm* llm, const std::vector<std::string>& prompts, int max_token_number, bool is_prompt_cache) {
     if (prompts.size() < 3) {
         MNN_ERROR("Need larger than 3 inputs\n");
@@ -82,6 +120,9 @@ static int benchmark(Llm* llm, const std::vector<std::string>& prompts, int max_
         // step 2: prefill prefix prompt
         llm->response(prompt_base, &std::cout, nullptr, 0);
 
+        // Verify: sync files should exist after first response (completePrefixWrite)
+        verifySyncFiles("prefixcache", "model_prompt_config_mnnversion");
+
 
         auto prompt_len   = context->prompt_len;
         auto decode_len   = context->gen_seq_len;

transformers/llm/engine/include/llm/llm.hpp

@@ -233,6 +233,7 @@ class MNN_PUBLIC Llm {
     int mCallIndex;
     int mPrefixLength;
     bool mIsPrefixFileExist = false;
+    void completePrefixWrite();
 };
 
 // Embedding start

transformers/llm/engine/src/llm.cpp

@@ -855,7 +855,9 @@ std::vector<int> Llm::generate(const std::vector<int>& input_ids, int max_tokens
             if(hidden_states == nullptr) {
                 return {};
             }
-            return generate(hidden_states, max_tokens);
+            auto result = generate(hidden_states, max_tokens);
+            completePrefixWrite();
+            return result;
         }
         int total_size = (int)input_ids.size();
         int loop_size = UP_DIV(total_size, mBlockSize);
@@ -872,6 +874,7 @@ std::vector<int> Llm::generate(const std::vector<int>& input_ids, int max_tokens
             }
             generate(input_embeds, 0);
         }
+        completePrefixWrite();
     } else {
         // update states
         updateContext((int)input_ids.size(), 0);
@@ -1081,6 +1084,32 @@ bool Llm::setPrefixCacheFile(const std::string& filename, int flag) {
     return mIsPrefixFileExist;
 }
 
+void Llm::completePrefixWrite() {
+    if (!mPrefixCacheMode || mCallIndex != 1 || mIsPrefixFileExist) {
+        return;
+    }
+    mMeta->file_flag = KVMeta::NoChange;
+    mMeta->file_name = "";
+    mMeta->layer_index = 0;
+    // Create sync files to mark prefix cache as valid
+    auto prefixDir = mConfig->prefix_cache_path();
+    for (int i = 0; i < mConfig->layer_nums(); i++) {
+        auto base = MNNFilePathConcat(prefixDir, mPrefixCacheFileName) + "_" + std::to_string(i);
+        auto k_file = base + ".k";
+        if (MNNFileExist(k_file.c_str())) {
+            auto k_sync = base + "_sync.k";
+            auto fd = MNNCreateFile(k_sync.c_str());
+            if (fd != INVALID_FILE) { MNNCloseFile(fd); }
+        }
+        auto v_file = base + ".v";
+        if (MNNFileExist(v_file.c_str())) {
+            auto v_sync = base + "_sync.v";
+            auto fd = MNNCreateFile(v_sync.c_str());
+            if (fd != INVALID_FILE) { MNNCloseFile(fd); }
+        }
+    }
+}
+
 bool Llm::reuse_kv() { return mConfig->reuse_kv(); }
 
 static inline bool needNewVar(VARP var, int axis, int seq_len, int kv_seq_len = 0) {