//
//  CPUKVCacheManager.cpp
//  MNN
//
//  Created by MNN on 2024/08/05.
//  Copyright © 2018, Alibaba Group Holding Limited
//

#ifdef MNN_SUPPORT_TRANSFORMER_FUSE

#include "CPUKVCacheManager.hpp"
#include "core/Concurrency.h"
#include "core/TensorUtils.hpp"

namespace MNN {

static inline int c4Offset(int token, int channel, int seqLen, int pack) {
    return (channel / pack) * seqLen * pack + token * pack + (channel % pack);
}

/*
**  @brief  Expand the size of kvcache and copy it from the old tensor in memory to the new tensor in memory
**          Finally reset the pointer to the new tensor
*/
void CPUKVCacheManager::expandKVCacheInMem(int oldMaxLength) {
    /*===================================  Key  ===================================*/
    auto new_key = Tensor::createDevice<int8_t>({mKvNumHead, (int)mCurrentKeySizePerHead});
    mBackend->onAcquireBuffer(new_key, Backend::STATIC);
    if (mKeyQuantMode != KVQuantMode::None) {
        memset(new_key->host<int8_t>(), 0, mKvNumHead * mCurrentKeySizePerHead);
    }
    for (int h = 0; h < mKvNumHead; h++) {
        memcpy(new_key->host<int8_t>() + h * mCurrentKeySizePerHead, mPastKey->host<int8_t>() + h * mPastKey->stride(0),
               mPastKey->stride(0));
        if (mKeyQuantMode != KVQuantMode::Int8 && (new_key->stride(0) - mPastKey->stride(0)) > 0) {
            memset(new_key->host<int8_t>() + h * new_key->stride(0) + mPastKey->stride(0), 0,
                   (new_key->stride(0) - mPastKey->stride(0)));
        }
    }
    mPastKey.reset(new_key);
    /*===================================  Value  ===================================*/
    auto newValue = Tensor::createDevice<int8_t>({mKvNumHead, (int)mCurrentValueSizePerHead});
    mBackend->onAcquireBuffer(newValue, Backend::STATIC);

    if (mUseFlashAttention) { // [mKvNumHead, UP_DIV(mMaxLength, mFlashAttentionUpperKv), UP_DIV(mHeadDim, hP),
                              // UP_DIV(mFlashAttentionUpperKv, lP), hP, lP]
        for (int h = 0; h < mKvNumHead; h++) {
            memset(newValue->host<int8_t>() + h * newValue->stride(0), 0, newValue->stride(0));
            memcpy(newValue->host<int8_t>() + h * newValue->stride(0),
                   mPastValue->host<int8_t>() + h * mPastValue->stride(0), mPastValue->stride(0));
        }
    } else {
        if (mValueQuantMode == KVQuantMode::Int8) { // [mKvNumHead, UP_DIV(mHeadDim, hP8), (UP_DIV(mMaxLength,
                                                    // lP8)*hP8*lP8+2*hP8*sizeof(float)) ]
            auto currentWeightInside = ROUND_UP(mMaxLength, lP8) * hP8;
            auto currentStride1 = currentWeightInside + 2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES;
            auto currentStride0 = currentStride1 * UP_DIV(mHeadDim, hP8);

            auto prevWeightInside = ROUND_UP(oldMaxLength, lP8) * hP8;
            auto prevStride1 = prevWeightInside + 2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES;
            auto prevStride0 = prevStride1 * UP_DIV(mHeadDim, hP8);
            for (int h = 0; h < mKvNumHead; ++h) {
                for (int d = 0; d < UP_DIV(mHeadDim, hP8); ++d) {
                    auto dstPtr = newValue->host<int8_t>() + h * currentStride0 + d * currentStride1;
                    auto srcPtr = mPastValue->host<int8_t>() + h * prevStride0 + d * prevStride1;

                    // initialize 0 for weightInt8
                    memset(dstPtr, 0, currentWeightInside);
                    // copy inner side weightInt8
                    memcpy(dstPtr, srcPtr, prevWeightInside);
                    // copy hP8 scale&bias
                    memcpy(dstPtr + currentWeightInside, srcPtr + prevWeightInside,
                           2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES);
                }
            }
        } else { // [mKvNumHead, UP_DIV(mHeadDim, hP), UP_DIV(mMaxLength, lP), hP, lP]
            auto currentStride1 = ROUND_UP(mMaxLength, lP) * hP * mBytes;
            auto currentStride0 = ROUND_UP(mMaxLength, lP) * hP * UP_DIV(mHeadDim, hP) * mBytes;

            auto prevStride1 = ROUND_UP(oldMaxLength, lP) * hP * mBytes;
            auto prevStride0 = ROUND_UP(oldMaxLength, lP) * hP * UP_DIV(mHeadDim, hP) * mBytes;
            for (int h = 0; h < mKvNumHead; ++h) {
                for (int d = 0; d < UP_DIV(mHeadDim, hP); ++d) {
                    auto dstPtr = newValue->host<int8_t>() + h * currentStride0 + d * currentStride1;
                    auto srcPtr = mPastValue->host<int8_t>() + h * prevStride0 + d * prevStride1;

                    // initialize 0 for weight
                    if (lP > 1) {
                        memset(dstPtr, 0, currentStride1);
                    }
                    // copy inner side weight
                    memcpy(dstPtr, srcPtr, prevStride1);
                }
            }
        }
    }
    mPastValue.reset(newValue);
}

/*
**  @brief  Move the kvcache from memory to the memory-mapped kvcache files in disk
**          Then release the memory buffer of old kvcache
*/
void CPUKVCacheManager::moveKVCacheFromMemToDisk(int oldMaxLength) {
    /*===================================  Key  ===================================*/
    size_t prevKeySizePerHead = 0;
    if (mKeyQuantMode == KVQuantMode::Int8) {
        prevKeySizePerHead = ROUND_UP(oldMaxLength, hP8) * ROUND_UP(mHeadDim, lP8) +
                             2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(oldMaxLength, hP8);
    } else {
        prevKeySizePerHead = UP_DIV(oldMaxLength, hP) * ROUND_UP(mHeadDim, lP) * hP * mBytes;
    }
    if (mHeadDim % lP || (mKeyQuantMode == KVQuantMode::Int8)) {
        memset(mMapKeyAddr, 0, mKvNumHead * mCurrentKeySizePerHead);
    }
    for (int h = 0; h < mKvNumHead; h++) {
        memcpy(mMapKeyAddr + h * mCurrentKeySizePerHead, mPastKey->host<int8_t>() + h * prevKeySizePerHead,
               prevKeySizePerHead);
    }
    mBackend->onReleaseBuffer(mPastKey.get(), Backend::STATIC);
    mPastKey.reset();
    /*===================================  Value  ===================================*/
    {
        size_t prevValueSizePerHead = 0;
        if (mValueQuantMode == KVQuantMode::Int8) {
            prevValueSizePerHead = UP_DIV(oldMaxLength, mFlashAttentionUpperKv) *
                                   (ROUND_UP(mHeadDim, hP8) * ROUND_UP(mFlashAttentionUpperKv, lP8) +
                                    2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(mHeadDim, hP8));
        } else {
            prevValueSizePerHead = UP_DIV(oldMaxLength, mFlashAttentionUpperKv) *
                                   (ROUND_UP(mHeadDim, hP) * ROUND_UP(mFlashAttentionUpperKv, lP) * mBytes);
        }
        if (lP > 1 || (mValueQuantMode == KVQuantMode::Int8)) {
            memset(mMapValueAddr, 0, mKvNumHead * mCurrentValueSizePerHead);
        }

        if (mUseFlashAttention) {
            for (int h = 0; h < mKvNumHead; h++) {
                memcpy(mMapValueAddr + h * mCurrentValueSizePerHead,
                       mPastValue->host<int8_t>() + h * prevValueSizePerHead, prevValueSizePerHead);
            }
        } else {
            if (mValueQuantMode == KVQuantMode::Int8) { // [mKvNumHead, UP_DIV(mHeadDim, hP8), (UP_DIV(mMaxLength,
                                                        // lP8)*hP8*lP8+2*hP8*sizeof(float)) ]
                auto currentWeightInside = ROUND_UP(mMaxLength, lP8) * hP8;
                auto currentStride1 = currentWeightInside + 2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES;
                auto currentStride0 = currentStride1 * UP_DIV(mHeadDim, hP8);

                auto prevWeightInside = ROUND_UP(oldMaxLength, lP8) * hP8;
                auto prevStride1 = prevWeightInside + 2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES;
                auto prevStride0 = prevStride1 * UP_DIV(mHeadDim, hP8);
                for (int h = 0; h < mKvNumHead; ++h) {
                    for (int d = 0; d < UP_DIV(mHeadDim, hP8); ++d) {
                        auto dstPtr = mMapValueAddr + h * currentStride0 + d * currentStride1;
                        auto srcPtr = mPastValue->host<int8_t>() + h * prevStride0 + d * prevStride1;

                        // initialize 0 for weightInt8
                        memset(dstPtr, 0, currentWeightInside);
                        // copy inner side weightInt8
                        memcpy(dstPtr, srcPtr, prevWeightInside);
                        // copy hP8 scale&bias
                        memcpy(dstPtr + currentWeightInside, srcPtr + prevWeightInside,
                               2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES);
                    }
                }
            } else { // [mKvNumHead, UP_DIV(mHeadDim, hP), UP_DIV(mMaxLength, lP), hP, lP]
                auto currentStride1 = ROUND_UP(mMaxLength, lP) * hP * mBytes;
                auto currentStride0 = ROUND_UP(mMaxLength, lP) * hP * UP_DIV(mHeadDim, hP) * mBytes;

                auto prevStride1 = ROUND_UP(oldMaxLength, lP) * hP * mBytes;
                auto prevStride0 = ROUND_UP(oldMaxLength, lP) * hP * UP_DIV(mHeadDim, hP) * mBytes;
                for (int h = 0; h < mKvNumHead; ++h) {
                    for (int d = 0; d < UP_DIV(mHeadDim, hP); ++d) {
                        auto dstPtr = mMapValueAddr + h * currentStride0 + d * currentStride1;
                        auto srcPtr = mPastValue->host<int8_t>() + h * prevStride0 + d * prevStride1;

                        // initialize 0 for weight
                        if (lP > 1) {
                            memset(dstPtr, 0, currentStride1);
                        }
                        // copy inner side weight
                        memcpy(dstPtr, srcPtr, prevStride1);
                    }
                }
            }
        }
        mBackend->onReleaseBuffer(mPastValue.get(), Backend::STATIC);
        mPastValue.reset();
    }
}

/*
**  @brief  Expand the size of kvcache files in disk
*/
void CPUKVCacheManager::expandKVCacheInDisk(size_t oldMaxLength, size_t oldKeySize, size_t oldValueSize, size_t keySize,
                                            size_t valueSize, file_t specKeyFile, file_t specValueFile) {
    // Step 1: Copy the old kvcache from files to temporary buffers in memory
    auto prevKeySizePerHead = oldKeySize / mKvNumHead;
    auto prevValueSizePerHead = oldValueSize / mKvNumHead;
    std::shared_ptr<Tensor> prevKey, prevValue;
    // NOTE: Tensor shape is std::vector<int>. prevKeySizePerHead is size_t after
    // the int -> size_t fix; cast back to int here. In realistic models a single
    // head's cache is far below 2 GB, so the narrowing is safe in practice.
    prevKey.reset(Tensor::createDevice<int8_t>({mKvNumHead, (int)prevKeySizePerHead}));
    prevValue.reset(Tensor::createDevice<int8_t>({mKvNumHead, (int)prevValueSizePerHead}));

    mBackend->onAcquireBuffer(prevKey.get(), Backend::STATIC);
    mBackend->onAcquireBuffer(prevValue.get(), Backend::STATIC);
    if (mHeadDim % lP) {
        memset(prevKey->host<uint8_t>(), 0, prevKey->length(0) * prevKey->stride(0));
    }
    if (lP > 1) {
        // can't be mMaxLenth % lP, since mMaxLength may be larger than seq_len for prefilling, we should ensure the
        // (mMaxLength - seq_len)'s buffer is 0. computing L is seq_len
        memset(prevValue->host<uint8_t>(), 0, prevValue->length(0) * prevValue->stride(0));
    }
    mmapKVCache(oldKeySize, oldValueSize, specKeyFile, specValueFile);
    memcpy(prevKey->host<int8_t>(), mMapKeyAddr, oldKeySize);
    memcpy(prevValue->host<int8_t>(), mMapValueAddr, oldValueSize);
    // Step 2: Resize the kvcache files and remap them
    unmapKVCache(oldKeySize, oldValueSize);
    resetKVCacheFileSize(keySize, valueSize);
    mmapKVCache(keySize, valueSize);
    // Step 3: Move the kvcache from temporary buffers in memory to disk
    memset(mMapKeyAddr, 0, keySize);
    memset(mMapValueAddr, 0, valueSize);

    for (int h = 0; h < mKvNumHead; h++) {
        memcpy(mMapKeyAddr + h * mCurrentKeySizePerHead, prevKey->host<int8_t>() + h * prevKeySizePerHead,
               prevKeySizePerHead);
    }

    if (mUseFlashAttention) {
        for (int h = 0; h < mKvNumHead; h++) {
            memcpy(mMapValueAddr + h * mCurrentValueSizePerHead, prevValue->host<int8_t>() + h * prevValueSizePerHead,
                   prevValueSizePerHead);
        }
    } else {
        if (mValueQuantMode == KVQuantMode::Int8) {
            auto currentWeightInside = ROUND_UP(mMaxLength, lP8) * hP8;
            auto currentStride1 = currentWeightInside + 2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES;
            auto currentStride0 = currentStride1 * UP_DIV(mHeadDim, hP8);

            auto prevWeightInside = ROUND_UP(oldMaxLength, lP8) * hP8;
            auto prevStride1 = prevWeightInside + 2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES;
            auto prevStride0 = prevStride1 * UP_DIV(mHeadDim, hP8);

            for (int h = 0; h < mKvNumHead; ++h) {
                for (int d = 0; d < UP_DIV(mHeadDim, hP8); ++d) {
                    auto dstPtr = mMapValueAddr + h * currentStride0 + d * currentStride1;
                    auto srcPtr = prevValue->host<int8_t>() + h * prevStride0 + d * prevStride1;

                    // initialize 0 for weightInt8
                    memset(dstPtr, 0, currentWeightInside);
                    // copy inner side weightInt8
                    memcpy(dstPtr, srcPtr, prevWeightInside);
                    // copy hP8 scale&bias
                    memcpy(dstPtr + currentWeightInside, srcPtr + prevWeightInside,
                           2 * mConfig.mBlockNum * hP8 * QUANT_INFO_BYTES);
                }
            }
        } else {
            auto currentStride1 = ROUND_UP(mMaxLength, lP) * hP * mBytes;
            auto currentStride0 = ROUND_UP(mMaxLength, lP) * hP * UP_DIV(mHeadDim, hP) * mBytes;

            auto prevStride1 = ROUND_UP(oldMaxLength, lP) * hP * mBytes;
            auto prevStride0 = ROUND_UP(oldMaxLength, lP) * hP * UP_DIV(mHeadDim, hP) * mBytes;
            for (int h = 0; h < mKvNumHead; ++h) {
                for (int d = 0; d < UP_DIV(mHeadDim, hP); ++d) {
                    auto dstPtr = mMapValueAddr + h * currentStride0 + d * currentStride1;
                    auto srcPtr = prevValue->host<int8_t>() + h * prevStride0 + d * prevStride1;

                    // initialize 0 for weight
                    if (lP > 1) {
                        memset(dstPtr, 0, currentStride1);
                    }
                    // copy inner side weight
                    memcpy(dstPtr, srcPtr, prevStride1);
                }
            }
        }
    }

    // Step 4: Release the temporary buffers
    mBackend->onReleaseBuffer(prevKey.get(), Backend::STATIC);
    mBackend->onReleaseBuffer(prevValue.get(), Backend::STATIC);
}

void CPUKVCacheManager::onResize(int kv_num_head, int head_dim) {
    mKvNumHead = kv_num_head;
    mHeadDim = head_dim;
    auto core = static_cast<CPUBackend*>(mBackend)->functions();
    core->MNNGetMatMulPackMode(&eP, &lP, &hP);
    mBytes = core->bytes;
    mThreadNum = static_cast<CPUBackend*>(mBackend)->threadNumber();
    if (mThreadNum > mKvNumHead) {
        mThreadNum = mKvNumHead;
    }

    static_cast<CPUBackend*>(mBackend)->int8Functions()->MNNGetGemmUnit(&hP8, &lP8, &eP8);
    mQuantKeyFunc = core->MNNQuantAttentionKey;
    mQuantValueFunc = core->MNNQuantAttentionValue;
}

void CPUKVCacheManager::onAlloc(KVMeta* meta, int seq_len) {
    mMeta = meta;

    // load disk prefix kvcache
    if (mMeta != nullptr && mMeta->file_name.size() > 0 && mMeta->file_flag == KVMeta::PendingRead) {
        // create new files
        std::string pathk = MNNFilePathConcat(mConfig.mPrefixCacheDir, mMeta->file_name) + "_" +
                            std::to_string(mMeta->layer_index) + ".k";
        std::string pathv = MNNFilePathConcat(mConfig.mPrefixCacheDir, mMeta->file_name) + "_" +
                            std::to_string(mMeta->layer_index++) + ".v";
        mMeta->layer_index = mMeta->layer_index % mMeta->layer_nums;
        auto old_key_fd = MNNOpenFile(pathk.c_str(), MNN_FILE_WRITE);
        auto old_value_fd = MNNOpenFile(pathv.c_str(), MNN_FILE_WRITE);
        if (old_key_fd == INVALID_FILE) {
            MNN_PRINT("Failed to open the file: %s\n", pathk.c_str());
        }
        if (old_value_fd == INVALID_FILE) {
            MNN_PRINT("Failed to open the file: %s\n", pathv.c_str());
        }

        // get kv cache file info
        auto oldKeySize = MNNGetFileSize(old_key_fd);
        auto oldValueSize = MNNGetFileSize(old_value_fd);

        size_t oldMaxLength = 0;
        if (mKeyQuantMode != KVQuantMode::None || mValueQuantMode != KVQuantMode::None) {
            MNN_ERROR("[Error]: Currently, kvcache save in disk not support quantized key/value\n");
        } else {
            size_t oldKeyMaxLength = oldKeySize / (mKvNumHead * ROUND_UP(mHeadDim, lP) * mBytes);
            size_t oldValueMaxLength = oldValueSize / (mKvNumHead * ROUND_UP(mHeadDim, hP) * mBytes);
            oldMaxLength = ALIMIN(oldKeyMaxLength, oldValueMaxLength);
        }
        if (oldMaxLength < meta->seqlen_in_disk) {
            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);
        }
        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);

