Method setPrefixCacheFile

transformers/llm/engine/src/llm.cpp:1262–1365

}

bool Llm::setPrefixCacheFile(const std::string& filename, int flag) {
    // Prefix cache requires kvcache_mmap. Without it mKVCacheDir is unset and
    // the disk-backed alloc in CPUKVCacheManager::onAlloc crashes with a
    // nullptr mmap. Reject up-front instead of crashing deep in the backend.
    if (!mConfig->kvcache_mmap()) {
        MNN_ERROR(
            "setPrefixCacheFile requires kvcache_mmap=true in config. "
            "Prefix cache is a disk-backed feature; please enable kvcache_mmap.\n");
        return false;
    }

    mPrefixCacheFileName = filename;
    mCallIndex = 0;
    mPrefixCacheMode = true;


    mIsPrefixFileExist = true;
    // check kvcache, validate file existence
    for(int i = 0; i < mConfig->layer_nums(); i++) {
        auto k_file = MNNFilePathConcat(mConfig->prefix_cache_path(), mPrefixCacheFileName) + "_" + std::to_string(i) + "_sync.k";
        if(!MNNFileExist(k_file.c_str())) {
            mIsPrefixFileExist = false;
            break;
        }
        auto v_file = MNNFilePathConcat(mConfig->prefix_cache_path(), mPrefixCacheFileName) + "_" + std::to_string(i) + "_sync.v";
        if(!MNNFileExist(v_file.c_str())) {
            mIsPrefixFileExist = false;
            break;
        }
    }

    // Further check: the real .k/.v data files must exist, be non-empty, and have consistent
    // sizes across layers. This guards against cases where _sync markers remain but the
    // actual data files were deleted / truncated / partially written, which would otherwise
    // lead to a crash in CPUKVCacheManager::onAlloc (e.g. memset on a null mmap address).
    if (mIsPrefixFileExist) {
        // Only "mix" (LA + regular attention) has intrinsically non-uniform
        // per-layer KV sizes. Everything else (full / sliding / unset → "full")
        // stays under the cross-layer uniformity check that catches partial or
        // corrupted writes.
        const bool isHybrid = (mConfig->attention_type() == "mix");
        size_t refKeySize = 0;
        size_t refValueSize = 0;
        for (int i = 0; i < mConfig->layer_nums(); i++) {
            auto base = MNNFilePathConcat(mConfig->prefix_cache_path(), mPrefixCacheFileName) + "_" + std::to_string(i);
            auto k_data = base + ".k";
            auto v_data = base + ".v";
            if (!MNNFileExist(k_data.c_str()) || !MNNFileExist(v_data.c_str())) {
                MNN_PRINT("Prefix cache data file missing: %s or %s\n", k_data.c_str(), v_data.c_str());
                mIsPrefixFileExist = false;
                break;
            }
            auto kfd = MNNOpenFile(k_data.c_str(), MNN_FILE_READ);
            auto vfd = MNNOpenFile(v_data.c_str(), MNN_FILE_READ);
            size_t kSize = (kfd == INVALID_FILE) ? INVALID_SIZE : MNNGetFileSize(kfd);
            size_t vSize = (vfd == INVALID_FILE) ? INVALID_SIZE : MNNGetFileSize(vfd);
            if (kfd != INVALID_FILE)
                MNNCloseFile(kfd);