        keySize = ALIMAX(keySize, oldKeySize);
        valueSize = ALIMAX(valueSize, oldValueSize);

        if (mKeyQuantMode == KVQuantMode::TQ3) {
            int tq3BytesPerSeq = (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
            mCurrentKeySizePerHead = (size_t)mMaxLength * tq3BytesPerSeq;
        } else if (mKeyQuantMode == KVQuantMode::TQ4) {
            int tq4BytesPerSeq = (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
            mCurrentKeySizePerHead = (size_t)mMaxLength * tq4BytesPerSeq;
        } else if (mKeyQuantMode == KVQuantMode::Int8) {
            mCurrentKeySizePerHead = ROUND_UP(mMaxLength, hP8) * ROUND_UP(mHeadDim, lP8) +
                                     2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(mMaxLength, hP8);
        } else {
            mCurrentKeySizePerHead = ROUND_UP(mMaxLength, hP) * ROUND_UP(mHeadDim, lP) * mBytes;
        }
        if (mValueQuantMode == KVQuantMode::TQ3) {
            int tq3BytesPerSeq = (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
            mCurrentValueSizePerHead = (size_t)mMaxLength * tq3BytesPerSeq;
        } else if (mValueQuantMode == KVQuantMode::TQ4) {
            int tq4BytesPerSeq = (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
            mCurrentValueSizePerHead = (size_t)mMaxLength * tq4BytesPerSeq;
        } else if (mValueQuantMode == KVQuantMode::Int8) {
            mCurrentValueSizePerHead = UP_DIV(mMaxLength, mFlashAttentionUpperKv) *
                                       (ROUND_UP(mHeadDim, hP8) * ROUND_UP(mFlashAttentionUpperKv, lP8) +
                                        2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(mHeadDim, hP8));
        } else {
            mCurrentValueSizePerHead = UP_DIV(mMaxLength, mFlashAttentionUpperKv) *
                                       (ROUND_UP(mHeadDim, hP) * ROUND_UP(mFlashAttentionUpperKv, lP) * mBytes);
        }

        createKVCacheFile();
        resetKVCacheFileSize(keySize, valueSize);
        expandKVCacheInDisk(oldMaxLength, oldKeySize, oldValueSize, keySize, valueSize, old_key_fd, old_value_fd);
        mPastLength = meta->seqlen_in_disk;
        mKVCacheInDisk = true;

        return;
    }

    // Do not use mMeta->add, because in VL models or Qnn case, mMeta->add is 0 or mMeta is nullptr.
    int kv_seq_len = seq_len;
    mMaxLength = kv_seq_len + mConfig.mExpandChunk;
    if (mUseFlashAttention) {
        setFlashAttentionUpperKv(MNN_FLASH_ATTENTION_BLOCK_SIZE);
    } else {
        setFlashAttentionUpperKv(mMaxLength);
    }

    // 1. compute size
    if (mKeyQuantMode == KVQuantMode::TQ3) {
        mCurrentKeySizePerHead = (size_t)mMaxLength * (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
    } else if (mKeyQuantMode == KVQuantMode::TQ4) {
        mCurrentKeySizePerHead = (size_t)mMaxLength * (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
    } else if (mKeyQuantMode == KVQuantMode::Int8) {
        mCurrentKeySizePerHead = ROUND_UP(mMaxLength, hP8) * ROUND_UP(mHeadDim, lP8) +
                                 2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(mMaxLength, hP8);
    } else {
        mCurrentKeySizePerHead = ROUND_UP(mMaxLength, hP) * ROUND_UP(mHeadDim, lP) * mBytes;
    }
    if (mValueQuantMode == KVQuantMode::TQ3) {
        mCurrentValueSizePerHead = (size_t)mMaxLength * (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
    } else if (mValueQuantMode == KVQuantMode::TQ4) {
        mCurrentValueSizePerHead = (size_t)mMaxLength * (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
    } else if (mValueQuantMode == KVQuantMode::Int8) {
        mCurrentValueSizePerHead = UP_DIV(mMaxLength, mFlashAttentionUpperKv) *
                                   (ROUND_UP(mHeadDim, hP8) * ROUND_UP(mFlashAttentionUpperKv, lP8) +
                                    2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(mHeadDim, hP8));
    } else {
        mCurrentValueSizePerHead = UP_DIV(mMaxLength, mFlashAttentionUpperKv) *
                                   (ROUND_UP(mHeadDim, hP) * ROUND_UP(mFlashAttentionUpperKv, lP) * mBytes);
    }
    size_t keySize = (size_t)mKvNumHead * mCurrentKeySizePerHead;
    size_t valueSize = (size_t)mKvNumHead * mCurrentValueSizePerHead;

    // 2. allocate buffer

    // case1: key&value size exceeds the limited size
    // case2: multi prompts share a common prefix kv cache info
    bool storeKvInDisk = !mConfig.mKVCacheDir.empty();
    bool sharePrefixKv = mMeta != nullptr && mMeta->file_name.size() > 0 && mMeta->file_flag == KVMeta::PendingWrite;

    if (sharePrefixKv) {
        mSaveShareKvPrefix = true;
        if (!MNNCreateDir(mConfig.mPrefixCacheDir.c_str())) {
            MNN_PRINT("Failed to create prefix cache file dir: %s\n", mConfig.mPrefixCacheDir.c_str());
        }
    }
    if (storeKvInDisk || sharePrefixKv) { // store kv in disk
        std::string keyStoredDst = "";
        std::string valueStoredDst = "";
        if (mMeta != nullptr) {
            mBasePrefixFileName =
                MNNFilePathConcat(mConfig.mPrefixCacheDir, mMeta->file_name) + "_" + std::to_string(mMeta->layer_index);
            keyStoredDst = sharePrefixKv ? mBasePrefixFileName + ".k" : "";
            valueStoredDst = sharePrefixKv ? mBasePrefixFileName + ".v" : "";
            mMeta->layer_index++;
            mMeta->layer_index = mMeta->layer_index % mMeta->layer_nums;
        }
        createKVCacheFile(keyStoredDst, valueStoredDst);
        resetKVCacheFileSize(keySize, valueSize);
        mmapKVCache(keySize, valueSize);
        mKVCacheInDisk = true;
    } else { // store kv in memory
        mPastKey.reset(Tensor::createDevice<int8_t>({mKvNumHead, (int)mCurrentKeySizePerHead}));
        mPastValue.reset(Tensor::createDevice<int8_t>({mKvNumHead, (int)mCurrentValueSizePerHead}));

        mBackend->onAcquireBuffer(mPastKey.get(), Backend::STATIC);
        mBackend->onAcquireBuffer(mPastValue.get(), Backend::STATIC);

        // initilize 0
        if (mHeadDim % lP || mKeyQuantMode != KVQuantMode::None) {
            memset(mPastKey->host<int8_t>(), 0, mPastKey->length(0) * mPastKey->stride(0));
        }
        if (lP > 1 || mValueQuantMode != KVQuantMode::None) {
            memset(mPastValue->host<int8_t>(), 0, mPastValue->length(0) * mPastValue->stride(0));
        }
    }
    // scale, zero point and sum of key for quantization
    if (mKeyQuantMode == KVQuantMode::Int8) { // quant K
        mKeySum.reset(Tensor::createDevice<int8_t>({mKvNumHead, ROUND_UP(mMaxLength, hP8) * QUANT_INFO_BYTES}));
        mKeyMax.reset(Tensor::createDevice<int8_t>({mKvNumHead, mHeadDim * QUANT_INFO_BYTES}));
        mBackend->onAcquireBuffer(mKeySum.get(), Backend::STATIC);
        mBackend->onAcquireBuffer(mKeyMax.get(), Backend::STATIC);

        for (int ks = 0; ks < mKvNumHead * mHeadDim; ++ks) {
            mKeyMax->host<float>()[ks] = std::numeric_limits<float>::lowest();
        }
        if (mBytes == 2) {
            auto core = static_cast<CPUBackend*>(mBackend)->functions();
            core->MNNFp32ToLowp(mKeyMax->host<float>(), (int16_t*)(mKeyMax->host<float>()), mKvNumHead * mHeadDim);
        }
    }
    if (mValueQuantMode == KVQuantMode::Int8) {
        mValueSum.reset(Tensor::createDevice<int8_t>(
            {mKvNumHead, (int)UP_DIV(mMaxLength, mFlashAttentionUpperKv), ROUND_UP(mHeadDim, hP8) * QUANT_INFO_BYTES}));
        mBackend->onAcquireBuffer(mValueSum.get(), Backend::STATIC);
        memset(mValueSum->host<int8_t>(), 0, mValueSum->stride(0) * mValueSum->length(0));
    }
}

void CPUKVCacheManager::onRealloc(KVMeta* meta) {
    auto kv_seq_len = meta->previous + meta->add - meta->remove + meta->computeReverseSize();
    if (kv_seq_len > mMaxLength) {
        // Realloc
        int oldMaxLength = mMaxLength;
        mMaxLength = (int)kv_seq_len + mConfig.mExpandChunk;
        if (mUseFlashAttention) {
            setFlashAttentionUpperKv(MNN_FLASH_ATTENTION_BLOCK_SIZE);
        } else {
            setFlashAttentionUpperKv(mMaxLength);
        }
        size_t oldKeySize = (size_t)mKvNumHead * mCurrentKeySizePerHead;
        size_t oldValueSize = (size_t)mKvNumHead * mCurrentValueSizePerHead;

        // update current key size per head
        if (mKeyQuantMode == KVQuantMode::TQ3) {
            mCurrentKeySizePerHead = (size_t)mMaxLength * (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
        } else if (mKeyQuantMode == KVQuantMode::TQ4) {
            mCurrentKeySizePerHead = (size_t)mMaxLength * (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
        } else if (mKeyQuantMode == KVQuantMode::Int8) {
            mCurrentKeySizePerHead = ROUND_UP(mMaxLength, hP8) * ROUND_UP(mHeadDim, lP8) +
                                     2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(mMaxLength, hP8);
        } else {
            mCurrentKeySizePerHead = UP_DIV(mMaxLength, hP) * ROUND_UP(mHeadDim, lP) * hP * mBytes;
        }
        // update current value size per head
        if (mValueQuantMode == KVQuantMode::TQ3) {
            mCurrentValueSizePerHead = (size_t)mMaxLength * (mHeadDim / TQ3_BLOCK_SIZE) * TQ3_BYTES_PER_BLOCK;
        } else if (mValueQuantMode == KVQuantMode::TQ4) {
            mCurrentValueSizePerHead = (size_t)mMaxLength * (mHeadDim / TQ4_BLOCK_SIZE) * TQ4_BYTES_PER_BLOCK;
        } else if (mValueQuantMode == KVQuantMode::Int8) {
            mCurrentValueSizePerHead = UP_DIV(mMaxLength, mFlashAttentionUpperKv) *
                                       (ROUND_UP(mHeadDim, hP8) * ROUND_UP(mFlashAttentionUpperKv, lP8) +
                                        2 * QUANT_INFO_BYTES * mConfig.mBlockNum * ROUND_UP(mHeadDim, hP8));
        } else {
            mCurrentValueSizePerHead = UP_DIV(mMaxLength, mFlashAttentionUpperKv) *
                                       (ROUND_UP(mHeadDim, hP) * ROUND_UP(mFlashAttentionUpperKv, lP) * mBytes);
        }
        size_t keySize = (size_t)mKvNumHead * mCurrentKeySizePerHead;
        size_t valueSize = (size_t)mKvNumHead * mCurrentValueSizePerHead;

        /*==== No limit for kvcache ====*/
        if (mKVCacheInDisk == false) {
            expandKVCacheInMem(oldMaxLength);
        } else {
            expandKVCacheInDisk(oldMaxLength, oldKeySize, oldValueSize, keySize, valueSize);
        }
        /* No matter where is the kvcache, the scales and zero points are always in memory, since their size is very
         * small */
        if (mKeyQuantMode == KVQuantMode::Int8) {
            auto newKeySumTensor = Tensor::createDevice<int32_t>({mKvNumHead, UP_DIV(mMaxLength, hP8), hP8});
            mBackend->onAcquireBuffer(newKeySumTensor, Backend::STATIC);
            for (int h = 0; h < mKvNumHead; h++) {
                memcpy(newKeySumTensor->host<int8_t>() + h * UP_DIV(mMaxLength, hP8) * hP8 * 4,
                       mKeySum->host<int8_t>() + h * UP_DIV(oldMaxLength, hP8) * hP8 * 4,
                       UP_DIV(oldMaxLength, hP8) * hP8 * 4);
            }
            mKeySum.reset(newKeySumTensor);
        }
        if (mValueQuantMode == KVQuantMode::Int8) {
            auto newValueSumTensor =
                Tensor::createDevice<int8_t>({mKvNumHead, (int)UP_DIV(mMaxLength, mFlashAttentionUpperKv),
                                              ROUND_UP(mHeadDim, hP8) * QUANT_INFO_BYTES});
            mBackend->onAcquireBuffer(newValueSumTensor, Backend::STATIC);
            auto remainSizePerHead = mValueSum->stride(0);
            auto increSizePerHead = newValueSumTensor->stride(0) - mValueSum->stride(0);
            for (int h = 0; h < mKvNumHead; ++h) {
                memcpy(newValueSumTensor->host<int8_t>() + h * newValueSumTensor->stride(0),
                       mValueSum->host<int8_t>() + h * mValueSum->stride(0), remainSizePerHead);
                // memset 0
                if (increSizePerHead > 0) {
                    memset(newValueSumTensor->host<int8_t>() + h * newValueSumTensor->stride(0) + remainSizePerHead, 0,
                           increSizePerHead);
                }
            }
            mValueSum.reset(newValueSumTensor);
        }
    }
    // Remove
    auto start = mPastLength - meta->remove;
    if (0 == meta->n_reserve || mKeyQuantMode != KVQuantMode::None ||
        mValueQuantMode != KVQuantMode::None) { // n_reserve > 0 is not currently supported when K or V is quantized.
        mPastLength = start;
        return;
    }
#if 1
    auto dstIndex = start;
    for (int n = 0; n < meta->n_reserve; ++n) {
        auto begin = meta->reserve[2 * n];
        auto size = meta->reserve[2 * n + 1];
        auto srcIndex = start + begin;
        if (mBytes == 2) {
            moveKV<FLOAT16_T>(srcIndex, dstIndex, size);
        } else {
            moveKV<float>(srcIndex, dstIndex, size);
        }
        dstIndex += size;
    }
    mPastLength = dstIndex;
#else
    // Don't support not align reserve
    auto align = hP;
    auto dstStart = start;
    auto lastValidSrcEnd = start;
    for (int n = 0; n < meta->n_reserve; ++n) {
        auto lastEndAlign = UP_DIV(lastValidSrcEnd, align) * align;
        auto begin = meta->reserve[2 * n];
        auto size = meta->reserve[2 * n + 1];
        auto startAlign = ((begin + start) / align) * align;
        if (startAlign <= lastEndAlign) {
            // Fullly reserve
            dstStart = dstStart + size;
            lastValidSrcEnd = begin + start + size;
            continue;
        }
        auto end = begin + start + size;
        auto endAlign = UP_DIV(end, align) * align;

        auto sizeUnit = (endAlign - startAlign) / align;
        auto dstStartAlign = UP_DIV(dstStart, align) * align;

        // TODO: Support Quant
        //        mPastKey.reset(Tensor::createDevice<float>({mKvNumHead, UP_DIV(mMaxLength, hP), mHeadDim, hP}));

        // Move K
        auto keyStride = UP_DIV(mMaxLength, align) * align * ROUND_UP(mHeadDim, lP);
        auto dstKAddr = keyAddr() + dstStartAlign * ROUND_UP(mHeadDim, lP) * mBytes;
        auto srcKAddr = keyAddr() + startAlign * ROUND_UP(mHeadDim, lP) * mBytes;
        for (int i = 0; i < mKvNumHead; ++i) {
            auto dst = dstKAddr + i * keyStride * mBytes;
            auto src = srcKAddr + i * keyStride * mBytes;
            for (int j = 0; j < sizeUnit; ++j) {
                ::memcpy(dst + j * align * ROUND_UP(mHeadDim, lP) * mBytes,
                         src + j * align * ROUND_UP(mHeadDim, lP) * mBytes, align * ROUND_UP(mHeadDim, lP) * mBytes);
            }
        }

        // Move V
        auto dstVAddr = valudAddr() + dstStartAlign * align * mBytes;
        auto srcVAddr = valudAddr() + startAlign * align * mBytes;
        auto number = mKvNumHead * UP_DIV(mHeadDim, align);
        for (int i = 0; i < number; ++i) {
            auto dst = dstVAddr + i * ROUND_UP(mMaxLength, lP) * align * mBytes;
            auto src = srcVAddr + i * ROUND_UP(mMaxLength, lP) * align * mBytes;
            for (int j = 0; j < sizeUnit; ++j) {
                ::memcpy(dst + j * align * align * mBytes, src + j * align * align * mBytes, align * align * mBytes);
            }
        }
        dstStart = dstStart + size;
        lastValidSrcEnd = begin + start + size;
    }
    mPastLength = dstStart;
#endif
}

void CPUKVCacheManager::saveKVCacheInDisk() {
    // get original kv cache info
    auto keySize = MNNGetFileSize(mKeyCacheFD);
    auto valueSize = MNNGetFileSize(mValueCacheFD);
    mmapKVCache(keySize, valueSize);
    if (!MNNCreateDir(mConfig.mPrefixCacheDir.c_str())) {
        MNN_PRINT("Failed to create prefix cache file dir: %s\n", mConfig.mPrefixCacheDir.c_str());
    }

    // create new files
    std::string pathk =
        MNNFilePathConcat(mConfig.mPrefixCacheDir, mMeta->file_name) + "_" + std::to_string(mMeta->layer_index) + ".k";
    std::string pathv = MNNFilePathConcat(mConfig.mPrefixCacheDir, mMeta->file_name) + "_" +
                        std::to_string(mMeta->layer_index++) + ".v";
    mMeta->layer_index = mMeta->layer_index % mMeta->layer_nums;

    auto new_key_fd = MNNCreateFile(pathk.c_str());
    auto new_value_fd = MNNCreateFile(pathv.c_str());
    if (new_key_fd == INVALID_FILE) {
        MNN_PRINT("Failed to create the file: %s\n", pathk.c_str());
    }
    if (new_value_fd == INVALID_FILE) {
        MNN_PRINT("Failed to create the file: %s\n", pathv.c_str());
    }
    // set new file size
    if (MNNSetFileSize(new_key_fd, keySize) != MNN::NO_ERROR ||
        MNNSetFileSize(new_value_fd, valueSize) != MNN::NO_ERROR) {
        MNN_PRINT("Failed to resize the kvcache files!\n");
    }
    // mmap files
    int8_t* mMapNewKeyAddr = (int8_t*)MNNMmapFile(new_key_fd, keySize);
    if (mMapNewKeyAddr == nullptr) {
        MNN_PRINT("Failed to memory-map the new kvcache!\n");
    }
    int8_t* mMapNewValueAddr = (int8_t*)MNNMmapFile(new_value_fd, valueSize);
    if (mMapNewValueAddr == nullptr) {
        MNN_PRINT("Failed to memory-map the kvcache!\n");
    }

    // copy
    memcpy(mMapNewKeyAddr, mMapKeyAddr, keySize);
    memcpy(mMapNewValueAddr, mMapValueAddr, valueSize);

    // unmap new files
    if (mMapNewKeyAddr != nullptr) {
        MNNUnmapFile(mMapNewKeyAddr, keySize);
        mMapNewKeyAddr = nullptr;
    }
    if (mMapNewValueAddr != nullptr) {
        MNNUnmapFile(mMapNewValueAddr, valueSize);
        mMapNewValueAddr = nullptr;
    }
    // close file
    if (new_key_fd != INVALID_FILE) {
        MNNCloseFile(new_key_fd);
        new_key_fd = INVALID_FILE;
    }
    if (new_value_fd != INVALID_FILE) {
        MNNCloseFile(new_value_fd);
        new_value_fd = INVALID_FILE;
    }
}

void CPUKVCacheManager::onClear() {
    if (mKVCacheInDisk) {
        // mSaveShareKvPrefix also need unmap file
        unmapKVCache(mCurrentKeySizePerHead * (size_t)mKvNumHead, mCurrentValueSizePerHead * (size_t)mKvNumHead);
        if (!mSaveShareKvPrefix) {
            // delete temp kvcache file
            removeKVCacheFile();
        }
        mKVCacheInDisk = false;
    }
    mPastKey.reset();
    mPastValue.reset();
    mKeySum.reset();
    mKeyMax.reset();
    mValueSum.reset();
    mMaxLength = mPastLength = 0;
}

template <typename T>
void CPUKVCacheManager::ProcessKey(const Tensor* key, int seqLen, int kvHead) {
    if ((mKeyQuantMode == KVQuantMode::TQ3) || (mKeyQuantMode == KVQuantMode::TQ4)) {
        int bytesPerBlock = (mKeyQuantMode == KVQuantMode::TQ3) ? TQ3_BYTES_PER_BLOCK : TQ4_BYTES_PER_BLOCK;
        int blockSize = (mKeyQuantMode == KVQuantMode::TQ3) ? TQ3_BLOCK_SIZE : TQ4_BLOCK_SIZE;
        int bytesPerSeq = (mHeadDim / blockSize) * bytesPerBlock;
        uint8_t* keyDst = (uint8_t*)addrOfKey(kvHead);
        for (int i = 0; i < seqLen; i++) {
            T* src = key->host<T>() + i * mKvNumHead * mHeadDim + kvHead * mHeadDim;
            uint8_t* dst = keyDst + (mPastLength + i) * bytesPerSeq;
            float block[TQ4_BLOCK_SIZE]; // TQ4_BLOCK_SIZE >= TQ3_BLOCK_SIZE
            for (int b = 0; b < mHeadDim / blockSize; b++) {
                for (int k = 0; k < blockSize; k++) {
                    block[k] = (float)src[b * blockSize + k];
                }
                if (mKeyQuantMode == KVQuantMode::TQ3) {
                    tq3_quantize_block(dst + b * bytesPerBlock, block);
                } else {
                    tq4_quantize_block(dst + b * bytesPerBlock, block);
                }
            }
        }
    } else if (mKeyQuantMode ==
               KVQuantMode::Int8) { // [seqLen, headDim] -> [maxlen/hP8, blockNum, (headDim/blockNum)/lP8, hP8, lP8]
        int8_t* keyDst = reinterpret_cast<int8_t*>(addrOfKey(kvHead));
        float* sumDst = reinterpret_cast<float*>(addrOfKeySum(kvHead));

        auto blockL = UP_DIV(mHeadDim, mConfig.mBlockNum);
        auto weightStride1 = ROUND_UP(blockL, lP8) * hP8;
        auto weightStride2 = lP8 * hP8;
        auto packedWeightStride1 = weightStride1 + 2 * QUANT_INFO_BYTES * hP8;

        T* keyMax = reinterpret_cast<T*>(addrOfKeyMax(kvHead));
        int32_t params[] = {mKvNumHead, seqLen, mHeadDim, mConfig.mBlockNum, eP8, lP8, hP8, mPastLength, kvHead};
        mQuantKeyFunc(keyDst, key->host<float>(), sumDst, (float*)keyMax, params);
    } else { // target: [maxlen/hP, headdim/lP, hP, lP]
        T* key_dst = reinterpret_cast<T*>(addrOfKey(kvHead));
        auto stride0 = ROUND_UP(mHeadDim, lP) * hP;
        auto stride1 = hP * lP;
        for (int i = 0; i < seqLen; i++) {
            T* key_src = key->host<T>() + i * mKvNumHead * mHeadDim + kvHead * mHeadDim;
            int out_index = (mPastLength + i) / hP;
            int in_index = (mPastLength + i) % hP;
            for (int j = 0; j < mHeadDim; j++) {
                key_dst[out_index * stride0 + (j / lP) * stride1 + in_index * lP + (j % lP)] = key_src[j];
            }
        }
    }
}

template <typename T>
void CPUKVCacheManager::ProcessValue(const Tensor* value, int seqLen, int kvHead) { // [headdim/hP, maxlen, hP]
    const bool valueC4 = TensorUtils::getDescribe(value)->dimensionFormat == MNN_DATA_FORMAT_NC4HW4;
    const int pack = static_cast<CPUBackend*>(mBackend)->functions()->pack;
    const auto valueSrc = value->host<T>();
    auto loadValue = [&](int token, int channel) {
        if (valueC4) {
            return valueSrc[c4Offset(token, channel, seqLen, pack)];
        }
        return valueSrc[token * mKvNumHead * mHeadDim + channel];
    };
    const int channelBase = kvHead * mHeadDim;
    if ((mValueQuantMode == KVQuantMode::TQ3) || (mValueQuantMode == KVQuantMode::TQ4)) {
        int bytesPerBlock = (mValueQuantMode == KVQuantMode::TQ3) ? TQ3_BYTES_PER_BLOCK : TQ4_BYTES_PER_BLOCK;
        int blockSize = (mValueQuantMode == KVQuantMode::TQ3) ? TQ3_BLOCK_SIZE : TQ4_BLOCK_SIZE;
        int bytesPerSeq = (mHeadDim / blockSize) * bytesPerBlock;
        uint8_t* valueDst = (uint8_t*)addrOfValue(kvHead);
        for (int i = 0; i < seqLen; i++) {
            uint8_t* dst = valueDst + (mPastLength + i) * bytesPerSeq;
            float block[TQ4_BLOCK_SIZE];
            for (int b = 0; b < mHeadDim / blockSize; b++) {
                for (int k = 0; k < blockSize; k++) {
                    block[k] = (float)loadValue(i, channelBase + b * blockSize + k);
                }
                if (mValueQuantMode == KVQuantMode::TQ3) {
                    tq3_quantize_block(dst + b * bytesPerBlock, block);
                } else {
                    tq4_quantize_block(dst + b * bytesPerBlock, block);
                }
            }
        }
    } else if (mValueQuantMode == KVQuantMode::Int8) {
        int8_t* valueDst = reinterpret_cast<int8_t*>(addrOfValue(kvHead));
        float* valueSum = reinterpret_cast<float*>(addrOfValueSum(kvHead));

        int32_t params[] = {mKvNumHead, seqLen,
                            mHeadDim,   mConfig.mBlockNum,
                            mMaxLength, lP8,
                            hP8,        mPastLength,
                            kvHead,     (int32_t)mFlashAttentionUpperKv};
        if (!valueC4) {
            mQuantValueFunc(valueDst, value->host<float>(), valueSum, params);
        } else {
            auto weightStride2 = lP8 * hP8;
            auto weightStride1 = UP_DIV((int32_t)mFlashAttentionUpperKv, lP8) * weightStride2;
            auto packedStride1 = (int)(weightStride1 + 2 * hP8 * sizeof(float));
            auto packedStride0 = UP_DIV(mHeadDim, hP8) * packedStride1;

            if (mPastLength == 0) {
                for (int d = 0; d < mHeadDim; ++d) {
                    float* scalePtr = (float*)(valueDst + (d / hP8) * packedStride1 + weightStride1) + (d % hP8);
                    float* biasPtr = scalePtr + hP8;

                    float dMax = (float)loadValue(0, channelBase + d);
                    float dMin = dMax;
                    for (int s = 0; s < seqLen; ++s) {
                        float data = (float)loadValue(s, channelBase + d);
                        dMax = ALIMAX(dMax, data);
                        dMin = ALIMIN(dMin, data);
                    }

                    float range = dMax - dMin;
                    if (range < 1e-6f) {
                        scalePtr[0] = 0.0f;
                        biasPtr[0] = dMax;
                    } else {
                        float scale = range / 255.0f;
                        float bias = range / 255.0f * 128.0f + dMin;
                        scalePtr[0] = scale;
                        biasPtr[0] = bias;
                    }
                }
            }

            // FA value int8 keeps one scale/bias per channel and copies it to each block. Decode within an existing
            // block intentionally reuses that scale; updating it requires re-quantizing cached values and valueSum.
            if (mPastLength == 0 || (mPastLength % mFlashAttentionUpperKv) == 0) {
                int32_t d0 = UP_DIV(mMaxLength, (int32_t)mFlashAttentionUpperKv);
                int32_t d1 = UP_DIV(mHeadDim, hP8);
                for (int k = 0; k < d0; ++k) {
                    for (int r = 0; r < d1; ++r) {
                        float* scalePtr = (float*)(valueDst + k * packedStride0 + r * packedStride1 + weightStride1);
                        float* biasPtr = scalePtr + hP8;
                        memcpy(scalePtr, valueDst + r * packedStride1 + weightStride1, hP8 * sizeof(float));
                        memcpy(biasPtr, valueDst + r * packedStride1 + weightStride1 + hP8 * sizeof(float),
                               hP8 * sizeof(float));
                    }
                }
            }

            for (int d = 0; d < mHeadDim; ++d) {
                int idxBase = (d / hP8) * packedStride1 + (d % hP8) * lP8;
                int8_t* dstBase = valueDst + idxBase;
                float* scaleBase = (float*)(valueDst + (d / hP8) * packedStride1 + weightStride1) + (d % hP8);
                float* biasBase = scaleBase + hP8;
                float* sumBase = valueSum + (d / hP8) * hP8 + (d % hP8);

                float qscale = scaleBase[0] < 1e-6f ? 0.0f : 1.0f / scaleBase[0];
                float qbias = scaleBase[0] < 1e-6f ? 0.0f : (-biasBase[0] / scaleBase[0]);
                for (int s = 0; s < seqLen; ++s) {
                    int kvSeqIndx = s + mPastLength;
                    int idxInner = (kvSeqIndx / (int32_t)mFlashAttentionUpperKv) * packedStride0 +
                                   (kvSeqIndx % (int32_t)mFlashAttentionUpperKv) / lP8 * weightStride2 +
                                   (kvSeqIndx % (int32_t)mFlashAttentionUpperKv) % lP8;
                    float xf = (float)loadValue(s, channelBase + d);
                    int8_t xq = ALIMAX(ALIMIN(127, static_cast<int32_t>(roundf(xf * qscale + qbias))), -128);
                    dstBase[idxInner] = xq;

                    int idxSum = (kvSeqIndx / (int32_t)mFlashAttentionUpperKv) * ROUND_UP(mHeadDim, hP8);
                    sumBase[idxSum] += ((float)xq * scaleBase[0] + biasBase[0]);
                }
            }
        }
    } else {
        // [mHeadDim/hP, mMaxLength/lP, hP, lP]
        auto stride0 = ROUND_UP(mMaxLength, lP) * hP;
        auto stride1 = hP * lP;

        auto weightStride2 = lP * hP;
        auto weightStride1 = UP_DIV((int32_t)mFlashAttentionUpperKv, lP) * weightStride2;
        auto weightStride0 = weightStride1 * UP_DIV(mHeadDim, hP);

        T* value_dst = reinterpret_cast<T*>(addrOfValue(kvHead));
        for (int i = 0; i < seqLen; i++) {
            // int seqLenOut = (mPastLength + i) / lP;
            // int seqLenIn = (mPastLength + i) % lP;

            int kvSeqIndx = mPastLength + i;
            int idxInner = (kvSeqIndx / (int32_t)mFlashAttentionUpperKv) * weightStride0 +
                           (kvSeqIndx % (int32_t)mFlashAttentionUpperKv) / lP * weightStride2 +
                           (kvSeqIndx % (int32_t)mFlashAttentionUpperKv) % lP;
            for (int j = 0; j < mHeadDim; j++) {
                int idxBase = (j / hP) * weightStride1 + (j % hP) * lP;
                int out_index = j / hP;
                int in_index = j % hP;
                // value_dst[out_index * stride0 + seqLenOut * stride1 + in_index * lP + seqLenIn] = value_src[j];
                value_dst[idxBase + idxInner] = loadValue(i, channelBase + j);
            }
        }
    }
}

size_t CPUKVCacheManager::keyIndex(int seq, int dim) const {
    return (seq / hP) * ROUND_UP(mHeadDim, lP) * hP + (dim / lP) * hP * lP + (seq % hP) * lP + (dim % lP);
}

size_t CPUKVCacheManager::valueIndex(int seq, int dim) const {
    auto stride1 = UP_DIV((int32_t)mFlashAttentionUpperKv, lP) * hP * lP;
    auto stride0 = stride1 * UP_DIV(mHeadDim, hP);
    auto seqInBlock = seq % (int32_t)mFlashAttentionUpperKv;
    return (seq / (int32_t)mFlashAttentionUpperKv) * stride0 +
           (dim / hP) * stride1 +
           (seqInBlock / lP) * hP * lP +
           (dim % hP) * lP +
           (seqInBlock % lP);
}

template <typename T>
void CPUKVCacheManager::moveKV(int src, int dst, int size) {
    for (int h = 0; h < mKvNumHead; ++h) {
        auto kPtr = reinterpret_cast<T*>(addrOfKey(h));
        auto vPtr = reinterpret_cast<T*>(addrOfValue(h));
        for (int i = 0; i < size; i++) {
            for (int j = 0; j < mHeadDim; j++) {
                kPtr[keyIndex(dst + i, j)] = kPtr[keyIndex(src + i, j)];
                vPtr[valueIndex(dst + i, j)] = vPtr[valueIndex(src + i, j)];
            }
        }
    }
}

void CPUKVCacheManager::onUpdateKV(const Tensor* key, const Tensor* value, int add) {
    auto core = static_cast<CPUBackend*>(mBackend)->functions();
    int seq_len = add;
    auto divPart = UP_DIV(mKvNumHead, 1);
    MNN_CONCURRENCY_BEGIN(tId, 1) {
        auto remainPart = mKvNumHead - tId * divPart;
        if (remainPart > 0) {
            remainPart = ALIMIN(divPart, remainPart);
            int startIdx = tId * divPart;
            int endIdx = startIdx + remainPart;
            for (int h = startIdx; h < endIdx; ++h) {
                if (mBytes == 2) {
                    ProcessKey<FLOAT16_T>(key, seq_len, h);
                    ProcessValue<FLOAT16_T>(value, seq_len, h);
                } else {
                    ProcessKey<float>(key, seq_len, h);
                    ProcessValue<float>(value, seq_len, h);
                }
            }
        }
    }
    MNN_CONCURRENCY_END();
    mPastLength += seq_len;
}

} // namespace MNN

#endif // MNN_SUPPORT_TRANSFORMER_FUSE